fix(model): stop building a DeepSeek query norm HF does not define

[bzantium]

What does this PR do ?

Fixes #5261

Stop building a per-layer query normalization that the HF DeepSeek architecture does not define, and stop the resulting unmapped parameter from crashing the weight load.

Bug and root cause

MCore derives Q and KV normalization from a single qk_layernorm flag. DeepSeek needs it enabled for kv_a_layernorm, which every checkpoint ships, so both bridges set provider.qk_layernorm = True unconditionally.

When q_lora_rank is None there is no query LoRA, and that same flag makes MCore fuse a query normalization into linear_q_proj instead:

[TE] q_lora_rank=None -> linear_q_proj=TELayerNormColumnParallelLinear, linear_q_up_proj=IdentityOp
[TE] q_lora_rank=1536 -> linear_q_proj=IdentityOp,                      linear_q_up_proj=TELayerNormColumnParallelLinear

Both lines come from the same resolver, mla_qk_norm_config.py::_resolve_mla_qk_layernorm, and this repo maps the second to self_attn.q_a_layernorm.weight (models/deepseek/common.py:57). HF defines no query-side norm when there is no query LoRA — modeling_deepseek_v3.py:384-389 builds a bare q_proj — so nothing can be loaded into the first.

Dumping the converted model showed the parameter sitting at its init while every neighbouring norm took its checkpoint value (source norms randomized to uniform(0.5, 1.5) to tell the two apart):

input_layernorm.weight                              requires_grad=True  min=0.5000 max=1.5000
self_attention.linear_q_proj.layer_norm_weight      requires_grad=True  min=1.0000 max=1.0000
self_attention.linear_kv_up_proj.layer_norm_weight  requires_grad=True  min=0.5625 max=1.4219
mlp.linear_fc1.layer_norm_weight                    requires_grad=True  min=0.5000 max=1.4922

requires_grad=True, so training optimized a degree of freedom the source model does not have, and export dropped it silently. On top of that the missing mapping produced a None conversion-task slot, which the load path dereferenced.

Changelog

• Add MLASelfAttentionWithoutQueryNorm, which keeps the KV norm and drops the query norm when q_lora_rank is None, plus get_deepseek_decoder_block_spec to install it.
• Point both DeepSeek bridges at that spec builder instead of get_gpt_decoder_block_spec.
• Mirror get_gpt_decoder_block_spec parameter for parameter, including vp_stage and pp_rank. GPTModelProvider.provide() inspects the callable's signature and only forwards vp_stage when it is declared, so omitting it routed interleaved pipeline parallelism into the layer-offset helper without a virtual stage.
• Share the attention substitution as replace_mla_self_attention and add a mtp_layer_spec_transform provider field, so a standalone MTP stage applies it on the fallback path where mtp_block_spec re-derives the layer spec from MCore.
• Require Transformer Engine when there is no query LoRA, with a message that names the setting at fault instead of surfacing an MCore internal error.
• Raise on a Megatron parameter with no mapping at all, naming the parameter, instead of warning and continuing.
• Add unit tests for each of the above, and a q_lora_rank=None state-dict round trip in the DeepSeek functional group. The existing toy fixture uses q_lora_rank=512, which is the unaffected control, so nothing in the suite reached the changed branch.

Scope of the None task slot

build_conversion_tasks has two producers of a None slot, and they are not the same
condition. A Megatron parameter with no registry entry is a wrong model: the parameter
stays at its initial value on import and is dropped on export. That now raises by name and
points at _HCAlphaSecondaryMapping as the way to declare a deliberate no-op.

The second producer is an HF checkpoint that does not carry a weight the mapping names. It
is benign and pre-existing, so the consumer-side skips stay. Removing them would turn a
missing optional weight into a crash in shared conversion code for every model family.

Backend support

The local backend cannot express this architecture, and could not before this PR.
_resolve_mla_qk_layernorm builds linear_q_proj from the backend's fused norm+linear
implementation whenever q_lora_rank is None, and LocalSpecProvider has none, so
_require_linear raises. Keeping the query norm instead trips _raise_unused_q_norm. Both
branches are closed, and DeepSeek needs qk_layernorm on for kv_a_layernorm.

Reproducing MCore's resolver inside the bridge to work around this would fork logic that
belongs upstream, so the bridge fails early with an explicit Transformer Engine
requirement.

qk_layernorm=False is not an alternative: it also removes linear_kv_up_proj.layer_norm_weight, which common.py:40 maps to a real HF weight. Two spec-level workarounds are closed as well — _reject_disabled_norm rejects a fused KV projection while QK norm is disabled, and _mla_fused_linear_or_default overrides an explicitly non-fused linear_q_proj back to the fused class. Overriding _resolve_qk_norm_config on the attention module is the one hook that composes, and it needs no Megatron-LM change (upstream main is byte-identical to the pinned submodule here).

Validation

Model construction, TP=1, tiny DeepseekV3ForCausalLM checkpoints:

checkpoint	linear_q_proj norm	linear_kv_up_proj norm	MTP params	MTP q_proj norm
q_lora_rank=None	0 (was 2)	2	–	–
q_lora_rank=64 (control)	0	2	–	–
q_lora_rank=None, MTP=1	0	3	22	0

No No mapping found for megatron_param warnings remain. The MTP row above is the shared-stage layout, where mtp_block_spec reuses the corrected decoder spec. A standalone MTP stage owns no decoder layers and re-derives its spec straight from MCore, which is why the substitution is also registered as mtp_layer_spec_transform rather than relying on the shared spec.

Logit parity, compare_hf_and_megatron/compare.py, TP=1:

variant	next token	cosine
q_lora_rank=64 (control)	match	0.999930
q_lora_rank=None, before this PR	match	0.999903
q_lora_rank=None, after this PR	match	0.999909

Removing the parameter does not move inference. The defect was never an inference-time divergence — it was a phantom trainable parameter — and an earlier revision of #5261 that claimed otherwise has been corrected.

Conversion-loop regression, against the parent commit:

python -m pytest tests/unit_tests/models/test_model_bridge.py -q -k skips_unmapped_task_slots

• Before: AttributeError: 'NoneType' object has no attribute 'megatron_module' and ... 'param_weight', 2 failed.
• After: 2 passed.

Unit tests:

python -m pytest tests/unit_tests/models/deepseek/ tests/unit_tests/models/test_model_bridge.py -q
94 passed

Run against the Megatron-Core commit pinned by this branch, not the submodule checkout,
so the QK-norm resolver under discussion is the one exercised.

The whole of tests/unit_tests/models/ was run on this branch and on the parent commit in
the same environment, to check that raising on an unmapped parameter does not disturb other
model families:

	failed	passed	skipped	errors
parent bd25f3ed3	66	2800	36	206
this branch	66	2809	36	206

Same failures and errors, nine more passes, which are the tests added here. The pre-existing
failures are environment-related and unrelated to this change.

State-dict round trip, one GPU:

python -m pytest tests/functional_tests/test_groups/models/deepseek/test_deepseek_conversion.py -k round_trip -q
1 passed

HF has q_proj.weight and kv_a_layernorm.weight and no query LoRA. Megatron has
linear_q_proj.weight and linear_kv_up_proj.layer_norm_weight and no
linear_q_proj.layer_norm_weight. Every original weight survives HF to Megatron to HF
bit-exactly, and no q_a_proj or q_a_layernorm key is invented. The structural
expectations come from the HF architecture, not from the conversion registry.

Coverage note: the unit tests reach the resolver and both affected bridges, but no test loads a real affected checkpoint. kakaocorp/kanana-2-30b-a3b-thinking is too large for the 2-GPU functional budget, and deepseek-ai/DeepSeek-V2-Lite cannot be loaded at all under the pinned transformers — its trust_remote_code modeling file calls DynamicCache.from_legacy_cache, which no longer exists. The model-construction and parity evidence above therefore comes from tiny DeepseekV3ForCausalLM checkpoints built to the same shape (q_proj present, q_a_proj / q_a_layernorm absent) rather than from CI. The round trip is now committed as TestDeepSeekWithoutQueryLoRA on a generated toy checkpoint, so the changed branch has end-to-end coverage even though no real affected checkpoint fits the budget.

Static checks: ruff check src tests passed, ruff format --check clean on changed files, git diff --check passed.

Environment caveat: these runs used a CUDA 12.9 / A100 container built from uv.lock, not the CI megatron-bridge image — nvidia-resiliency-ext publishes glibc 2.39 wheels only (this host is glibc 2.35), the CUDA source-build extras do not compile here, and Transformer Engine had to be built from the pinned commit with NVTE_WITH_NCCL_EP=0 for sm80. gradient_accumulation_fusion was disabled for the parity runs because APEX is absent; it selects the weight-gradient kernel and does not affect the forward pass. The toy checkpoints are 2-layer and randomly initialized, so they do not speak to drift across a deep trained stack. CI on the real image remains the authoritative run.

