From 75844a836139211dd554719a5bd9fdcec2d7673c Mon Sep 17 00:00:00 2001 From: biondizzle Date: Mon, 18 May 2026 18:35:34 +0000 Subject: [PATCH] Post-quant fix via Dockerfile patch to process_weights_after_loading MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Forward pre-hook approach didn't work — torch.compile and model wrappers bypass hooks. Instead, patch vLLM's utils.py to call model._post_quant_fix() at the end of process_weights_after_loading. This guarantees the fix runs AFTER quant methods set up their attrs. Dockerfile now patches: model_loader/utils.py → calls model._post_quant_fix() if it exists DeepseekV4ForCausalLM._post_quant_fix() dequantizes attention NVFP4 weights to BF16 and replaces quant_method. --- Dockerfile | 20 ++++++++ vllm/patches/deepseek_v4.py | 94 ++++++++++++++----------------------- 2 files changed, 55 insertions(+), 59 deletions(-) diff --git a/Dockerfile b/Dockerfile index bcc7b5d3..d00ff017 100644 --- a/Dockerfile +++ b/Dockerfile @@ -38,6 +38,26 @@ COPY vllm/nvfp4_cutedsl.py ${VLLM_MODELS_DIR}/nvfp4_cutedsl.py RUN sed -i 's/"DeepseekV32ForCausalLM": ("deepseek_v2", "DeepseekV3ForCausalLM"),/"DeepseekV32ForCausalLM": ("deepseek_v2", "DeepseekV3ForCausalLM"),\n "DeepseekV4ForCausalLM": ("deepseek_v4", "DeepseekV4ForCausalLM"),/' \ ${VLLM_MODELS_DIR}/registry.py +# Patch process_weights_after_loading to call model._post_quant_fix() after quant setup +ARG VLLM_LOADER_DIR=/usr/local/lib/python3.12/dist-packages/vllm/model_executor/model_loader +RUN python3 -c " +import re +path = '${VLLM_LOADER_DIR}/utils.py'.replace('\$', '') +with open(path) as f: + src = f.read() +# Add _post_quant_fix() call at end of process_weights_after_loading +old = ' if model_config.quantization == \"torchao\":' +new = ''' # Custom: allow models to run post-quant-init fixes + if hasattr(model, '_post_quant_fix'): + model._post_quant_fix() + + if model_config.quantization == \"torchao\":''' +src = src.replace(old, new, 1) +with open(path, 'w') as f: + f.write(src) +print('Patched process_weights_after_loading') +" + # Verify RUN python3 -c "import torch; print(f'PyTorch {torch.__version__} OK')" && \ python3 -c "import vllm; print('vLLM OK')" && \ diff --git a/vllm/patches/deepseek_v4.py b/vllm/patches/deepseek_v4.py index 182ba2bd..41c305e7 100644 --- a/vllm/patches/deepseek_v4.py +++ b/vllm/patches/deepseek_v4.py @@ -2344,8 +2344,6 @@ class DeepseekV4ForCausalLM(nn.Module): self.model._convert_nvfp4_post_load() print(" Warming up tilelang kernels...", flush=True) self._warmup_tilelang() - print(" Registering post-quant-init fix...", flush=True) - self._register_post_quant_fix() print(" NVFP4 model ready ✓", flush=True) return loaded_params @@ -2409,65 +2407,43 @@ class DeepseekV4ForCausalLM(nn.Module): del residual, fn, hc_scale, hc_base, x, post_mix, comb_mix torch.cuda.empty_cache() + def _post_quant_fix(self) -> None: + """Called by vLLM's process_weights_after_loading AFTER quant methods + have set up their attributes. Dequantizes attention NVFP4 weights to BF16 + because FlashInferCutlassNvFp4LinearKernel uses broken input_global_scale_inv.""" + from vllm.model_executor.layers.linear import UnquantizedLinearMethod + + E2M1_LUT = torch.tensor([0, 0.5, 1, 1.5, 2, 3, 4, 6], dtype=torch.bfloat16) + fixed = 0 + for layer_idx, layer in enumerate(self.model.layers): + attn = layer.attn + for proj_name in ["fused_wqa_wkv", "wq_b", "wo_b"]: + if not hasattr(attn, proj_name): + continue + mod = getattr(attn, proj_name) + if not hasattr(mod, "weight") or mod.weight.dtype not in (torch.uint8, torch.int8): + continue + # Dequantize to BF16 + self.model._dequant_nvfp4_to_bf16(mod, E2M1_LUT) + # Replace quant method + mod.quant_method = UnquantizedLinearMethod() + # Clean up NVFP4 attributes + for attr in ("weight_scale", "weight_scale_2", "input_scale", + "input_global_scale", "input_global_scale_inv", + "weight_global_scale", "alpha", + "weight_scale_inv"): + if hasattr(mod, attr): + try: + delattr(mod, attr) + except (AttributeError, TypeError): + pass + fixed += 1 + print(f" [CLAWMINE] Post-quant fix: {fixed} attention projections → BF16 ✓", flush=True) + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: return self.model.get_expert_mapping() def _register_post_quant_fix(self) -> None: - """Register a one-shot forward pre-hook that fixes attention NVFP4 - activation scales AFTER process_weights_after_loading has run. - - process_weights_after_loading (called by vLLM's model loader AFTER - load_weights returns) sets up FlashInferCutlassNvFp4LinearKernel - with broken input_global_scale_inv from the checkpoint. This hook - runs on the first forward call (guaranteed after all init) and: - 1. Dequantizes attention NVFP4 weights to BF16 - 2. Replaces quant_method with UnquantizedLinearMethod - 3. Removes itself (one-shot) - """ - import os - # Only needed for NVFP4 quantized models - quant_method_name = type(getattr(self.model.layers[0].attn.fused_wqa_wkv, 'quant_method', None)).__name__ - if 'NvFp4' not in quant_method_name and 'nvfp4' not in quant_method_name.lower(): - print(" No NVFP4 attention fix needed (quant_method={quant_method_name})", flush=True) - return - - def _fix_attn_nvfp4(module, args): - """One-shot hook: dequant attention NVFP4 to BF16 after quant init.""" - print(" [CLAWMINE] Running post-quant-init BF16 fix...", flush=True) - from vllm.model_executor.layers.linear import UnquantizedLinearMethod - - E2M1_LUT = torch.tensor([0, 0.5, 1, 1.5, 2, 3, 4, 6], dtype=torch.bfloat16) - fixed = 0 - for layer_idx, layer in enumerate(module.layers): - attn = layer.attn - for proj_name in ["fused_wqa_wkv", "wq_b", "wo_b"]: - if not hasattr(attn, proj_name): - continue - mod = getattr(attn, proj_name) - if not hasattr(mod, "weight") or mod.weight.dtype not in (torch.uint8, torch.int8): - continue - # Dequantize to BF16 - module._dequant_nvfp4_to_bf16(mod, E2M1_LUT) - # Replace quant method (quant method already ran, won't overwrite again) - mod.quant_method = UnquantizedLinearMethod() - # Clean up NVFP4 attributes that might confuse forward - for attr in ("weight_scale", "weight_scale_2", "input_scale", - "input_global_scale", "input_global_scale_inv", - "weight_global_scale", "alpha", - "weight_scale_inv"): - if hasattr(mod, attr): - try: - delattr(mod, attr) - except (AttributeError, TypeError): - pass - fixed += 1 - - print(f" [CLAWMINE] Fixed {fixed} attention projections → BF16", flush=True) - # Remove this hook (one-shot) - module._post_quant_fix_handle.remove() - del module._post_quant_fix_handle - print(" [CLAWMINE] Post-quant-init fix done ✓", flush=True) - - handle = self.model.register_forward_pre_hook(_fix_attn_nvfp4) - self._post_quant_fix_handle = handle + """No-op — we use _post_quant_fix() called from process_weights_after_loading.""" + pass