diff --git a/tests/test_nvfp4_mapper.py b/tests/test_nvfp4_mapper.py new file mode 100644 index 00000000..ef23cee2 --- /dev/null +++ b/tests/test_nvfp4_mapper.py @@ -0,0 +1,173 @@ +#!/usr/bin/env python3 +"""Unit test for the NVFP4 weights mapper. + +Validates that checkpoint key names from our ModelOpt-quantized +DeepSeek-V4-Pro checkpoint are correctly mapped to vLLM model +parameter names. + +This can run WITHOUT vLLM or CUDA — it only tests the mapper logic. +""" + +import re +import sys +from typing import Optional + + +class WeightsMapper: + """Simplified WeightsMapper for testing.""" + + def __init__( + self, + orig_to_new_prefix: Optional[dict] = None, + orig_to_new_regex: Optional[dict] = None, + orig_to_new_suffix: Optional[dict] = None, + orig_to_new_substr: Optional[dict] = None, + ): + self.prefix_map = orig_to_new_prefix or {} + self.regex_map = orig_to_new_regex or {} + self.suffix_map = orig_to_new_suffix or {} + self.substr_map = orig_to_new_substr or {} + + def map_name(self, name: str) -> str: + # 1. Prefix + for old, new in self.prefix_map.items(): + if name.startswith(old): + name = new + name[len(old):] + break + + # 2. Regex + for pattern, replacement in self.regex_map.items(): + name = pattern.sub(replacement, name) + + # 3. Suffix + for old, new in self.suffix_map.items(): + if name.endswith(old): + name = name[: -len(old)] + new + break + + # 4. Substr (ordered dict — specific before general) + for old, new in self.substr_map.items(): + if old in name: + name = name.replace(old, new, 1) + + return name + + +def _make_deepseek_v4_nvfp4_weights_mapper() -> WeightsMapper: + """Exact copy of the mapper from deepseek_v4.py.""" + expert_rename_regex = { + re.compile(r"(\.experts\.\d+\.)gate_proj\."): r"\1w1.", + re.compile(r"(\.experts\.\d+\.)up_proj\."): r"\1w3.", + re.compile(r"(\.experts\.\d+\.)down_proj\."): r"\1w2.", + } + + suffix_renames = {} + + substr_renames = { + ".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.", + ".self_attn.o_a_proj.": ".attn.wo_a.", + ".self_attn.o_b_proj.": ".attn.wo_b.", + ".self_attn.q_a_norm.": ".attn.q_a_norm.", + ".self_attn.kv_norm.": ".attn.kv_norm.", + ".self_attn.sinks": ".attn.sinks", + ".mlp.shared_experts.gate_proj.": ".ffn.shared_experts.w1.", + ".mlp.shared_experts.up_proj.": ".ffn.shared_experts.w3.", + ".mlp.shared_experts.down_proj.": ".ffn.shared_experts.down_proj.", + ".mlp.": ".ffn.", + ".self_attn.": ".attn.", + } + + return WeightsMapper( + orig_to_new_prefix={ + "layers.": "model.layers.", + "embed_tokens.": "model.embed_tokens.", + "norm.": "model.norm.", + "hc_head": "model.hc_head", + "mtp.": "model.mtp.", + }, + orig_to_new_regex=expert_rename_regex, + orig_to_new_suffix=suffix_renames, + orig_to_new_substr=substr_renames, + ) + + +TEST_CASES = [ + # Embedding & top-level + ("embed_tokens.weight", "model.embed_tokens.weight"), + ("norm.weight", "model.norm.weight"), + ("hc_head.hc_fn", "model.hc_head.hc_fn"), + ("hc_head.hc_base", "model.hc_head.hc_base"), + ("hc_head.hc_scale", "model.hc_head.hc_scale"), + ("lm_head.weight", "lm_head.weight"), + + # Attention — self_attn → attn + ("layers.0.self_attn.q_a_proj.weight", "model.layers.0.attn.wq_a.weight"), + ("layers.0.self_attn.q_a_proj.input_scale", "model.layers.0.attn.wq_a.input_scale"), + ("layers.0.self_attn.kv_proj.weight", "model.layers.0.attn.wkv.weight"), + ("layers.0.self_attn.kv_proj.input_scale", "model.layers.0.attn.wkv.input_scale"), + ("layers.0.self_attn.q_b_proj.weight", "model.layers.0.attn.wq_b.weight"), + ("layers.0.self_attn.o_a_proj.weight", "model.layers.0.attn.wo_a.weight"), + ("layers.0.self_attn.o_b_proj.weight", "model.layers.0.attn.wo_b.weight"), + ("layers.0.self_attn.o_b_proj.input_scale", "model.layers.0.attn.wo_b.input_scale"), + ("layers.0.self_attn.q_a_norm.weight", "model.layers.0.attn.q_a_norm.weight"), + ("layers.0.self_attn.kv_norm.weight", "model.layers.0.attn.kv_norm.weight"), + ("layers.0.self_attn.sinks", "model.layers.0.attn.sinks"), + + # Attention — compressor (inside self_attn) + ("layers.0.self_attn.compressor.kv_proj.weight", "model.layers.0.attn.mla_attn.compressor.kv_proj.weight"), + ("layers.0.self_attn.compressor.kv_proj.input_scale", "model.layers.0.attn.mla_attn.compressor.kv_proj.input_scale"), + ("layers.0.self_attn.compressor.gate_proj.weight", "model.layers.0.attn.mla_attn.compressor.gate_proj.weight"), + ("layers.0.self_attn.compressor.gate_proj.input_scale", "model.layers.0.attn.mla_attn.compressor.gate_proj.input_scale"), + ("layers.0.self_attn.compressor.kv_norm.weight", "model.layers.0.attn.mla_attn.compressor.kv_norm.weight"), + ("layers.0.self_attn.compressor.position_bias", "model.layers.0.attn.mla_attn.compressor.position_bias"), + + # MoE gate + ("layers.0.mlp.gate.tid2eid", "model.layers.0.ffn.gate.tid2eid"), + ("layers.0.mlp.gate.weight", "model.layers.0.ffn.gate.weight"), + + # Expert weights — gate_proj → w1, up_proj → w3, down_proj → w2 + ("layers.0.mlp.experts.0.gate_proj.weight", "model.layers.0.ffn.experts.0.w1.weight"), + ("layers.0.mlp.experts.0.up_proj.weight", "model.layers.0.ffn.experts.0.w3.weight"), + ("layers.0.mlp.experts.0.down_proj.weight", "model.layers.0.ffn.experts.0.w2.weight"), + ("layers.0.mlp.experts.0.gate_proj.input_scale", "model.layers.0.ffn.experts.0.w1.input_scale"), + ("layers.0.mlp.experts.0.gate_proj.weight_scale", "model.layers.0.ffn.experts.0.w1.weight_scale"), + ("layers.0.mlp.experts.0.gate_proj.weight_scale_2", "model.layers.0.ffn.experts.0.w1.weight_scale_2"), + ("layers.0.mlp.experts.255.down_proj.weight", "model.layers.0.ffn.experts.255.w2.weight"), + + # Shared experts — gate_proj → w1, up_proj → w3, down_proj stays + ("layers.0.mlp.shared_experts.gate_proj.weight", "model.layers.0.ffn.shared_experts.w1.weight"), + ("layers.0.mlp.shared_experts.up_proj.weight", "model.layers.0.ffn.shared_experts.w3.weight"), + ("layers.0.mlp.shared_experts.down_proj.weight", "model.layers.0.ffn.shared_experts.down_proj.weight"), + ("layers.0.mlp.shared_experts.gate_proj.input_scale", "model.layers.0.ffn.shared_experts.w1.input_scale"), + ("layers.0.mlp.shared_experts.down_proj.weight_scale", "model.layers.0.ffn.shared_experts.down_proj.weight_scale"), + ("layers.0.mlp.shared_experts.down_proj.weight_scale_2", "model.layers.0.ffn.shared_experts.down_proj.weight_scale_2"), + + # Layer norm + ("layers.0.post_attention_layernorm.weight", "model.layers.0.post_attention_layernorm.weight"), +] + + +def main(): + mapper = _make_deepseek_v4_nvfp4_weights_mapper() + passed = 0 + failed = 0 + + for ckpt_key, expected in TEST_CASES: + result = mapper.map_name(ckpt_key) + if result == expected: + passed += 1 + else: + failed += 1 + print(f"FAIL: {ckpt_key}") + print(f" expected: {expected}") + print(f" got: {result}") + + print(f"\n{passed} passed, {failed} failed") + return 0 if failed == 0 else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/vllm/patches/deepseek_v4.py b/vllm/patches/deepseek_v4.py index 73037577..77f1d227 100644 --- a/vllm/patches/deepseek_v4.py +++ b/vllm/patches/deepseek_v4.py @@ -1648,12 +1648,14 @@ def _make_deepseek_v4_nvfp4_weights_mapper() -> WeightsMapper: } suffix_renames = { - "head.weight": "lm_head.weight", - "embed.weight": "embed_tokens.weight", + # The NVFP4 checkpoint already uses lm_head / embed_tokens directly, + # no suffix renames needed (unlike the MXFP4 upstream format). } # NOTE: specific renames MUST come before general ones (applied in order) substr_renames = { + # Attention compressor (MUST come before .self_attn. → .attn.) + ".self_attn.compressor.": ".attn.mla_attn.compressor.", # Attention projections (specific before .self_attn. → .attn.) ".self_attn.q_a_proj.": ".attn.wq_a.", ".self_attn.kv_proj.": ".attn.wkv.", @@ -1663,7 +1665,6 @@ def _make_deepseek_v4_nvfp4_weights_mapper() -> WeightsMapper: ".self_attn.q_a_norm.": ".attn.q_a_norm.", ".self_attn.kv_norm.": ".attn.kv_norm.", ".self_attn.sinks": ".attn.sinks", - ".attn.compressor.": ".attn.mla_attn.compressor.", # Shared expert projections (specific before .mlp. → .ffn.) ".mlp.shared_experts.gate_proj.": ".ffn.shared_experts.w1.", ".mlp.shared_experts.up_proj.": ".ffn.shared_experts.w3.", @@ -1675,7 +1676,11 @@ def _make_deepseek_v4_nvfp4_weights_mapper() -> WeightsMapper: return WeightsMapper( orig_to_new_prefix={ - # Keep model. prefix — model params use it (e.g. model.embed_tokens.weight) + "layers.": "model.layers.", + "embed_tokens.": "model.embed_tokens.", + "norm.": "model.norm.", + "hc_head": "model.hc_head", + "mtp.": "model.mtp.", }, orig_to_new_regex=expert_rename_regex, orig_to_new_suffix=suffix_renames,