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

S11: Fixed substr mapping, stacking, suffix, and o_a_proj - loads weights but attention forward uses FP8 einsum incompatible with NVFP4

Browse Source
This commit is contained in:
biondizzle
2026-05-10 17:45:53 +00:00
parent 6fd03a0aa0
commit db16be8e5d
3 changed files with 3631 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

146
patches/deepseek_v4.py
View File

@@ -1473,6 +1473,103 @@ class DeepseekV4Model(nn.Module):
# Skip them silently.
continue
param = params_dict[name]
# Handle bf16 → uint8 mismatch for o_a_proj:
# modelopt didn't quantize o_a_proj (bf16, no scales),
# but ModelOptNvFp4Config creates wo_a with NVFP4 quant
# (uint8 weight + scales). We quantize the bf16 weight
# to NVFP4 at load time so the layer runs in NVFP4 path.
if (name.endswith(".weight")
and loaded_weight.dtype != torch.uint8
and param.data.dtype == torch.uint8):
# Quantize bf16 → NVFP4 (E2M1 packed uint8 + scales)
w_bf16 = loaded_weight
out_dim, in_dim = w_bf16.shape
block_size = 16
assert in_dim % block_size == 0
n_blocks = in_dim // block_size
# Reshape into blocks
w_blocks = w_bf16.reshape(out_dim, n_blocks, block_size)
# Compute per-block amax
amax = w_blocks.abs().amax(dim=-1) # [out, n_blocks]
# Global scale (weight_scale_2): max amax / (6.0 * 448.0)
global_amax = amax.max()
# Use 448.0 as the max e4m3 value for scale computation
weight_scale_2_val = global_amax / (6.0 * 448.0)
weight_scale_2 = weight_scale_2_val.to(torch.float32)
# Per-block scale (weight_scale): fp8 e4m3
# block_scale = amax / (6.0 * weight_scale_2)
block_scale = amax / (6.0 * weight_scale_2_val)
# Clamp to fp8 e4m3 range and cast
block_scale = block_scale.clamp(min=0, max=448.0)
weight_scale = block_scale.to(torch.float8_e4m3fn)
# Quantize to FP4 (E2M1)
# E2M1 LUT: 0, 0.5, 1, 1.5, 2, 3, 4, 6 (positive)
FP4_POS = torch.tensor(
[0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0],
dtype=torch.float32, device=w_bf16.device,
)
# For each block, dequantize the block scale from fp8
block_scale_f32 = weight_scale.to(torch.float32)
# Scale the weight values: normalized = w / (block_scale * weight_scale_2)
# We need to find the nearest FP4 value
scaled = w_blocks / (block_scale_f32.unsqueeze(-1) * weight_scale_2_val)
# Find nearest FP4 index (0-7 for magnitude)
# Use absolute value for matching, then apply sign
scaled_abs = scaled.abs()
# Find closest FP4 value
diff = (scaled_abs.unsqueeze(-1) - FP4_POS).abs()
fp4_idx = diff.argmin(dim=-1) # [out, n_blocks, block_size]
# Apply sign: negative values get bit 3 set
sign = (scaled < 0).int()
fp4_val = (sign << 3) | fp4_idx.int()
# Pack: 2 FP4 values per uint8 byte
# Even positions → lower nibble, Odd → upper nibble
fp4_flat = fp4_val.reshape(out_dim, -1) # [out, in_dim]
assert fp4_flat.shape[1] % 2 == 0
even = fp4_flat[:, 0::2] # lower nibble
odd = fp4_flat[:, 1::2] # upper nibble
packed = (odd << 4) | even
weight_packed = packed.to(torch.uint8)
# Reshape weight_scale to [out, n_blocks]
weight_scale_2d = weight_scale.reshape(out_dim, n_blocks)
# Load the quantized weight into the uint8 param
weight_loader = param.weight_loader
weight_loader(param, weight_packed)
loaded_params.add(name)
# Load scales into sibling params
base = name.rsplit(".", 1)[0]
# weight_scale
ws_name = f"{base}.weight_scale"
if ws_name in params_dict:
ws_param = params_dict[ws_name]
ws_loader = getattr(ws_param, "weight_loader", default_weight_loader)
ws_loader(ws_param, weight_scale_2d)
loaded_params.add(ws_name)
# weight_scale_2
ws2_name = f"{base}.weight_scale_2"
if ws2_name in params_dict:
ws2_param = params_dict[ws2_name]
ws2_loader = getattr(ws2_param, "weight_loader", default_weight_loader)
ws2_loader(ws2_param, weight_scale_2.reshape(1))
loaded_params.add(ws2_name)
# input_scale: use 1.0 default (dynamic quant)
is_name = f"{base}.input_scale"
if is_name in params_dict:
is_param = params_dict[is_name]
is_loader = getattr(is_param, "weight_loader", default_weight_loader)
is_loader(is_param, torch.tensor(1.0, dtype=torch.float32))
loaded_params.add(is_name)
continue
weight_loader = getattr(
param, "weight_loader", default_weight_loader
)
@@ -1567,35 +1664,19 @@ def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper:
# Must match ORIGINAL checkpoint key names (before substr renaming).
fused_skip_regex = {
# Compressor projections → fused_wkv_wgate (stacked)
# Compressor uses UnquantizedLinearMethod (quant_config=None),
# so it only has a bf16 weight param — no scale params registered.
# We unpack the NVFP4 uint8 weights to bf16 at load time.
re.compile(r"\.compressor\.kv_proj\.weight_scale$"): None,
re.compile(r"\.compressor\.gate_proj\.weight_scale$"): None,
re.compile(r"\.compressor\.kv_proj\.weight_scale_2$"): None,
re.compile(r"\.compressor\.gate_proj\.weight_scale_2$"): None,
re.compile(r"\.compressor\.kv_proj\.input_scale$"): None,
re.compile(r"\.compressor\.gate_proj\.input_scale$"): None,
# Attention projections → fused_wqa_wkv (stacked)
re.compile(r"\.self_attn\.kv_proj\.weight_scale$"): None,
re.compile(r"\.self_attn\.q_a_proj\.weight_scale$"): None,
re.compile(r"\.self_attn\.q_b_proj\.weight_scale$"): None,
re.compile(r"\.self_attn\.o_a_proj\.weight_scale$"): None,
re.compile(r"\.self_attn\.o_b_proj\.weight_scale$"): None,
re.compile(r"\.self_attn\.kv_proj\.weight_scale_2$"): None,
re.compile(r"\.self_attn\.q_a_proj\.weight_scale_2$"): None,
re.compile(r"\.self_attn\.q_b_proj\.weight_scale_2$"): None,
re.compile(r"\.self_attn\.o_a_proj\.weight_scale_2$"): None,
re.compile(r"\.self_attn\.o_b_proj\.weight_scale_2$"): None,
re.compile(r"\.self_attn\.kv_proj\.input_scale$"): None,
re.compile(r"\.self_attn\.q_a_proj\.input_scale$"): None,
re.compile(r"\.self_attn\.q_b_proj\.input_scale$"): None,
re.compile(r"\.self_attn\.o_a_proj\.input_scale$"): None,
re.compile(r"\.self_attn\.o_b_proj\.input_scale$"): None,
# Shared expert gate_proj/up_proj → gate_up_proj (stacked)
re.compile(r"\.shared_experts\.gate_proj\.weight_scale$"): None,
re.compile(r"\.shared_experts\.up_proj\.weight_scale$"): None,
re.compile(r"\.shared_experts\.gate_proj\.weight_scale_2$"): None,
re.compile(r"\.shared_experts\.up_proj\.weight_scale_2$"): None,
re.compile(r"\.shared_experts\.gate_proj\.input_scale$"): None,
re.compile(r"\.shared_experts\.up_proj\.input_scale$"): None,
# Note: attention and shared expert scale tensors are NO LONGER
# skipped. After fixing substr mappings, they correctly map to the
# model's NVFP4 scale parameters (fused_wqa_wkv, wq_b, wo_a,
# wo_b, gate_up_proj). They load via the stacking logic.
}
# Routed expert projections: gate_proj→w1, up_proj→w3, down_proj→w2
# Regex (not substr) to match ONLY .experts.N. — not .shared_experts.
@@ -1620,7 +1701,6 @@ def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper:
},
orig_to_new_regex=merged_regex,
orig_to_new_suffix={
"head.weight": "lm_head.weight",
"embed.weight": "embed_tokens.weight",
".ffn.gate.bias": ".ffn.gate.e_score_correction_bias",
},
@@ -1628,16 +1708,16 @@ def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper:
".attn.compressor.": ".attn.mla_attn.compressor.",
".shared_experts.w2": ".shared_experts.down_proj",
# ── ModelOpt NVFP4 substr patches ──
# Attention: self_attn → attn.mla_attn
".self_attn.q_a_proj.": ".attn.mla_attn.wq_a.",
".self_attn.q_b_proj.": ".attn.mla_attn.wq_b.",
".self_attn.q_a_norm.": ".attn.mla_attn.q_norm.",
".self_attn.o_a_proj.": ".attn.mla_attn.wo_a.",
".self_attn.o_b_proj.": ".attn.mla_attn.wo_b.",
".self_attn.sinks": ".attn.mla_attn.attn_sink",
# Attention: self_attn → attn (projections at attn level, not mla_attn)
".self_attn.q_a_proj.": ".attn.wq_a.",
".self_attn.q_b_proj.": ".attn.wq_b.",
".self_attn.q_a_norm.": ".attn.q_norm.",
".self_attn.o_a_proj.": ".attn.wo_a.",
".self_attn.o_b_proj.": ".attn.wo_b.",
".self_attn.sinks": ".attn.attn_sink",
# kv_proj → wkv (for stacking into fused_wqa_wkv)
".self_attn.kv_proj.": ".attn.mla_attn.wkv.",
".self_attn.kv_norm.": ".attn.mla_attn.kv_norm.",
".self_attn.kv_proj.": ".attn.wkv.",
".self_attn.kv_norm.": ".attn.kv_norm.",
# kv_norm is at attention level, not compressor/mla_attn level in vllm
# Must come before the general compressor mapping
".self_attn.compressor.kv_norm.": ".attn.kv_norm.",

1719
patches/deepseek_v4.py.bak Normal file
View File

File diff suppressed because it is too large Load Diff

1799
patches/deepseek_v4.py.s11 Normal file
View File

File diff suppressed because it is too large Load Diff