[svcnemo-autobot]

Background

This PR merges main into dev for the 100-commit slice 850bf4db..0156c213. These are the oldest 100 commits of dev..main; commit 101 (91d774de) is excluded. The dev base is ae76c60e.

What changed

• Merge b6810054, with parents dev ae76c60e and main 0156c213.
• Squashed repair 73c827d0 (fix: post-CI corrections), 14 files, +98/-75:
  • Adapts tensor_parallel/random.py and MoE router.py/token_dispatcher.py to API changes from main. random.py keeps dev's [Dev] refactor(checkpoint): re-architect CheckpointWithoutOutput internals #8390 port of refactor(checkpoint): re-architect CheckpointWithoutOutput internals #8334.
  • Adjusts 8 unit tests and removes duplicate entries from the gb200 unit-test recipe.
  • examples/mimo_bagel/train.py now reads get_train_state().iteration, because refactor(training): make TrainState authoritative #7932 no longer sets args.iteration (review 5609274485).
  • pyproject.toml adds the megatron.core.allocator csrc/*.cpp package-data entry from Fix concurrent VMM allocator extension builds #7494.
• Commit graph: 100 original main commits + 1 merge + 1 repair = 102.
• Python lines: +40287 / -6292 across 410 files

Details

• CODEOWNERS, uv.lock and Dockerfile.ci.dev are identical to dev. pyproject.toml differs from dev only by the Fix concurrent VMM allocator extension builds #7494 entry.
• The advisory flagged 4387 dev lines as possibly removed:
  • 4268 exist at the merge-base, so they are main-side edits.
  • 11 are whitespace-only.
  • 108 lines in 37 files trace to main commits in the slice (refactor(training): initialize runtime services after config validation #8272, refactor(training): make TrainState authoritative #7932, Wire CP into MIMO #7758 and others) or to the intended mimo_bagel fix.
  • Nothing was removed by accident, and nothing was restored.
• Remerge-diff covers 146 files: git show --remerge-diff b6810054.
• The MoE megakernel × moe_virtual_expert_load_balance question stays deferred to the MoE owners, with no code change.

100 source commits (oldest first)

6ae0711 fe3481a f10b91d 3bb482c f9ff958 c4334de ecd9d1b fb65908 7a40475 7a0024e de003e9 1e3c0c6 1c3726f 9460129 c919b5e 86edf3d 0c7d417 07046cc 80781ab 4e23fd6 1e0bb06 6bfec46 f87f092 56a36e5 eb00152 cea0980 51e299e 9bd7ba4 9131593 0101df9 915b1d2 e7c9336 3e95097 a54dd73 9eadd78 c13b6e7 3914483 0adca65 ee6baa2 d7924d3 b597bd1 90dfb4b 1ddc6e4 6e7e6ed 649994e 06b8419 54d99a3 3888b34 438ef00 fb011f2 3e2807c 466cf25 a666bd6 f5186bf adcd90f 502078a 1fc08e5 09462ef c6e38c6 a4b2bee 3ef7c7d 36731ac a4fb6cd 4ab3af1 0900d92 8442db1 e6913b8 a747c1b 31265d3 cca1a74 b1d9750 65ac821 1effc75 4dc7e41 db3a6cd 85c950c 62c063f b607c1b ce852bf d56953f e8a302a 60547a3 66a6d3f e592c04 1bb190d 6c8f14a ebab48f d37b3a8 0e34650 4f86e2c 74c9420 e564873 894b511 d2c430a ebab8e2 474b1e9 7aa615b 13a9d46 fa322ce 0156c21

Tested

• Local H100 run with image megatron-lm@sha256:565e34e2b11f5fd6feeccc75f32253acb86f53d06000a01172a92248c38d8d46 (profile megatron-lm v15) on 73c827d0:
  • The test rows this change affects passed: Unit 32/32 and Functional 8/8. No rows were skipped upstream.
  • The other rows (out of 61 unit and 83 functional) keep their earlier passes from 0b27aa77/98b43a0c. Their code and test inputs did not change since then.
• gpt3_mcore_te_tp1_pp2_dsv4_hybrid_mhc_mtp uses lightweight mode in its dev config, so its pytest golden-value step is skipped by design. Training exited with code 0, and its losses matched between runs.
• PR CI on 73c827d0: CICD Megatron-LM, Installation Test, Build/validate Megatron-Core, Copyright check and Process group usage check passed.
• The DCO gate fails because some original main commits have a Signed-off-by identity that differs from the commit author. Those commits stay unchanged, so this needs a maintainer DCO override.
• cicd-mbridge-testing was skipped because the Run MBridge tests label is not set.

[ko3n1g]

/review model=claude mode=strict

[svcnvidia-nemo-ci]

/ok to test 0b27aa7

[ko3n1g]

/review mode=strict model=claude

[svcnvidia-nemo-ci]

Strict review: nightly sync main → dev (2026-10-20)

Findings: 0 CRITICAL · 0 IMPORTANT · 1 SUGGESTION

Overall risk: low for the parts I checked, which include all of the repair commit's source changes. The rest of this 981-file sync was not checked, so I am leaving a comment rather than an approval.

Finding

• [SUGGESTION Simplification] The repair commit added a precomputed_indices parameter to InferenceTopKRouter._compiled_topk_routing in router.py. Nothing ever passes it a value. The only caller passes precomputed_indices only when it is set, and in exactly that case it uses the uncompiled topk_routing_with_score_function instead. The base code already had this caller unchanged, so the parameter does not fix a mismatch. Details are in the inline comment.

What I checked and found no issue with

• MHCCheckpointManager now subclasses CheckpointWithoutOutputManager (random.py). This matches the internals from main #8334. The RNG helpers it relies on (_get_all_rng_states, _set_all_rng_states, _fork_rng) exist and are also used by RecomputeSegment in model_chunk_schedule_plan.py.
• _HybridEPManager.setup_metadata (token_dispatcher.py). The token budget is handled correctly in all three cases:
  • The rank-capacity branch sets the budget.
  • The branch with no capacity setting resets num_permuted_tokens to None.
  • drop-and-pad still overrides it with the capacity value, and combine keeps the static budgets.
• Startup ownership after main #8272. initialize_megatron now calls initialize_runtime_services, which resolves the tokenizer vocab size. That service setup is not safe to run twice. Every changed entry point (BERT, MIMO, inference examples, detxoify, and the dev-only multimodal_dev) correctly dropped its own explicit call. pretrain(..., *, build_tokenizer=...) exists (training.py), so pretrain_mimo's build_tokenizer=False is valid.
• MIMO changes.
  • DistributedInitConfig has align_grad_reduce (default True) and high_priority_stream_groups, which grad_sync.py and pretrain_mimo.py now read.
  • get_data_lane_rank reads pg_collection.dp_gtp_remat, which topology.py populates for every grid.
• Multimodal dataloader resume. The new (args, iteration) signature is used consistently by dataloader_provider and its unit test.
• CI.
  • The new reusable _unit-tests.yml workflow runs the full bucket on GPU unless selective mode is enforce.
  • The coverage gate requires an empty-selection plan for every bucket.
  • Removing A100 from the GitLab stages is consistent across configure, run, cleanup, notify and publish.
  • The gb200 unit-test recipe now has a single test_gated_delta_net.py bucket.
• Router quantile balancing. The router implements both the micro_batch and global_batch estimation scopes, so having the virtual-expert tests select micro_batch only changes the tests.

What this review did not cover

I prioritised source code over config and generated files, following the strict-mode guidance for PRs with more than 50 files. I inspected the repair-commit sources, the inlined CI/docs/examples hunks, and the cross-component paths they touch (training startup, the MIMO topology and helpers, and the MoE router and dispatcher).

I did not individually inspect:

• the large omitted core diffs, such as MFSDP v2, disaggregated inference, GTP sequence sharding and dist-checkpointing;
• most of the about 500 golden-value and recipe files.

Those rely on the author's reported local runs and pending CI. reviewed_files lists only what I inspected.

Review generated by claude. This was a findings-only review. Request a full review with /review reset=true.

[svcnvidia-nemo-ci]

[SUGGESTION Simplification] Strict mode checks for unused arguments: the new precomputed_indices=None parameter of InferenceTopKRouter._compiled_topk_routing is never given a value.

Problem: The only caller is _forward (L1476-L1500). It adds precomputed_indices to routing_kwargs only when self.qb_beta is not None. In exactly that case it picks the uncompiled topk_routing_with_score_function (routing = topk_routing_with_score_function if precomputed_indices is not None or is_batch_invariant_mode_enabled() else self._compiled_topk_routing). So when _compiled_topk_routing runs, routing_kwargs == {} and precomputed_indices is always None.

The base caller is identical (base L1454-L1478) and never passed the argument to the compiled function. The commit message's reason ("matching the base router call site") therefore does not correspond to an actual mismatch. The dead parameter also suggests the quantile-balancing path is compiled when it is not.

Fix (smallest change): remove the parameter and the forwarded keyword so the compiled helper matches its real contract.

         router_replay,
         dense_output,
-        precomputed_indices=None,
     ):
         return topk_routing_with_score_function(
             ...
             dense_output=dense_output,
-            precomputed_indices=precomputed_indices,
         )

If you actually want quantile-balanced (qb_beta) inference routing to run through torch.compile, make that change on purpose in _forward's routing choice instead of leaving a parameter that is never used. That alternative changes behaviour and compilation (an optional tensor argument can trigger recompiles), so it needs its own validation.

Behaviour to keep:

1. With qb_beta set, _forward keeps calling the uncompiled topk_routing_with_score_function(..., precomputed_indices=...).
2. Without qb_beta and outside batch-invariant mode, _forward keeps calling _compiled_topk_routing with the same 11 arguments it uses today.
3. In batch-invariant mode, the uncompiled path stays in use.
4. Hash-layer routing is unaffected.

How to verify:

• Static check: after the change, confirm that nothing passes precomputed_indices to _compiled_topk_routing. I found no tests that reference _compiled_topk_routing.
• Tests to run (I could not run them): tests/unit_tests/transformer/moe/test_routers.py and tests/unit_tests/inference/test_moe_dispatching_and_routing.py. They should still pass, in particular the inference-router and batch-invariant training/inference parity cases (assert_close(..., atol=0, rtol=0)).

[svcnemo-autobot]

Reply to review 5602919574 / inline finding 4339310383 (megatron/core/transformer/moe/router.py:1440, precomputed_indices)

I'm keeping the parameter.

• tests/unit_tests/inference/contexts/test_routing_metadata.py::TestRoutingMetadata::test_async_banks_use_one_graph_and_exclude_mtp comes from main Support routing replay with asynchronous scheduling #7310 and is unchanged by this sync.
• That test calls _compiled_topk_routing(..., precomputed_indices=None) directly, for both compiled and CUDA-graph execution.
• Main 13a9d46d declares the parameter.
• On merge 9d779c05, which lacks the parameter, this test raised TypeError. With it restored, the inference/** unit row passes on 0b27aa77.

The repair commit message says "matching the base router call site", which is imprecise. The parameter is needed by main's test contract, not by the production caller. I'm not amending the commit only for that wording.

Status for maintainers (the bot has taken no CI action):

• DCO gate (run 39572202780) failed at 08:49Z because it ran before the DCO sign-off was manually approved. The DCO check succeeded at 08:50Z. Re-running the failed job should pick up that result.
• MBridge: cicd-mbridge-testing was skipped because the Run MBridge tests label is not set. The main-to-dev sync process requires it.

[FDecaYed]

Approving the sync at 0b27aa77. I re-reviewed the conflict resolutions and the repair commit, and audited dropped changes in both directions. I found no blocking merge errors. Two small merge-introduced issues should ideally be fixed before merge, or in a quick follow-up PR:

1. examples/mimo_bagel still reads args.curr_iteration / args.iteration (L351, L583). Main #7932 stopped setting these on args, and this dev-only example wasn't updated. With PP>1 the data seed is now the same every step, so the numpy/random draws repeat each iteration. Fix: use get_train_state().iteration.

2. Main #7494's package-data entry is missing from pyproject.toml. Dev's pyproject.toml was kept, so [tool.setuptools.package-data] "megatron.core.allocator" = ["csrc/*.cpp"] didn't come over. Wheel and sdist installs won't ship vmm_symm_allocator.cpp, so the VMM allocator build fails there. Source-tree and editable CI runs are unaffected.

One combination for the MoE owners to confirm: dev's MOK megakernel (moe_megakernel_backend="mok") and main's moe_virtual_expert_load_balance meet for the first time in this sync, and neither validation rejects the other. The megakernel path returns early (L881), so the virtual-expert wrap/finalize may not pair up. If the combination isn't supported, a mutual-exclusion check would be the safe option.

Agree with keeping precomputed_indices: main #7310's test calls it directly. MBridge tests didn't run because the label is missing.

[svcnemo-autobot]

Reply to review 5609274485 (@FDecaYed). Thanks for the approval.

• F1 is fixed in 98b43a0c. examples/mimo_bagel/train.py now uses get_train_state().iteration at L351 and L583, with the import updated. Main refactor(training): make TrainState authoritative #7932 (0900d928) no longer sets args.iteration or args.curr_iteration.
• F2 is fixed in 98b43a0c. I added [tool.setuptools.package-data] "megatron.core.allocator" = ["csrc/*.cpp"] from main Fix concurrent VMM allocator extension builds #7494 (6c8f14a3).
• F3 (MoE megakernel with moe_virtual_expert_load_balance) is unchanged in this sync. I'm deferring it to the MoE owners.
• Thanks for agreeing to keep precomputed_indices.

Both fixes went into the single squashed repair commit. The graph is still 98 imports + 1 merge + 1 repair.

To maintainers: this PR still needs the Run MBridge tests label, unless MBridge is waived for this sync.

[svcnvidia-nemo-ci]

/ok to test 98b43a0

[svcnvidia-nemo-ci]

/ok to test 73c827d