diff --git a/scripts/quantize_nvfp4.py b/scripts/quantize_nvfp4.py index 23034b6..e386f20 100644 --- a/scripts/quantize_nvfp4.py +++ b/scripts/quantize_nvfp4.py @@ -48,23 +48,9 @@ CALIB_SAVE_PATH = "/root/nvidia-meeting/v4_nvfp4_calibrated_state.pt" def apply_patches(): """Apply runtime patches for V4 compatibility.""" - from modelopt.torch.quantization.nn.modules import quant_module, tensor_quantizer as tq_module + from modelopt.torch.quantization.nn.modules import tensor_quantizer as tq_module - # 1. Patch quant_module.py for V4's ModuleList expert quantizers - orig_iter = quant_module._QuantFusedExperts.iter_weights_for_calibration - - def patched_iter_weights_for_calibration(self, **kwargs): - """Handle V4's nn.ModuleList expert quantizers (vs singular TensorQuantizer).""" - for name, quantizer in self.named_modules(): - if not isinstance(quantizer, quant_module.TensorQuantizer): - continue - if quantizer.is_enabled: - yield name, quantizer - - quant_module._QuantFusedExperts.iter_weights_for_calibration = patched_iter_weights_for_calibration - print("✓ Patched _QuantFusedExperts.iter_weights_for_calibration for V4 ModuleList") - - # 2. Patch TensorQuantizer.export_amax to move _amax to CPU before reading + # 1. Patch TensorQuantizer.export_amax to move _amax to CPU before reading orig_export_amax = tq_module.TensorQuantizer.export_amax def patched_export_amax(self): @@ -77,7 +63,7 @@ def apply_patches(): tq_module.TensorQuantizer.export_amax = patched_export_amax print("✓ Patched TensorQuantizer.export_amax (CPU safety)") - # 3. Patch NVFP4QTensor.get_activation_scaling_factor for graceful degradation + # 2. Patch NVFP4QTensor.get_activation_scaling_factor for graceful degradation from modelopt.torch.quantization.qtensor import nvfp4_tensor orig_get_asf = nvfp4_tensor.NVFP4QTensor.get_activation_scaling_factor