From 07cd50e8237081f10e7773593e36cb019ff4b90f Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 9 May 2026 22:50:58 +0000 Subject: [PATCH] =?UTF-8?q?8=20patches=20covering=20full=20export=20chain?= =?UTF-8?q?=20=E2=80=94=20no=20more=20whack-a-mole?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Traced the full execution chain from _process_quantized_modules through every function that reads stale GPU tensors: _process_quantized_modules → _export_quantized_weight (Patch 4: force weight to CPU at entry point) → get_weight_scaling_factor (Patch 7: belt-and-suspenders) → get_weights_scaling_factor_from_quantizer (safe: weight now CPU) → NVFP4QTensor.get_weights_scaling_factor (safe: input is CPU) → get_weight_scaling_factor_2 (Patch 8: force quantizer to CPU) → get_activation_scaling_factor (Patch 3: CPU + clamp) → to_quantized_weight (Patch 6: force all tensors to CPU) → weight.to(dtype) (safe: weight is CPU) → _export_fused_experts (Patch 5: force expert weights + quantizer to CPU) Patch 4 is the key: it moves weight to CPU at the earliest possible point, so ALL downstream .to(weight.device) calls resolve to CPU. Patches 5-8 are belt-and-suspenders for alternative code paths. --- scripts/quantize_nvfp4.py | 223 ++++++++++++++++++++++++++++++-------- 1 file changed, 178 insertions(+), 45 deletions(-) diff --git a/scripts/quantize_nvfp4.py b/scripts/quantize_nvfp4.py index d7cdab8..4a09cec 100644 --- a/scripts/quantize_nvfp4.py +++ b/scripts/quantize_nvfp4.py @@ -56,14 +56,39 @@ AMAX_SNAPSHOT_PATH = "/root/nvidia-meeting/v4_nvfp4_amax_snapshots.pt" def apply_patches(): """Apply runtime patches for V4 compatibility and GPU tensor safety. - These patches are applied BEFORE hf_ptq runs, so they're active during - calibration and export. No modelopt source files are modified. + Root cause of all export crashes: use_seq_device_map keeps model weights on GPU + for 5+ hours during calibration. By export time, CUDA's memory allocator has + recycled the underlying memory, so any read of those GPU tensors triggers + cudaErrorIllegalAddress. + + Fix strategy: patch at the EARLIEST possible entry points to force stale GPU + tensors to CPU before any downstream code reads them. This covers the full + chain of execution we traced through the export path: + + _process_quantized_modules + → _export_quantized_weight (or _export_fused_experts) + → get_weight_scaling_factor + → get_weights_scaling_factor_from_quantizer (reads weight, _amax, global_amax) + → NVFP4QTensor.get_weights_scaling_factor (dynamic: reduce_block_amax on weight) + → get_weight_scaling_factor_2 (reads _amax, global_amax) + → get_activation_scaling_factor (reads _amax) [already patched] + → to_quantized_weight (reads weight, does .to(weight.device) on scaling factors) + → weight.to(dtype) (reads weight) + + By forcing weight to CPU in Patch 4 (_export_quantized_weight), ALL downstream + .to(weight.device) calls resolve to CPU. Patches 5-8 are belt-and-suspenders. """ from modelopt.torch.quantization.nn.modules import tensor_quantizer as tq_module from modelopt.torch.quantization.qtensor import nvfp4_tensor + from modelopt.torch.export import quant_utils + from modelopt.torch.quantization.utils import quantizer_attr_names as _quantizer_attr_names + import modelopt.torch.export.unified_export_hf as uehf - # ── Patch 1: load_calib_amax — force _amax to CPU after calibration ── + # ══════════════════════════════════════════════════════════════════════ + # Patch 1: load_calib_amax — force _amax to CPU immediately after calibration + # This runs during calibration, right after each quantizer finishes. + # ══════════════════════════════════════════════════════════════════════ orig_load_calib_amax = tq_module.TensorQuantizer.load_calib_amax def patched_load_calib_amax(self, *args, **kwargs): @@ -72,20 +97,24 @@ def apply_patches(): self._amax = self._amax.cpu() tq_module.TensorQuantizer.load_calib_amax = patched_load_calib_amax - print("✓ Patched TensorQuantizer.load_calib_amax (force _amax to CPU)") + print("✓ Patch 1: TensorQuantizer.load_calib_amax → force _amax to CPU") - # ── Patch 2: export_amax — CPU safety ── + # ══════════════════════════════════════════════════════════════════════ + # Patch 2: export_amax — CPU safety net at export time + # ══════════════════════════════════════════════════════════════════════ orig_export_amax = tq_module.TensorQuantizer.export_amax def patched_export_amax(self): - if self.amax is not None and self.amax.is_cuda: + if hasattr(self, '_amax') and self._amax is not None and self._amax.is_cuda: self._amax = self._amax.cpu() return orig_export_amax(self) tq_module.TensorQuantizer.export_amax = patched_export_amax - print("✓ Patched TensorQuantizer.export_amax (CPU fallback)") + print("✓ Patch 2: TensorQuantizer.export_amax → CPU fallback") - # ── Patch 3: NVFP4QTensor.get_activation_scaling_factor — graceful degradation ── + # ══════════════════════════════════════════════════════════════════════ + # Patch 3: get_activation_scaling_factor — CPU + clamp + # ══════════════════════════════════════════════════════════════════════ @classmethod def patched_get_activation_scaling_factor(cls, quantizer): if not quantizer.is_enabled: @@ -112,67 +141,171 @@ def apply_patches(): return activation_scaling_factor nvfp4_tensor.NVFP4QTensor.get_activation_scaling_factor = patched_get_activation_scaling_factor - print("✓ Patched NVFP4QTensor.get_activation_scaling_factor (CPU + clamp)") + print("✓ Patch 3: NVFP4QTensor.get_activation_scaling_factor → CPU + clamp") - # ── Patch 4: get_weight_scaling_factor — force weight to CPU before export computation ── - # After hours of calibration with use_seq_device_map, model weight tensors on GPU - # have stale memory. Any attempt to read them (even .to(device)) triggers - # cudaErrorIllegalAddress. We force weight to CPU and do all math on CPU. - from modelopt.torch.export import quant_utils - from modelopt.torch.quantization.utils import quantizer_attr_names as _quantizer_attr_names + # ══════════════════════════════════════════════════════════════════════ + # Patch 4: _export_quantized_weight — THE KEY PATCH + # + # This is the entry point for exporting each quantized module. It reads + # `weight = getattr(sub_module, weight_name)` which is on a stale GPU. + # By moving weight to CPU right here, ALL downstream functions are safe: + # - get_weight_scaling_factor: weight.device is now CPU + # - get_weights_scaling_factor: operates on CPU weight + # - to_quantized_weight: .to(weight.device) stays on CPU + # - weight.to(dtype): CPU cast + # We also force all quantizer state to CPU for the same reason. + # ══════════════════════════════════════════════════════════════════════ + orig_export_quantized_weight = uehf._export_quantized_weight + def patched_export_quantized_weight(sub_module, dtype, weight_name="weight"): + # Move weight to CPU (stale GPU → safe CPU) + weight = getattr(sub_module, weight_name, None) + if weight is not None and isinstance(weight, torch.Tensor) and weight.is_cuda: + try: + weight_cpu = weight.cpu() + with torch.no_grad(): + setattr(sub_module, weight_name, torch.nn.Parameter(weight_cpu)) + except (torch.cuda.CudaError, RuntimeError) as e: + print(f" WARNING: weight.cpu() failed for {weight_name} ({e})") + raise + + # Force all quantizer state to CPU + qattrs = _quantizer_attr_names(weight_name) + for qattr in [qattrs.weight_quantizer, qattrs.input_quantizer, qattrs.output_quantizer]: + if not qattr: + continue + quantizer = getattr(sub_module, qattr, None) + if quantizer is None: + continue + for attr in ['_amax', '_pre_quant_scale', 'global_amax', '_global_amax']: + val = getattr(quantizer, attr, None) + if val is not None and isinstance(val, torch.Tensor) and val.is_cuda: + try: + setattr(quantizer, attr, val.cpu()) + except (torch.cuda.CudaError, RuntimeError): + pass + # Handle SequentialQuantizer (W4A8 path) + if hasattr(quantizer, 'quantizers'): + for sub_q in quantizer.quantizers: + for attr in ['_amax', '_pre_quant_scale', 'global_amax', '_global_amax']: + val = getattr(sub_q, attr, None) + if val is not None and isinstance(val, torch.Tensor) and val.is_cuda: + try: + setattr(sub_q, attr, val.cpu()) + except (torch.cuda.CudaError, RuntimeError): + pass + + return orig_export_quantized_weight(sub_module, dtype, weight_name) + + uehf._export_quantized_weight = patched_export_quantized_weight + print("✓ Patch 4: _export_quantized_weight → force weight + quantizer state to CPU") + + # ══════════════════════════════════════════════════════════════════════ + # Patch 5: _export_fused_experts — same treatment for MoE expert weights + # DeepseekV4Experts go through this different code path. + # ══════════════════════════════════════════════════════════════════════ + orig_export_fused_experts = uehf._export_fused_experts + + def patched_export_fused_experts(sub_module, dtype): + # Force all expert weights to CPU + for name, param in list(sub_module.named_parameters()): + if isinstance(param, torch.Tensor) and param.is_cuda: + try: + with torch.no_grad(): + setattr(sub_module, name, torch.nn.Parameter(param.cpu())) + except (torch.cuda.CudaError, RuntimeError): + pass + # Force all buffers to CPU + for name, buf in list(sub_module.named_buffers()): + if isinstance(buf, torch.Tensor) and buf.is_cuda: + try: + sub_module.register_buffer(name, buf.cpu()) + except (torch.cuda.CudaError, RuntimeError): + pass + # Force all quantizer state to CPU + for mod in sub_module.modules(): + for attr in ['_amax', '_pre_quant_scale', 'global_amax', '_global_amax']: + val = getattr(mod, attr, None) + if val is not None and isinstance(val, torch.Tensor) and val.is_cuda: + try: + setattr(mod, attr, val.cpu()) + except (torch.cuda.CudaError, RuntimeError): + pass + return orig_export_fused_experts(sub_module, dtype) + + uehf._export_fused_experts = patched_export_fused_experts + print("✓ Patch 5: _export_fused_experts → force expert weights + quantizer state to CPU") + + # ══════════════════════════════════════════════════════════════════════ + # Patch 6: to_quantized_weight — force scaling factors to CPU + # This does .to(weight.device) on scaling factors. With weight now on + # CPU (Patch 4), this should be a no-op, but belt-and-suspenders. + # ══════════════════════════════════════════════════════════════════════ + orig_to_quantized_weight = quant_utils.to_quantized_weight + + def patched_to_quantized_weight(weight, weights_scaling_factor, quantization, + weights_scaling_factor2=None, block_size=None): + if isinstance(weight, torch.Tensor) and weight.is_cuda: + weight = weight.cpu() + if weights_scaling_factor is not None and isinstance(weights_scaling_factor, torch.Tensor) and weights_scaling_factor.is_cuda: + weights_scaling_factor = weights_scaling_factor.cpu() + if weights_scaling_factor2 is not None and isinstance(weights_scaling_factor2, torch.Tensor) and weights_scaling_factor2.is_cuda: + weights_scaling_factor2 = weights_scaling_factor2.cpu() + return orig_to_quantized_weight(weight, weights_scaling_factor, quantization, + weights_scaling_factor2, block_size) + + quant_utils.to_quantized_weight = patched_to_quantized_weight + print("✓ Patch 6: to_quantized_weight → force all tensors to CPU") + + # ══════════════════════════════════════════════════════════════════════ + # Patch 7: get_weight_scaling_factor — force weight + quantizer to CPU + # Belt and suspenders: Patch 4 should handle this, but this is also + # called from other code paths. + # ══════════════════════════════════════════════════════════════════════ orig_get_weight_scaling_factor = quant_utils.get_weight_scaling_factor def patched_get_weight_scaling_factor(module, weight_name="weight"): - """Force weight and all quantizer state to CPU before computing scaling factors.""" - # Move weight to CPU if on GPU (avoids stale GPU tensor reads) - weight = getattr(module, weight_name) - if isinstance(weight, torch.Tensor) and weight.is_cuda: + weight = getattr(module, weight_name, None) + if weight is not None and isinstance(weight, torch.Tensor) and weight.is_cuda: try: - weight_cpu = weight.cpu() - # Update the module parameter so downstream code uses CPU version with torch.no_grad(): - setattr(module, weight_name, torch.nn.Parameter(weight_cpu)) - weight = weight_cpu + setattr(module, weight_name, torch.nn.Parameter(weight.cpu())) except (torch.cuda.CudaError, RuntimeError) as e: - print(f" WARNING: weight.cpu() failed for {weight_name} ({e}), using zeros") - weight = torch.zeros_like(weight, device='cpu') - - # Move quantizer amax to CPU + print(f" WARNING: weight.cpu() failed in get_weight_scaling_factor ({e})") + raise weight_quantizer = getattr(module, _quantizer_attr_names(weight_name).weight_quantizer, None) if weight_quantizer is not None: for attr in ['_amax', '_pre_quant_scale', 'global_amax', '_global_amax']: - if hasattr(weight_quantizer, attr): - val = getattr(weight_quantizer, attr) - if val is not None and isinstance(val, torch.Tensor) and val.is_cuda: - try: - setattr(weight_quantizer, attr, val.cpu()) - except (torch.cuda.CudaError, RuntimeError): - pass - + val = getattr(weight_quantizer, attr, None) + if val is not None and isinstance(val, torch.Tensor) and val.is_cuda: + try: + setattr(weight_quantizer, attr, val.cpu()) + except (torch.cuda.CudaError, RuntimeError): + pass return orig_get_weight_scaling_factor(module, weight_name) quant_utils.get_weight_scaling_factor = patched_get_weight_scaling_factor - print("✓ Patched get_weight_scaling_factor (force weight + quantizer to CPU)") + print("✓ Patch 7: get_weight_scaling_factor → force weight + quantizer to CPU") - # ── Patch 5: get_weight_scaling_factor_2 — force quantizer state to CPU ── + # ══════════════════════════════════════════════════════════════════════ + # Patch 8: get_weight_scaling_factor_2 — force quantizer to CPU + # ══════════════════════════════════════════════════════════════════════ orig_get_weight_scaling_factor_2 = quant_utils.get_weight_scaling_factor_2 def patched_get_weight_scaling_factor_2(module, weight_name="weight"): weight_quantizer = getattr(module, _quantizer_attr_names(weight_name).weight_quantizer, None) if weight_quantizer is not None: for attr in ['_amax', '_pre_quant_scale', 'global_amax', '_global_amax']: - if hasattr(weight_quantizer, attr): - val = getattr(weight_quantizer, attr) - if val is not None and isinstance(val, torch.Tensor) and val.is_cuda: - try: - setattr(weight_quantizer, attr, val.cpu()) - except (torch.cuda.CudaError, RuntimeError): - pass + val = getattr(weight_quantizer, attr, None) + if val is not None and isinstance(val, torch.Tensor) and val.is_cuda: + try: + setattr(weight_quantizer, attr, val.cpu()) + except (torch.cuda.CudaError, RuntimeError): + pass return orig_get_weight_scaling_factor_2(module, weight_name) quant_utils.get_weight_scaling_factor_2 = patched_get_weight_scaling_factor_2 - print("✓ Patched get_weight_scaling_factor_2 (force quantizer to CPU)") + print("✓ Patch 8: get_weight_scaling_factor_2 → force quantizer to CPU") def snapshot_amax_to_cpu(model, snapshot_path):