Fix for sequence-level aux MoE loss being dependent on batch size by OlegSudakov · Pull Request #5798 · NVIDIA/Megatron-LM

OlegSudakov · 2026-07-14T07:57:01Z

I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Fix MBS-dependent seq_aux_loss gradient under --calculate-per-token-loss

What

We have observed a gradual increase in logged sequence-level aux MoE loss when increasing the MBS from 1 to 2. Investigation findings - under --calculate-per-token-loss, the seq_load_balancing_loss gradient is
scaled by a spurious 1/MBS factor, so increasing the micro-batch size silently
deprioritizes load balancing and shifts the training trajectory (the logged loss
barely moves since it is recorded before the scaling).

Root cause

_apply_seq_aux_loss reshapes routing_map to [seq_length, bsz*num_experts],
so local_num_tokens = routing_map.shape[0] = seq_length — independent of MBS.
It passes this as valid_token_count, which attach_and_log_load_balancing_loss
uses to pre-multiply the aux gradient and cancel the global 1/total_tokens
division in finalize_model_grads (total_tokens = GBS*seq_len, constant).
Because it supplies seq_len instead of the micro-batch total seq_len*MBS, the
MBS factor never cancels. Plain aux_loss is unaffected (its routing_map keeps
shape [seq_len*MBS, num_experts]).

Fix

Pass valid_token_count=local_num_tokens * bsz in _apply_seq_aux_loss. Correct
with and without padding (with padding, local_num_tokens is the mean valid
tokens/seq, so * bsz recovers the total valid tokens). Only the
calculate_per_token_loss=True path changes.

Test

Adds test_seq_aux_loss_mbs_invariant_per_token_loss: asserts equal accumulated
router gradients at MBS=1 vs MBS=N (with and without padding). Fails pre-fix.

Issue tracking

For PRs from open-source community contributors:

New features: a linked issue is required. Please open a feature request and reference it here before submitting the PR.
Small updates (bug fixes, minor improvements): a linked issue is recommended and will accelerate the PR review process.

Linked issue:

Contribution process

Pre-checks

I have added relevant unit tests
I have added relevant functional tests
I have added proper typing to my code Typing guidelines
I have added relevant documentation
I have run the autoformatter.sh on my PR

Code review

Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

When your PR is ready, click Ready for Review.
An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
- Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

copy-pr-bot · 2026-07-14T07:57:04Z

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.

github-actions · 2026-07-14T07:57:11Z

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:

Add the oncall reviewer (optional reviewer)
Add required review teams based on your changes

See the contribution guide for more details.

Signed-off-by: Oleg Sudakov <osudakov@nvidia.com>

Signed-off-by: Oleg Sudakov <oleg.sudakov@outlook.com>

YangFei1990 · 2026-07-16T18:51:06Z

/ok to test 9c83c5f

YangFei1990 · 2026-07-17T03:00:26Z

/ok to test 6c1b91a

svcnvidia-nemo-ci · 2026-07-17T04:53:35Z

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/29555939384

svcnvidia-nemo-ci · 2026-07-17T06:24:19Z

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/29559932466

…IDIA#5798) Signed-off-by: Oleg Sudakov <osudakov@nvidia.com> Signed-off-by: Oleg Sudakov <oleg.sudakov@outlook.com> Co-authored-by: Fei Wu <33940270+YangFei1990@users.noreply.github.com> Signed-off-by: mchochowski <mchochowski@nvidia.com>

