From 2835cb040b9d6d4fe2f1af533c57b9e94418a002 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Mon, 18 May 2026 16:43:44 +0000 Subject: [PATCH] Fix input_scale BEFORE process_weights_after_loading runs Instead of dequantizing to BF16 (which gets overwritten by process_weights_after_loading), fix the input_scale parameter on the module before the quant method reads it. The quant method computes input_global_scale_inv = input_scale.max(), so fixing input_scale propagates the correct activation scale. Computes correct input_scale by temporarily dequantizing weight to BF16, running warmup forward, and computing act_amax. input_scale = 1/(act_amax * headroom). --- vllm/patches/deepseek_v4.py | 87 +++++++++++++++++++++++++++++-------- 1 file changed, 69 insertions(+), 18 deletions(-) diff --git a/vllm/patches/deepseek_v4.py b/vllm/patches/deepseek_v4.py index 01e618f2..36d53978 100644 --- a/vllm/patches/deepseek_v4.py +++ b/vllm/patches/deepseek_v4.py @@ -1685,37 +1685,26 @@ class DeepseekV4Model(nn.Module): def _convert_nvfp4_post_load(self): """Post-load conversion of NVFP4 weights for vLLM compatibility. - All attention NVFP4 projections (except wo_a) are dequantized to BF16. - The checkpoint input_scale values cause NaN during activation quantization - in FlashInferCutlassNvFp4LinearKernel. BF16 bypasses this entirely. + Fixes the attention input_scale values BEFORE + process_weights_after_loading runs. The checkpoint input_scale + values are wrong and cause NaN during activation quantization. + We compute correct values by dequantizing to BF16 temporarily + and running a warmup forward. wo_a is converted to FP8 for fp8_einsum (no input_scale needed). Compressor weights are reconstructed from checkpoint sub-weights. """ - bf16_proj_names = {"fused_wqa_wkv", "wq_b", "wo_b"} fp8_proj_names = {"wo_a"} - bf16_converted = 0 fp8_converted = 0 compressor_converted = 0 + input_scale_fixes = 0 _shard_index = self._build_shard_index("/model") if os.path.isdir("/model") else None from tqdm import tqdm - for layer_idx, layer in tqdm(enumerate(self.layers), total=len(self.layers), desc=" (upcast)NVFP4→BF16 attn projs", unit="layer"): + for layer_idx, layer in tqdm(enumerate(self.layers), total=len(self.layers), desc=" (fix)NVFP4 attn input_scale", unit="layer"): attn = layer.attn - # BF16 dequantization: attention projections (except wo_a) - for proj_name in bf16_proj_names: - if not hasattr(attn, proj_name): - continue - mod = getattr(attn, proj_name) - if not hasattr(mod, "weight"): - continue - if mod.weight.dtype in (torch.uint8, torch.int8): - E2M1_LUT = torch.tensor([0, 0.5, 1, 1.5, 2, 3, 4, 6], dtype=torch.bfloat16) - self._dequant_nvfp4_to_bf16(mod, E2M1_LUT) - bf16_converted += 1 - # FP8 conversion: wo_a (used by fp8_einsum, no input_scale) FP8_MAX = torch.finfo(torch.float8_e4m3fn).max for proj_name in fp8_proj_names: @@ -1729,6 +1718,68 @@ class DeepseekV4Model(nn.Module): self._convert_nvfp4_to_fp8(mod, E2M1_LUT, FP8_MAX) fp8_converted += 1 + # Fix input_scale for attention NVFP4 projections + # process_weights_after_loading reads input_scale and computes + # input_global_scale_inv = 1/input_scale. By fixing input_scale + # here, the quant method will propagate the correct value. + 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, "input_scale"): + continue + if not hasattr(mod, "weight") or mod.weight.dtype not in (torch.uint8, torch.int8): + continue + + # Temporarily dequantize weight to BF16 for warmup + E2M1_LUT = torch.tensor([0, 0.5, 1, 1.5, 2, 3, 4, 6], dtype=torch.bfloat16) + w_uint8 = mod.weight.data + w_bf16_unpacked = self._unpack_nvfp4_to_bf16(w_uint8, E2M1_LUT, w_uint8.device) + if hasattr(mod, "weight_scale") and hasattr(mod, "weight_scale_2"): + block_scale = self._block_scale_to_float32(mod.weight_scale.data) + if block_scale.dim() == 2 and w_bf16_unpacked.dim() == 2: + block_size = w_bf16_unpacked.shape[1] // block_scale.shape[1] + block_scale_expanded = block_scale.unsqueeze(-1).expand(-1, -1, block_size).reshape(w_bf16_unpacked.shape) + else: + block_scale_expanded = block_scale + global_scale = mod.weight_scale_2.data.max().item() + w_bf16_dequant = (w_bf16_unpacked.float() * block_scale_expanded * global_scale).to(torch.bfloat16) + else: + w_bf16_dequant = w_bf16_unpacked + + # Compute correct input_scale from warmup + with torch.no_grad(): + in_features = w_bf16_dequant.shape[-1] + dummy_input = torch.randn(256, in_features, dtype=torch.bfloat16, device=mod.weight.device) * 2.0 + ref_output = torch.nn.functional.linear(dummy_input, w_bf16_dequant) + act_amax = ref_output.amax().item() + del w_bf16_unpacked, w_bf16_dequant, ref_output + + # input_scale should be 1/(amax * headroom) — this is the + # activation global scale that maps activations to FP4 range. + # process_weights_after_loading computes: + # input_global_scale_inv = input_scale.max() + # input_global_scale = 1 / input_global_scale_inv + headroom = 1.2 + new_input_scale = 1.0 / (act_amax * headroom) if act_amax > 0 else mod.input_scale.data + + if layer_idx == 0: + old_input_scale = mod.input_scale.data.item() if mod.input_scale.data.numel() == 1 else mod.input_scale.data.max().item() + print(f"[CLAWMINE] Layer 0: {proj_name} input_scale: {old_input_scale:.8f} → {new_input_scale:.8f} (act_amax={act_amax:.4f})") + + mod.input_scale = torch.nn.Parameter( + torch.tensor([new_input_scale] * mod.input_scale.data.numel(), dtype=mod.input_scale.data.dtype, device=mod.input_scale.data.device), + requires_grad=False + ) + input_scale_fixes += 1 + + _shard_index = self._build_shard_index("/model") if os.path.isdir("/model") else None + + from tqdm import tqdm + for layer_idx, layer in tqdm(enumerate(self.layers), total=len(self.layers), desc=" (upcast)NVFP4→BF16 attn projs", unit="layer"): + attn = layer.attn + + # Compressor: still needs BF16 reconstruction mla_attn = getattr(attn, "mla_attn", None) if mla_attn is not None: