Repository navigation
feat(optimizer): AUT-2316 reapply DeepSeek-V3 NVFP4 training support - #7328
Merged
balasaajay merged 6 commits intoSep 24, 2026
Merged
Conversation
svcnvidia-nemo-ci
marked this pull request as draft
September 15, 2026 00:20
Contributor
|
This PR has been automatically converted to draft because all PRs must start as drafts. When you are ready for review, click Ready for Review to begin the review process. This will:
See the contribution guide for more details. |
svcnvidia-nemo-ci
added
the
Run full unit tests
Contributor
|
/ok to test e5c89ba |
Contributor
|
/ok to test 7b9ac78 |
Contributor
|
/ok to test 2fc3a99 |
Contributor
|
/ok to test ac2564e |
balasaajay
approved these changes
Sep 15, 2026
balasaajay
marked this pull request as ready for review
September 15, 2026 19:32
Phlip79
approved these changes
Sep 21, 2026
balasaajay
removed
the
Run full unit tests
Reapply "Add DeepSeek-V3 NVFP4 training support (NVIDIA#6841)". This reverts the revert commit 014508b ("Revert 'Add DeepSeek-V3 NVFP4 training support (NVIDIA#6841)'") from PR NVIDIA#7326, restoring the original change. Signed-off-by: svcnemo-autobot <svcnemo-autobot@nvidia.com>
…load test The re-applied DeepSeek-V3 NVFP4 change (NVIDIA#6841) adds test_distributed_optimizer_reload_main_params_from_fused_mla_canonical_state_dict, whose precondition asserted the checkpoint canonicalizes the fused input LayerNorm to 'decoder.layers.0.input_layernorm.'. The layer's sharded_state_dict_keys_map rewrites 'self_attention.linear_qkv_down_proj.layer_norm_' to 'input_layernorm.' *within* the layer, so the canonical checkpoint key is 'decoder.layers.0.self_attention.input_layernorm.' (the layer's input_layernorm is a sublayer of the transformer layer). Correct the asserted prefix so the test exercises the reload path instead of failing its own precondition. Signed-off-by: svcnemo-autobot <svcnemo-autobot@nvidia.com>
Signed-off-by: svcnemo-autobot <svcnemo-autobot@nvidia.com>
Signed-off-by: svcnemo-autobot <svcnemo-autobot@nvidia.com>
…eload test The captured canonical sharded_state_dict for the fused-MLA layer keeps the input LayerNorm fused on linear_qkv_down_proj.layer_norm_weight rather than remapping it to a separate self_attention.input_layernorm.* key. Fix the over-strict assertion (and stale docstring) to match the real canonical keys, and drop the temporary diagnostic prints. Signed-off-by: svcnemo-autobot <svcnemo-autobot@nvidia.com>
…re-apply The revert-of-revert of NVIDIA#6841 predates and stripped several additions main made to gpt_model.py after NVIDIA#6841 merged: the freeze_base_model_for_mtp / mtp_hsm / mtp_hybrid_override_pattern input validation, the output-logits observe_tensor call, and the conditional decoder name. Their removal broke main's GPTModel tests (test_hsm_requires_two_mtp_layers, test_rejects_hybrid_mtp_override). Restore main's content for gpt_model.py so the NVFP4 re-apply no longer regresses those validations; gpt_model.py is now unchanged from main. Signed-off-by: svcnemo-autobot <svcnemo-autobot@nvidia.com>
nemo-autobot-origin
Bot
force-pushed
the
revert/deepseek-nvfp4-6841
branch
from
September 22, 2026 15:57
899707a to
efb182d
Compare
balasaajay
enabled auto-merge
September 22, 2026 16:15
svcnvidia-nemo-ci
added
the
Run full unit tests
Contributor
|
/ok to test efb182d |
zhongbozhu
approved these changes
Sep 23, 2026
jaredcasper
approved these changes
Sep 23, 2026
svcnvidia-nemo-ci
added
Approved
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/35937923675 |
This branch was previously deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Background
PR #7326 ("ci: Revert recently merged changes as they cause CI unit test failures") reverted three recently-merged changes. This PR re-applies one of them: "Add DeepSeek-V3 NVFP4 training support (#6841)", reverted by commit
014508b.What changed
Reverts the revert commit
014508ba8d2158d33b39a9462b630efc60e25cfd, restoring DeepSeek-V3 NVFP4 training support across the distributed optimizer, GPT model, multi-token prediction, and transformer block/layer, plus the related unit tests. Rebased onto currentmain.Details
Rebased onto
main@18a991265. One conflict resolved at the root: main addedhash_moe_layer_thresholdtoMultiTokenPredictionBlock's pattern-build path where this change renames MTP layers tolayers.{layer_number - 1}(NVFP4 checkpoint-key namespacing). The resolution keeps main's new kwarg and this change's layer naming.gpt_model.pyis restored identical to currentmainso the re-apply no longer regresses main's MTP input validation. The fused-MLA reload test asserts the canonical fused key (...self_attention.linear_qkv_down_proj.layer_norm_weight).Tested
Rebase conflict resolution verified: post-rebase tree-vs-main diff is identical to the pre-rebase change for all touched files (same 6 files, 167 insertions / 26 deletions) modulo the intended
hash_moe_layer_thresholdcontext line and line-number drift. Local lint/test toolchain unavailable in the rootless authoring workstation; GitHub CI on the pushed head is the authoritative gate.