biondizzle/deepseek-v4-quant
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

8 patches covering full export chain — no more whack-a-mole

Browse Source 
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.
This commit is contained in:
biondizzle
2026-05-09 22:50:58 +00:00
parent efc111a11f
commit 07cd50e823
1 changed files with 178 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files

223
scripts/quantize_nvfp4.py
View File

@@ -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):