Add shared experts to post-quant BF16 dequant fix
Shared experts also use FlashInferCutlassNvFp4LinearKernel with broken input_scale. They need the same BF16 dequant treatment. gate_up_proj and down_proj on ffn.shared_experts.
This commit is contained in:
@@ -2409,35 +2409,51 @@ class DeepseekV4ForCausalLM(nn.Module):
|
||||
|
||||
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."""
|
||||
have set up their attributes. Dequantizes NVFP4 weights to BF16 for
|
||||
attention projections and shared experts 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
|
||||
# Attention projections
|
||||
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"):
|
||||
"weight_global_scale", "alpha", "weight_scale_inv"):
|
||||
if hasattr(mod, attr):
|
||||
try:
|
||||
delattr(mod, attr)
|
||||
except (AttributeError, TypeError):
|
||||
pass
|
||||
try: delattr(mod, attr)
|
||||
except: pass
|
||||
fixed += 1
|
||||
|
||||
# Shared expert projections (also NVFP4 with broken input_scale)
|
||||
ffn = layer.ffn
|
||||
if hasattr(ffn, 'shared_experts'):
|
||||
for proj_name in ["gate_up_proj", "down_proj"]:
|
||||
se = ffn.shared_experts
|
||||
if not hasattr(se, proj_name):
|
||||
continue
|
||||
mod = getattr(se, proj_name)
|
||||
if not hasattr(mod, "weight") or mod.weight.dtype not in (torch.uint8, torch.int8):
|
||||
continue
|
||||
self.model._dequant_nvfp4_to_bf16(mod, E2M1_LUT)
|
||||
mod.quant_method = UnquantizedLinearMethod()
|
||||
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: 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]]:
|
||||
|
||||
Reference in New Issue
Block a user