Post-merge corrections for the nightly main -> dev sync, in the single rolling fix commit on top of the immutable merge commit. Each addresses a CI failure caused by the 3-way merge (kept main's usage but dev's narrower import line; paired main's caller with a dev-renamed symbol; duplicated a definition; or dropped main's numeric fix while keeping main's test). 1. Dropped imports (merge kept main's usages, dev's narrower import line won): - inference/engines/dynamic_engine.py: restore CUDAGraphBatchDimensionBuilder, InferenceBatchDimensions. - inference/contexts/dynamic_context.py: restore MAX_INTERMEDIATE_OFFSETS_PER_REQUEST. - optimizer/optimizer.py: restore "from itertools import chain". 2. API mismatch from an override-from-main file (main's caller used a symbol dev renamed): - training/training.py: get_hybrid_data_context_parallel_groups -> get_dynamic_data_context_parallel_groups (import + call site). 3. Duplicate definitions (autodoc2.dup_item / build warnings): - packed_seq_params.py: drop main's duplicate "pad_between_seqs: bool = None", keep dev's Optional[bool] field. - transformer/moe/fused_a2a.py: drop the duplicated HYBRIDEP_TOKEN_ALIGNMENT = 64. 4. Golden-value nan sentinels: drop the step-1 warmup "nan" iteration-time key and set start_step=2 on three golden files (bert tp1_pp2, bert tp1_pp4_vp2, moe deepseek_proxy_fsdp_ep2_fsdp2) for tools/check_golden_values.py. 5. inference/engines/dynamic_engine.py: restore the __init__ assignment self.cuda_graph_all_prefills = inference_config.cuda_graph_all_prefills. The merge kept main's usage in _cg_admission_gating_active but dropped the assignment, so DynamicInferenceEngine raised AttributeError. 6. tests/.../inference/contexts/test_dynamic_prefix_caching.py: restore self.cuda_graph_all_prefills = False in the _StubEngine that bypasses super().__init__, matching fix 5. 7. transformer/moe/router.py: restore valid_token_count=local_num_tokens * bsz in the seq_load_balancing_loss branch. Fix e3fe550 (#5798) made the sequence-level aux loss batch-size-independent on main; the merge reverted the source line to dev's pre-fix value while keeping main's test, producing grad_mbs1 == 4x grad_mbsN. 8. optimizer/layer_wise_optimizer.py: reconcile the LayerWise Muon layout stack. The file was taken wholesale from main (skill override list, reason "constructor signature"), but the signatures are now identical and the rest of the retained tree is dev's: optimizer_config (use_layer_wise_param_layout), param_and_grad_buffer (the layout producer), and the base optimizer.py (save/load_state_dict_from_file + _layer_wise_non_distopt_child). Main's layout consumer against dev's producer raised "Optimizer param (shape=(4,16)) not found in any param layout". Dev is the more-evolved side (PRs #5388/#5470/#5684 land after the merge-base; main has #4967 GTP + #6048). Restored dev's layer_wise_optimizer.py and grafted main's GTP replica-fold support (_build_gtp_replica_fold, _fold_replica_id, the sharded_state_dict fold call) so main's incoming GTP (#4967) checkpoint correctness is preserved; _fold_replica_id reduces to dev's DP-zero reset when GTP is off. Empirically verified test_optimizer.py, transformer/**, transformer/moe/**, mfsdp_v2/**, and inference/** are green on recent dev CI (PRs #6020/#5967/#6040), confirming these buckets are sync-caused, not pre-existing. Signed-off-by: svcnvidia-nemo-ci <svcnvidia-nemo-ci@nvidia.com>

…IDIA#5798) Signed-off-by: Oleg Sudakov <osudakov@nvidia.com> Signed-off-by: Oleg Sudakov <oleg.sudakov@outlook.com> Co-authored-by: Fei Wu <33940270+YangFei1990@users.noreply.github.com>

…IDIA#5798) Signed-off-by: Oleg Sudakov <osudakov@nvidia.com> Signed-off-by: Oleg Sudakov <oleg.sudakov@outlook.com> Co-authored-by: Fei Wu <33940270+YangFei1990@users.noreply.github.com> Signed-off-by: Dmytro Pykhtar <dpykhtar@nvidia.com>

OlegSudakov requested review from a team as code owners July 14, 2026 07:57

svcnvidia-nemo-ci marked this pull request as draft July 14, 2026 07:57

OlegSudakov force-pushed the osudakov/seq_aux_loss_mbs_fix branch 5 times, most recently from c87b718 to f40aef0 Compare July 14, 2026 08:24

Fix for sequence-level aux moe loss being dependent on batch size

27f4986

Signed-off-by: Oleg Sudakov <osudakov@nvidia.com>

OlegSudakov force-pushed the osudakov/seq_aux_loss_mbs_fix branch from f40aef0 to 27f4986 Compare July 14, 2026 08:33

Increase numerical precision for test

95eac2c

Signed-off-by: Oleg Sudakov <oleg.sudakov@outlook.com>

YangFei1990 marked this pull request as ready for review July 16, 2026 18:50

Merge branch 'main' into osudakov/seq_aux_loss_mbs_fix

9c83c5f

svcnvidia-nemo-ci added the complexity: low label Jul 16, 2026

copy-pr-bot Bot temporarily deployed to public July 16, 2026 18:51 Inactive

copy-pr-bot Bot temporarily deployed to test July 16, 2026 18:52 Inactive

copy-pr-bot Bot temporarily deployed to public July 16, 2026 18:54 Inactive

copy-pr-bot Bot temporarily deployed to public July 16, 2026 18:55 Inactive

copy-pr-bot Bot temporarily deployed to public July 16, 2026 19:05 Inactive

YangFei1990 approved these changes Jul 16, 2026

View reviewed changes

Phlip79 approved these changes Jul 16, 2026

View reviewed changes

yaoyu-33 approved these changes Jul 16, 2026

View reviewed changes

yashaswikarnati approved these changes Jul 16, 2026

View reviewed changes

yaox12 approved these changes Jul 17, 2026

View reviewed changes

svcnvidia-nemo-ci added the Final Review PR is in the "final review" stage label Jul 17, 2026

YangFei1990 enabled auto-merge July 17, 2026 01:29

jiemingz approved these changes Jul 17, 2026

View reviewed changes

svcnvidia-nemo-ci added Approved All necessary approvals have been made and removed Final Review PR is in the "final review" stage labels Jul 17, 2026

Merge branch 'main' into osudakov/seq_aux_loss_mbs_fix

6c1b91a

copy-pr-bot Bot temporarily deployed to public July 17, 2026 03:00 Inactive

copy-pr-bot Bot temporarily deployed to test July 17, 2026 03:01 Inactive

copy-pr-bot Bot temporarily deployed to public July 17, 2026 03:04 Inactive

copy-pr-bot Bot temporarily deployed to public July 17, 2026 03:05 Inactive

copy-pr-bot Bot temporarily deployed to public July 17, 2026 03:14 Inactive

Victarry approved these changes Jul 17, 2026

View reviewed changes

YangFei1990 added this pull request to the merge queue Jul 17, 2026

Merged via the queue into NVIDIA:main with commit e3fe550 Jul 17, 2026
90 checks passed

lhb8125 mentioned this pull request Aug 18, 2026

Megatron-LM Monthly Activity Report — July 2026 #6607

Open

Navigation Menu

Uh oh!

Fix for sequence-level aux MoE loss being dependent on batch size - #5798

Fix for sequence-level aux MoE loss being dependent on batch size#5798
YangFei1990 merged 4 commits into
NVIDIA:mainfrom
OlegSudakov:osudakov/seq_aux_loss_mbs_fix

OlegSudakov commented Jul 14, 2026

Uh oh!

copy-pr-bot Bot commented Jul 14, 2026

Uh oh!

github-actions Bot commented Jul 14, 2026

Uh oh!

YangFei1990 commented Jul 16, 2026

Uh oh!

YangFei1990 commented Jul 17, 2026

Uh oh!

svcnvidia-nemo-ci commented Jul 17, 2026

Uh oh!

svcnvidia-nemo-ci commented Jul 17, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

9 participants

Uh oh!

Conversation

OlegSudakov commented Jul 14, 2026

What does this PR do?

What

Root cause

Fix

Test

Issue tracking

Contribution process

Pre-checks

Code review

Step 1: Mark PR as "Ready for Review"

Step 2: Final Review

Step 3: Approved

Merge

Uh oh!

copy-pr-bot Bot commented Jul 14, 2026

Uh oh!

github-actions Bot commented Jul 14, 2026

Uh oh!

YangFei1990 commented Jul 16, 2026

Uh oh!

YangFei1990 commented Jul 17, 2026

Uh oh!

svcnvidia-nemo-ci commented Jul 17, 2026

Uh oh!

svcnvidia-nemo-ci commented Jul 17, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

9 participants