Repository navigation
Scope layerwise optimizer collectives to be per module for mulitmodal mimo training - #6192
Merged
yashaswikarnati merged 3 commits intoAug 24, 2026
Conversation
yashaswikarnati
marked this pull request as ready for review
August 4, 2026 00:23
yashaswikarnati
force-pushed
the
agent/mimo-layerwise-optimizer-support
branch
2 times, most recently
from
August 23, 2026 00:23
8ad7ddd to
588b3ac
Compare
yashaswikarnati
marked this pull request as draft
August 23, 2026 00:24
Signed-off-by: ykarnati <ykarnati@nvidia.com>
yashaswikarnati
force-pushed
the
agent/mimo-layerwise-optimizer-support
branch
from
August 23, 2026 23:52
a82946c to
680fc4e
Compare
yashaswikarnati
marked this pull request as ready for review
August 23, 2026 23:53
Contributor
Author
|
/ok to test 680fc4e |
Signed-off-by: ykarnati <ykarnati@nvidia.com>
yashaswikarnati
force-pushed
the
agent/mimo-layerwise-optimizer-support
branch
from
August 24, 2026 00:02
680fc4e to
145297c
Compare
Contributor
Author
|
/ok to test 145297c |
yashaswikarnati
changed the title
Signed-off-by: ykarnati <ykarnati@nvidia.com>
deepakn94
approved these changes
Aug 24, 2026
Contributor
Author
|
/ok to test e047dfb |
Phlip79
approved these changes
Aug 24, 2026
svcnvidia-nemo-ci
added
Approved
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/32750099346 |
ksivaman
pushed a commit
to ksivaman/Megatron-LM
that referenced
this pull request
Sep 1, 2026
… mimo training (NVIDIA#6192) Signed-off-by: ykarnati <ykarnati@nvidia.com>
xuwchen
pushed a commit
to xuwchen/Megatron-LM
that referenced
this pull request
Oct 14, 2026
… mimo training (NVIDIA#6192) Signed-off-by: ykarnati <ykarnati@nvidia.com> (cherry picked from commit 0560fbb) Signed-off-by: xuwenc <xuwenc@nvidia.com>
xuwchen
pushed a commit
to xuwchen/Megatron-LM
that referenced
this pull request
Oct 14, 2026
… mimo training (NVIDIA#6192) Signed-off-by: ykarnati <ykarnati@nvidia.com> (cherry picked from commit 0560fbb) Signed-off-by: xuwenc <xuwenc@nvidia.com>
nemo-autobot-origin Bot
pushed a commit
to svcnemo-autobot/Megatron-LM
that referenced
this pull request
Oct 19, 2026
Single repair commit for the 2026-10-16 main->dev sync of the 100-commit main slice ending at 302ed23 onto dev dee456a. Dependencies - pyproject.toml: add gigatoken to the training and mlm extras, which code imported from main requires. - uv.lock: regenerated with uv 0.7.2 (Dockerfile pin); adds gigatoken 0.10.0, awkward 2.14.0 and awkward-cpp 57; no removals. Tests adapted from main's API to dev's API - transformer/test_multi_token_prediction.py: dev's MTPLossLoggingHelper exposes reduced values via reduce_loss_in_tracker()/tracker["values"]. - data/test_get_batch.py: dev's pretrain_hybrid.get_batch reads args.linear_cp_layout and args.context_parallel_size. - inference/contexts/test_dynamic_prefix_caching.py: on dev the engine records the Mamba match count before admission; set it like the engine. - models/mimo/test_mimo_1f1b_schedule.py: pass use_layer_wise_param_layout so DistributedDataParallelConfig agrees with the layer-wise optimizer. - transformer/experimental_attention_variant/test_attention_variant_dsa.py: dynamic context parallel on dev requires max_seqlen_per_dp_cp_rank. - transformer/test_multi_token_prediction.py (main NVIDIA#6119 speculative-decoding capture test): dev stores the MTP input on model._decoder_hidden_states_cache, not on inference_context.mtp_decoder_hidden_states. Golden values - moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon dev H100 golden refreshed from the exact merge-candidate run: main 0560fbb (NVIDIA#6192) makes the LayerWise Muon child grad statistics span the intra distributed-optimizer group; reverting only that line reproduces dev's previous golden values exactly. - tests: set max_seqlen_per_dp_cp_rank in dev dynamic-CP/scheduler test configs (transformer_layer, transformer_config_dynamic_cp, dsv4_hybrid_attention_cp) to satisfy main NVIDIA#5902 validation; dev arguments.py already derives it when a sequence-packing scheduler is set. - tests: transformer/test_thd_cuda_graph.py balanced DSA dynamic-pack case with max_seqlen_per_dp_cp_rank=None now expects main NVIDIA#5902's earlier ModelParallelConfig validation message. - tests: extend main NVIDIA#6284 mimo overlap-lifecycle optimizer mock with the dev _optimizer_state_offloader/_layer_wise_non_distopt_child attributes (NVIDIA#6244). - tests: extend main NVIDIA#6119 HybridModel MTP forward stub with dev moe_n_hash_layers and pg_collection.dp_cp used by the dev forward. Megatron-FSDP v2 symmetric memory (requester-approved alignment with main) - experimental/dbuffer.py, experimental/parameter_group.py: take main's version. Symmetric-memory MFSDP raises "requires PyTorch 2.12 or later" on PyTorch builds without torch.distributed._symmetric_memory .is_symm_mem_tensor, as on main. This drops dev's allocation-flag fallback (ef8d090, 4a534df); with it, the symmetric path's losses diverged across ranks on dev's PyTorch 2.12 pre-release image and test_fully_shard_symmetric_memory_matches_default_and_profiles_nccl[3] failed on the merge candidate. - mfsdp_v2 tests: skip the symmetric-memory cases on such builds (test_symmetric_memory.py, the symmetric_memory params of test_overlap.py and test_cuda_graph.py, and test_mcore_adapter.py nccl_ub). They run again once the image provides is_symm_mem_tensor. Hybrid context parallelism (review 5570704991) - models/hybrid/hybrid_block.py: HybridStack rejects linear_cp_layout='contiguous' with context parallelism when the segment contains GDN ('G') or KDA ('K') layers. main NVIDIA#6776 feeds contiguous shards to linear layers, but dev's GDN/KDA chunkwise and headwise paths still read their input layout from cp_partition_mode (zigzag by default) and would scramble token order. Contiguous Mamba/GDP and the zigzag GDN/KDA paths are unchanged. - ssm/test_hybrid_block.py: cover the GDN/KDA rejection and the zigzag GDN and contiguous Mamba builds under CP=2. Signed-off-by: svcnemo-autobot <svcnemo-autobot@nvidia.com>
nemo-autobot-origin Bot
pushed a commit
to svcnemo-autobot/Megatron-LM
that referenced
this pull request
Oct 19, 2026
Single repair commit for the 2026-10-16 main->dev sync of the 100-commit main slice ending at 302ed23 onto dev dee456a. Dependencies - pyproject.toml: add gigatoken to the training and mlm extras, which code imported from main requires. - uv.lock: regenerated with uv 0.7.2 (Dockerfile pin); adds gigatoken 0.10.0, awkward 2.14.0 and awkward-cpp 57; no removals. Tests adapted from main's API to dev's API - transformer/test_multi_token_prediction.py: dev's MTPLossLoggingHelper exposes reduced values via reduce_loss_in_tracker()/tracker["values"]. - data/test_get_batch.py: dev's pretrain_hybrid.get_batch reads args.linear_cp_layout and args.context_parallel_size. - inference/contexts/test_dynamic_prefix_caching.py: on dev the engine records the Mamba match count before admission; set it like the engine. - models/mimo/test_mimo_1f1b_schedule.py: pass use_layer_wise_param_layout so DistributedDataParallelConfig agrees with the layer-wise optimizer. - transformer/experimental_attention_variant/test_attention_variant_dsa.py: dynamic context parallel on dev requires max_seqlen_per_dp_cp_rank. - transformer/test_multi_token_prediction.py (main NVIDIA#6119 speculative-decoding capture test): dev stores the MTP input on model._decoder_hidden_states_cache, not on inference_context.mtp_decoder_hidden_states. Golden values - moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon dev H100 golden refreshed from the exact merge-candidate run: main 0560fbb (NVIDIA#6192) makes the LayerWise Muon child grad statistics span the intra distributed-optimizer group; reverting only that line reproduces dev's previous golden values exactly. - tests: set max_seqlen_per_dp_cp_rank in dev dynamic-CP/scheduler test configs (transformer_layer, transformer_config_dynamic_cp, dsv4_hybrid_attention_cp) to satisfy main NVIDIA#5902 validation; dev arguments.py already derives it when a sequence-packing scheduler is set. - tests: transformer/test_thd_cuda_graph.py balanced DSA dynamic-pack case with max_seqlen_per_dp_cp_rank=None now expects main NVIDIA#5902's earlier ModelParallelConfig validation message. - tests: extend main NVIDIA#6284 mimo overlap-lifecycle optimizer mock with the dev _optimizer_state_offloader/_layer_wise_non_distopt_child attributes (NVIDIA#6244). - tests: extend main NVIDIA#6119 HybridModel MTP forward stub with dev moe_n_hash_layers and pg_collection.dp_cp used by the dev forward. Megatron-FSDP v2 symmetric memory (requester-approved alignment with main) - experimental/dbuffer.py, experimental/parameter_group.py: take main's version. Symmetric-memory MFSDP raises "requires PyTorch 2.12 or later" on PyTorch builds without torch.distributed._symmetric_memory .is_symm_mem_tensor, as on main. This drops dev's allocation-flag fallback (ef8d090, 4a534df); with it, the symmetric path's losses diverged across ranks on dev's PyTorch 2.12 pre-release image and test_fully_shard_symmetric_memory_matches_default_and_profiles_nccl[3] failed on the merge candidate. - mfsdp_v2 tests: skip the symmetric-memory cases on such builds (test_symmetric_memory.py, the symmetric_memory params of test_overlap.py and test_cuda_graph.py, and test_mcore_adapter.py nccl_ub). They run again once the image provides is_symm_mem_tensor. Hybrid context parallelism (review 5570704991) - models/hybrid/hybrid_block.py: HybridStack rejects linear_cp_layout='contiguous' with context parallelism when the segment contains GDN ('G') or KDA ('K') layers. main NVIDIA#6776 feeds contiguous shards to linear layers, but dev's GDN/KDA chunkwise and headwise paths still read their input layout from cp_partition_mode (zigzag by default) and would scramble token order. Contiguous Mamba/GDP and the zigzag GDN/KDA paths are unchanged. - ssm/test_hybrid_block.py: cover the GDN/KDA rejection and the zigzag GDN and contiguous Mamba builds under CP=2. - tests: transformer/test_multi_token_prediction.py TestMultiTokenPredictionHybrid (main NVIDIA#6119) stub config sets cross_entropy_loss_fusion=False and cross_entropy_fusion_impl="native" (the TransformerConfig defaults); dev NVIDIA#8105 HybridModel.forward reads them when labels are given (review 5573599695). Signed-off-by: svcnemo-autobot <svcnemo-autobot@nvidia.com>
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.
Summary
intra_dist_optgroup;dp_cporexpt_dpgroup;intra_dist_optgroup;