Scope

• No public API, CLI, config, dependency, lockfile, workflow, or MCore changes.
• Model construction changes only for MLA with q_lora_rank is None; the q_lora_rank path resolves identically to before.
• Conversion behaviour changes for one condition: a Megatron parameter with no registry entry now raises by name instead of warning and continuing. _HCAlphaSecondaryMapping in deepseek_v4_bridge.py is the existing way to declare a deliberate no-op. The separate No mapping found for global_name warning that test_fp8_param_export.py asserts is a different site and is untouched, as are the consumer-side skips for a weight the HF checkpoint does not carry.
• Checkpoints trained before this change still contain the extra parameter; loading them into the corrected model is not addressed here.

GitHub Actions CI

See the CI section in the Contributing doc for how to trigger the CI. A Nvidia developer will need to approve and trigger the CI for external contributors.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?
• Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  • Reviewer: Does the PR have correct import guards for all optional libraries?

Additional Information

• Fixes [bug] DeepSeek MLA with q_lora_rank=null builds a trainable Q-norm the HF architecture does not define #5261

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[yaoyu-33]

The central TE/TP=1/no-MTP query-norm substitution is independently validated, including a bit-exact affected-branch HF→MCore→HF round trip. I am requesting changes for the remaining P1 issues: preserve and test virtual-stage routing; retain strict failure for genuinely unmapped global parameters instead of silently skipping them; support or intentionally gate the optional local backend; and commit a real affected state-dict/round-trip regression test. The standalone-MTP fallback is also still open as P2. Please update the linked inline threads and request re-review on the exact new head; CI approval should wait for that revision.

[bzantium]

All five threads are answered on the new head, b67b8d9, and the description is updated to match.

Four are implemented as asked. The local backend is the exception: MCore has no non-fused path for linear_q_proj when there is no query LoRA, so I took your second option and gated it with an explicit Transformer Engine requirement rather than reimplementing the resolver in the bridge. Details in that thread.

[yaoyu-33]

/ok to test ee9f7d9

[kamran-nvidia]

/ok to test b6a6612

[bzantium]

The three functional failures turned out to be the same problem in different places, and ee9f7d9 was right to flag all of them. Fixed in abd5d21.

Each of those toys is built from a HF class whose current save layout no longer matches the released checkpoint the bridge maps, so the mapped weights were not in the file at all and the runs were comparing parameters that had never been loaded. That is what the new check exists to catch, so I fixed the toys rather than the check.

exaone was the clearest of them. ExaoneMoeForCausalLM declares _keys_to_ignore_on_load_unexpected = [r"mtp.*"], so released checkpoints carry MTP weights but the class never builds them. The toy config set num_nextn_predict_layers, which made the provider build MTP layers and the bridge register 15 MTP mappings against weights the toy could not possibly save, so I dropped MTP from the toy.

The other two are naming. nemotronh saves backbone.embedding.weight while the released Nemotron-3-Nano index has backbone.embeddings.weight, and qwen_audio nests the language model one level deeper than Qwen/Qwen2-Audio-7B-Instruct does, language_model.model.model.* instead of language_model.model.*. Both are corrected right after save. The audio tower, projector and head were already fine.

I ran what I could locally on two A100s. qwen_audio passes outright, exaone's roundtrip matches 89/89 weights, and nemotronh clears the validation and then stops on the APEX fused wgrad extension, which my container lacks and CI has. gemma_vl only failed in that one run, with EADDRINUSE on port 52177 and no mapping error anywhere in the log, so I left it alone.

What I did not touch is that this mismatch is not only a test artifact. Anyone who re-saves a Nemotron-H or Qwen2-Audio checkpoint with the pinned transformers ends up with a file these bridges cannot map, and before ee9f7d9 it would have loaded quietly with those weights left at their initialization. Whether the bridges should accept both layouts felt bigger than this PR and not really mine to decide for those families, but I can pick it up if you would rather have it here.

[adityavavreNVDA]

/ok to test bdbe7eb

[bzantium]

Everything is green now, including the four jobs that were failing before. The exaone, nemotronh and qwen_audio fixes went in with the last branch update, and gemma_vl passed on the rerun, so that port collision was a one off.

@yaoyu-33 @kamran-nvidia @adityavavreNVDA this is only waiting on review at this point. The earlier review shows as dismissed, so it needs another pass from whoever has a moment. Thanks for keeping the branch current and running CI while I was chasing down the fixtures.