Fix compressor path: attn.mla_attn.compressor (not attn.compressor)

The compressor is inside mla_attn, not directly on the attention wrapper.
Debug output confirmed: layers.0.attn.mla_attn.compressor.fused_wkv_wgate.*
This commit is contained in:
2026-05-19 03:47:26 +00:00
parent 9d41419e9f
commit 4cf5b8b751
2 changed files with 7 additions and 11 deletions

View File

@@ -71,7 +71,7 @@ def make_nvfp4_mapper() -> WeightsMapper:
"compressor.gate_proj.": "compressor.wgate.",
"compressor.kv_norm.": "compressor.norm.",
"compressor.position_bias": "compressor.ape",
".self_attn.compressor.": ".attn.compressor.",
".self_attn.compressor.": ".attn.mla_attn.compressor.",
".self_attn.q_a_proj.": ".attn.wq_a.",
".self_attn.kv_proj.": ".attn.wkv.",
".self_attn.q_b_proj.": ".attn.wq_b.",
@@ -119,11 +119,11 @@ def test_mapper():
("layers.0.self_attn.sinks", "model.layers.0.attn.attn_sink"),
# Compressor (non-indexer)
("layers.0.self_attn.compressor.kv_proj.weight", "model.layers.0.attn.compressor.wkv.weight"),
("layers.0.self_attn.compressor.kv_proj.input_scale", "model.layers.0.attn.compressor.wkv.input_scale"),
("layers.0.self_attn.compressor.gate_proj.weight", "model.layers.0.attn.compressor.wgate.weight"),
("layers.0.self_attn.compressor.kv_norm.weight", "model.layers.0.attn.compressor.norm.weight"),
("layers.0.self_attn.compressor.position_bias", "model.layers.0.attn.compressor.ape"),
("layers.0.self_attn.compressor.kv_proj.weight", "model.layers.0.attn.mla_attn.compressor.wkv.weight"),
("layers.0.self_attn.compressor.kv_proj.input_scale", "model.layers.0.attn.mla_attn.compressor.wkv.input_scale"),
("layers.0.self_attn.compressor.gate_proj.weight", "model.layers.0.attn.mla_attn.compressor.wgate.weight"),
("layers.0.self_attn.compressor.kv_norm.weight", "model.layers.0.attn.mla_attn.compressor.norm.weight"),
("layers.0.self_attn.compressor.position_bias", "model.layers.0.attn.mla_attn.compressor.ape"),
# Indexer
("layers.2.self_attn.compressor.indexer.q_b_proj.weight", "model.layers.2.attn.indexer.wq_b.weight"),

View File

@@ -1444,10 +1444,6 @@ class DeepseekV4Model(nn.Module):
("compressor.fused_wkv_wgate", "compressor.wgate", 1),
]
params_dict = dict(self.named_parameters())
# DEBUG: print compressor params
for k in sorted(params_dict):
if "compressor" in k:
print(f"[DEBUG] param: {k}")
loaded_params: set[str] = set()
# TP for attention
@@ -1672,7 +1668,7 @@ def _make_deepseek_v4_nvfp4_weights_mapper() -> WeightsMapper:
"compressor.gate_proj.": "compressor.wgate.",
"compressor.kv_norm.": "compressor.norm.",
"compressor.position_bias": "compressor.ape",
".self_attn.compressor.": ".attn.compressor.",
".self_attn.compressor.": ".attn.mla_attn.compressor.",
# Attention projections
".self_attn.q_a_proj.": ".attn.wq_a.",
".self_attn.kv_proj.": ".attn.wkv.",