Support the MIMO cross-grid path in training loop by yashaswikarnati · Pull Request #5373 · NVIDIA/Megatron-LM

yashaswikarnati · 2026-06-16T20:10:37Z

What

Three adjustments so the stock training loop supports the MIMO cross-grid path, all no-ops for non-MIMO runs:

get_forward_backward_func selects forward_backward_pipelining_without_interleaving when passed a MultiModuleProcessGroupCollection (the cross-grid bridge); train() calls the selector unconditionally instead of branching on p2p_communicator.
Guard the losses_reduced[0] access: MIMO encoder-grid ranks are their module's last PP stage but compute no language loss, so the list is empty there.
Skip the ModelOpt distillation shape hook unless modelopt_enabled; it reads parallel_state, which is uninitialized on the MIMO path.

Validation

8-GPU 20L Nemotron VLM e2e: trains + checkpoint save/resume (lm loss 12.18 → 11.54 across resume).

Three guarded adjustments to the stock training loop, all keyed on the already-threaded p2p_communicator argument (None for non-MIMO runs) or an empty losses_reduced, so stock single-grid behavior is unchanged: - Select forward_backward_pipelining_without_interleaving when a cross-grid p2p_communicator is present. The MIMO bridge cannot pick the schedule by parallel_state pp size, which is unset when mpu is not initialized; the without-interleaving schedule accepts both the p2p_communicator and the MultiModuleProcessGroupCollection. - Guard the losses_reduced[0] access with `and losses_reduced`. MIMO encoder-grid ranks are their own module's last PP stage but never compute the language loss, so losses_reduced is empty on those ranks. - Skip the ModelOpt distillation tensor-shape hook when p2p_communicator is present. MIMO does not use distillation, and the helper reads parallel_state pp groups that are not initialized without mpu. Validated in the 8-GPU 20L Nemotron VLM e2e (trains + checkpoint save/resume, lm loss 12.18->11.54 across resume). Signed-off-by: ykarnati <ykarnati@nvidia.com>

copy-pr-bot · 2026-06-16T20:10:41Z

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.

yashaswikarnati · 2026-06-16T22:15:04Z


-    if is_last_stage:
-        # Average loss across microbatches.
+    if is_last_stage and losses_reduced:


what exactly is last stage is here and how its computed? if its only LLM why we also need losses_reduced ?

yashaswikarnati · 2026-06-17T00:28:42Z


-        if has_nvidia_modelopt:
-            # [ModelOpt]: Pipeline-parallel Distillation stacks student and teacher tensors
+        if has_nvidia_modelopt and p2p_communicator is None:


why did we add p2p_communicator is None ? when is use_model_opt triggered ?

yashaswikarnati · 2026-06-17T00:29:21Z

    # Wrap forward_backward_func for Full iteration CUDA graph
-    forward_backward_func = get_forward_backward_func()
+    if p2p_communicator is not None:
+        # MIMO cross-grid bridge: select the schedule by the presence of a


verbose doc strings

Address review feedback: the three MIMO cross-grid guard comments were multi-paragraph; trim to terse one-liners consistent with surrounding training.py comment density. No logic change. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: ykarnati <ykarnati@nvidia.com>

Replace the p2p_communicator-based skip with the modelopt_enabled args flag, which expresses the real precondition (the distillation shape-adjust helper reads parallel_state) without assuming the MIMO path. Trim the losses_reduced guard comment to a one-liner. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: ykarnati <ykarnati@nvidia.com>

Remove the explanatory comment on the MIMO schedule-selection branch; the branch logic is unchanged. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: ykarnati <ykarnati@nvidia.com>

Move cross-grid bridge schedule selection into the canonical get_forward_backward_func selector: when a MultiModuleProcessGroupCollection is passed, it returns forward_backward_pipelining_without_interleaving (the only schedule that accepts the cross-grid communicator and pg collection). This replaces the implicit p2p_communicator-based branch in train() with an explicit, single unconditional call, keying off the actual multi-grid marker so colocated and stock single-grid runs are unaffected. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: ykarnati <ykarnati@nvidia.com>

yashaswikarnati · 2026-06-22T16:32:03Z

/ok to test 649065d

svcnvidia-nemo-ci · 2026-06-22T19:16:24Z

🔄 Merge queue validation started!

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

Signed-off-by: ykarnati <ykarnati@nvidia.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Vitaly Kurin <vitalyk@nvidia.com>

Signed-off-by: ykarnati <ykarnati@nvidia.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

Signed-off-by: ykarnati <ykarnati@nvidia.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Jon Barker <jbarker@aws-cmh-slurm-1-vscode-02.cm.cluster>

Signed-off-by: ykarnati <ykarnati@nvidia.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

Signed-off-by: ykarnati <ykarnati@nvidia.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Dmytro Pykhtar <dpykhtar@nvidia.com>

…shard mapping GTP shards dim0 of the fused [gate|up] fc1 weight, but the swiglu checkpoint factory used to run per-shard: each shard was chunked into [gate_i|up_i], so the checkpoint layout depended on the save-time GTP degree and, worse, the storage mapping it implied was interleaved -- the all-gathered weight was NOT the logical TP-local tensor, and every consume site needed a de-interleave permutation (plus the matching wgrad-side inverse when syncing gradients). Those two runtime transforms cost extra kernels on every layer consume and were the prime suspect in the GTP-vs-MFSDP perf asymmetry. Replace the runtime-transform approach with the same storage re-layout that Mamba/GatedDeltaNet in_proj already uses: pin the mapping "GTP shard i == contiguous rows [i*n, (i+1)*n) of the logical TP-local weight" via the checkpoint wiring, and delete both runtime transforms entirely. - New megatron/core/tensor_parallel/gtp_ckpt.py hosts the gather/slice pair (moved verbatim from ssm/utils.py, which now re-exports them): on save, all-gather the shards back to the unpadded TP-local tensor, THEN run the standard apply_swiglu_sharded_factory on it -- identical on-disk format to a non-GTP run; on load, wrap merge_fn to cat -> re-pad -> slice this rank's contiguous rows. - MLP.sharded_state_dict wires the pair around the swiglu factory when fc1 is a GTP param (covers dense MLP and SharedExpertMLP). Grouped routed-expert fc1 under EGTP keeps its NotImplementedError guard until the grouped wiring is ported. - The gathered weight is now logically ordered by construction, so the consume-site de-interleave and the wgrad re-interleave (and their transformability checks) are deleted: zero per-iteration cost, and no transform pair to keep in sync. Disk format is unchanged (gate/up halves of the full TP-local weight), so checkpoints interop with 3D-parallel runs and existing checkpoints load correctly under the new mapping. Tests: TestGTPGatedFusedLayout now seeds shards as contiguous slices of a broadcast logical weight and asserts elementwise fwd equality vs F.linear plus per-rank wgrad slices (permutation-sensitive, unlike grad-norm checks); new TestGtpFc1SwigluDcp asserts the swiglu factory shapes/offsets, merge round-trip (pad rows re-zeroed) and DCP writer election under TP2xGTP2 with alignment padding. fix(gtp): give expert grads the GTP/EGTP share of the DP normalization GTP carves its ranks out of the data-parallel axis, so DDP's 1/DP scaling shrinks as GTP grows. Dense params recover the missing factor from the gtp_remat AVG; expert params only ever see compensation sized to the EGTP axis, leaving their gradients GTP/EGTP times too large. Expert layers were therefore stepped that many times too far while dense layers were correct, which no learning-rate change can compensate. The failure is silent -- no error, no NaN, just expert gradients off by a constant -- so a short run cannot see it and only convergence can. The CPU tests therefore pin the arithmetic directly with stand-in groups: the GTP/EGTP factor is applied to expert params (allreduce=False) and not to dense ones, nothing happens when GTP does not exceed EGTP or the groups are absent, calculate_per_token_loss is skipped, and params without a grad or with requires_grad=False are left alone. test(gtp): cover the multimodal vision->embedding->language chain ordering Multimodal models put the vision tower before the language embedding on one prefetch chain, and the embedding is a forward-only chain member: it links on its on-demand forward gather but its backward (GTPEmbeddingWeight) never enters all_gather_and_prefetch_bwd, so it produces no gather for its chain predecessor. The per-consume gather rule from #6242 is what keeps this ordering correct — the vision tail's backward consume finds no gather issued for it and falls back on demand. Pin that with bitwise ground-truth comparisons on every consume, in both directions, so a regression shows up as wrong data here rather than as a silent stale read in a full multimodal run. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> fix(gtp): import itertools.chain in the GTP sharded-param backfill _backfill_gtp_sharded_param_map (upstream #4967) iterates chain.from_iterable(float16_groups) but never imports chain, so the first distributed-Muon save with an unmatched GTP param dies with `NameError: name 'chain' is not defined` instead of backfilling. tests/unit_tests/generalized_tensor_parallel/test_gtp_muon_dcp.py hits it on its first optimizer sharded_state_dict call. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> feat(gtp): make the weight-remat alignment pad configurable GTPRematConfig.pad_for_alignment defaults to 16 and is only ever overridden by the three quantized branches in configure_gtp_remat_from_recipe, so a bf16 run pads without any way to say otherwise. The pad rows are stripped before every GEMM, but they are not invisible: they widen a GTP tensor's global shape in the checkpoint (which is why a GTP checkpoint cannot be read by a non-GTP run), and on load they are re-zeroed rather than round-tripped. Add --gtp-remat-pad-for-alignment so a run can pin or disable it. This is what lets an experiment separate 'the pad rows cause X' from 'something else does' without editing the source. CPU tests pin the precedence: each quantized recipe keeps its historical pad, bf16 leaves the pre-existing value untouched, and an explicit value (including 0 to disable) wins over every recipe. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> fix(gtp): make GTP produce a loadable fully-reshardable optimizer checkpoint dist_ckpt_optim_fully_reshardable is the only distributed-optimizer format whose state survives a change of parallel layout: every other format prefixes its entries with `optimizer.distributed.dp_group_idx_N.`, pinning them to the DP grouping in force at save time. GTP could neither write that format nor be read back from it, so a GTP run could not hand its optimizer state to a differently-sharded run in either direction. Two things were in the way, and fixing either alone leaves the capability unusable. 1. The save aborted. sharded_param_state_fully_reshardable matches a grad-buffer parameter to its model ShardedTensor by object identity, and GTP puts a different tensor in the model entry twice over: a dequantized BF16 copy for native-FP8 weights, and a factory exposing the GATHERED tensor for GDN/Mamba in_proj. The latter raised "Model param ...in_proj.weight ... not in model_sharded_state_dict" and stopped the save outright. Resolve both the way _backfill_gtp_sharded_param_map already does for the distributed-Muon path — follow the _gtp_dequant_src backlink, else rebuild the per-shard ShardedTensor — and refuse the expert-parallel rebuild that would write duplicate shards across EP groups. That helper documents distributed Adam as unaffected, which holds for dp_reshardable (it never consults the model entries) but not for this format, which does. 2. What it wrote could not be read. The model side gathers GDN in_proj back to TP-local before splitting into [z|x|B|C|dt], so its keys match a non-GTP run; the optimizer did not follow, and keyed the same state '...in_proj.weight' where a non-GTP run writes '...in_proj.weight.z'. Reuse the model's own factory and gather the optimizer state across the GTP group to the width it expects, stripping the alignment pad. The state arrives on CPU from the DP gather while the GTP group is NCCL, so the collective runs on the parameter's device and returns; mem_efficient is off under GTP, which means the world tensors reach every rank and the collective is symmetric. Name the parameter when the match fails. The error printed the tensor's values and neither the parameter nor what the identity-keyed map actually holds, which is what a reader needs to see that the model exposed a different tensor in its place. CPU tests pin the resolver's contract: plain params stay the caller's problem, the dequant backlink and the by-name factory match resolve, and the EP-unaware rebuild refuses expert params. Verified on the VL MTP proxy: the checkpoint saves with no dp_group_idx keys, a same-topology resume loads and continues, and a 3D-saved checkpoint loads into a GTP run. A no-GTP control saves cleanly in this format, which is what established the gap was GTP's rather than the format's. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> fix(gtp): checkpoint GDN input projection across GTP shards The semantic [q|k|v|z|beta|alpha] boundaries of GatedDeltaNet's in_proj can cross GTP shard boundaries, so a per-shard save would write a layout that depends on the save-time GTP degree. Save gathers the physical shards back to the logical TP-local projection before the factory split; load pads then slices back to this rank's shard. Both reuse the shared ssm/utils helpers that MambaMixer's in_proj already goes through, so the checkpoint layout is identical to a non-GTP run and independent of the topology it was written from, and the GTP rank rides in replica_id so writer election stays correct when dp_cp_group excludes the GTP axis. Upstream wires GTP into mamba_mixer but not gated_delta_net, which is the attention variant Qwen3.5 uses. Test builds a GatedDeltaNet under TP2 x GTP2 with nonzero alignment padding and verifies the saved factory splits into the six logical TP-local sections, the load-side merge reconstructs this rank's physical shard (pad rows re-zeroed), and exactly one GTP peer wins DCP writer election per chunk. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> refactor(ssm): share the GTP fused-projection checkpoint helpers Mamba's in_proj GTP checkpoint handling — gather the GTP shards back to the logical TP-local width before the semantic split on save, re-pad and slice the merged tensor back to this rank's shard on load — is exactly what GatedDeltaNet needs next, so hoist both blocks into megatron/core/ssm/utils.py alongside _split_tensor_factory and switch MambaMixer to the helpers. Two deliberate strengthenings ride along: - The GTP rank is folded into replica_id ((0, gtp_rank, dp_cp_rank) instead of (0, 0, dp_cp_rank)). The gathered tensor is replicated across the GTP peers, and when metadata['dp_cp_group'] is a GTP-excluded group (explicit pg_collection grids pass pg_collection.dp_cp) the peers share a dp_cp rank — without the fold, two ranks would both win DCP writer election for the same chunk. mpu-derived grids, where dp_cp subsumes the GTP axis, are unaffected: the extra slot only ever separates ranks that previously collided. - The load-side merge refuses a non-2-D tensor instead of padding/slicing a flattened buffer; only the unflattened model-weight factory is supported (optimizer state resolves through the per-shard rebuild, never this merge). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> test(gtp): pin the zero_out_wgrad contract for GTP's dummy wgrad The fix itself now comes from upstream #6565, which landed the same zero= pairing in _handle_megatron_grad_accum. These two unit tests are ours and upstream has no equivalent: one asserts the flag is forwarded to get_dummy_wgrad at all, the other guards the consequence by running DDP's accumulation by hand and checking main_grad stays finite. They are cheap single-process tests, unlike #6565's full MTP run. fix(gtp): recycle the wgrad reduce-scatter inputs, not its outputs wgrad_reduce_scatter() rebound one name to both the pool-owned full-size RS inputs and the shard-sized RS outputs, then returned the outputs to the pool that _wgrad_pool_get() serves full-size buffers from. Split into wgrad_inputs and reduced_wgrads so only the pool's own buffers go back to it. Test asserts the synchronous batched path returns the original input buffers to the wgrad pool. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> fix(gtp): restore replicated TE bias width after GTP presharding GTP remat shards weights only, but TE sizes both weight and bias from the pre-sharded out_features it is constructed with, leaving every bias at GTP-shard width. Resize each bias back to the logical TP-local width after construction. Refuses the two layouts this does not model — TE single_grouped_bias and parameters_split — rather than silently mis-sizing them. Upstream GTP has no bias-size handling at all, so any GTP run with add_bias_linear would carry short biases into the first forward. GPU-free tests cover the resize, the no-op when already correct, the disabled-bias placeholder, the grouped case, and each refusal path. The distributed test drives row/column/padded TE linears end to end through forward and backward, including the bias-disabled placeholder Qwen3.5 GDN uses. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> fix(multimodal): align Qwen3.5-VL RoPE wrappers Match the current MCore RoPE callable contract: both fp32 wrappers now take mla_rotary_interleaved / inverse / mla_output_remove_interleaving / max_seqlen in the canonical parameter order and forward them through the BSHD and THD paths. The wrappers replace the module-level apply_rotary_pos_emb, so a signature that diverged in position 7 (max_seqlen where upstream has mla_rotary_interleaved) could silently misbind a positional caller. Adapted from 2a00f99b52 on feat/qwen35-gtp-clean rather than taken verbatim: this tree has since gained a fused mRoPE THD fast path that the original body predates and would have removed. The signature contract is kept; the stale body is not. Guard the fused path on the new options. fused_apply_mrope_thd implements neither the inverse rotation nor MLA output de-interleaving and takes no mla_rotary_interleaved argument, so requesting any of them must fall through to the unfused wrapper instead of silently computing the wrong rotation. fix(multimodal): derive MockQwen35VLDataset samples from the sample index The dataset drew from ambient CPU RNG and ignored idx. Ranks share the CPU seed unless data_parallel_random_init is set, so every rank of a data-parallel group drew the same stream and the distinct samples in a global batch collapsed to the per-rank microbatch count — a number that changes with the parallel layout, so two configurations being compared saw different data. Each sample is now derived from (seed, split, idx). The split name keeps the train/valid/test datasets apart: sharing one stream would make them emit identical samples, trading duplication across ranks for duplication across splits. The provider passes args.seed through, which the ambient-RNG version got for free. The parts are hashed with blake2b rather than added. Addition overlaps by construction: idx=1 under seed S is idx=0 under seed S+1, and each split would replay the previous one shifted by one index. Since idx now determines content, out-of-range indices raise IndexError instead of silently hashing to a plausible sample from outside the dataset. (cherry picked from commit 716a9e7249cbcfa6a2197772a52c4cbe58522b05) fix(gtp+mtp): zero the placeholder wgrad when zero_out_wgrad is set (#6565) (cherry picked from commit d262dbed3d48a0b6a22b1f0ddc58c2cede2b9149) [fix] fix GTP+DCP ckpt saving/loading for GDP module (#6503) (cherry picked from commit 8ba5c60d2a) Import block adapted to this base: it keeps is_using_quantization_scales, which upstream's tree does not have, alongside the make_tp_sharded_tensor_for_checkpoint this change adds. [fix] GTP+DDP potential data-race: publish DDP params before GTP prefetch reads them (#6388) Co-authored-by: Jiangfei Duan <jiangfeid@nvidia.com> (cherry picked from commit 87cc590ba03b7490ca2ab27dfc2372cd4d24677d) Gated Delta Product (GDP) implementation (#6074) (cherry picked from commit 8e70f69827, adapted for the dev base: - hybrid_flops/gdn_layer_flops grafts re-done inside the dev-evolved functions (DSv4/MLA/MTP-aware signatures); gdn_layer_flops keeps the dev signature (no main-only use_gdn2) - main-only context tests (maybe_save_dataloader_state) not carried; tests/unit_tests/transformer/test_transformer_config.py recreated with only the GDP tests since dev deleted the file - _equivalent_pair hybrid fixture gains gdp_num_householder, matching what upstream did to _make_hybrid_args) Balance LayerWise optimizer shards by Newton-Schulz cost, not parameter size (#6379) Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> (cherry picked from commit 85344907be6a9985410bf528acbedaca77b2a020) GTP + PartialCG: Classify latent projections for partial CUDA graphs (#6446) (cherry picked from commit 81fe7c746f6556e5d8dc24196a1358ef0898dbf7) [fix] GTP+recompute: keep adjacent GTP recompute weights off the same gather buffer (#6407) (cherry picked from commit e22bcb09b91375639588217c12e96242e8e48c36) feat(gtp): opt-in GTP_remat sharding support for 'moe-latent-proj' (#6383) (cherry picked from commit fb24e8707c94aebe28fb919e4a6253b600ad0528) [feat] GTP+MTP (#6242) (cherry picked from commit 79e6b6a7f525a9fb42a8adcb347534b4fc791b0a) Perf: skip per-param copy_ dispatch in the MXFP8 param copy-back (#6094) (cherry picked from commit 8e57bb642344f60e04bec0d7a79e5fd66c5c0023) Enforce that the number of optimizer shards used in layout computation is the same used during the training iteration (#6048) (cherry picked from commit 8a424b83777d4bdc8c3c857b1cb1d9ce141818af) [GTP] Cleanup usage of pg_collection in gtp (#6234) Co-authored-by: ykarnati <ykarnati@nvidia.com> (cherry picked from commit 9f3fe8c925fb3588356c06300071cd077b90359a) Populate dp process group in auto-built ProcessGroupCollection in pipeline schedules (#5901) Co-authored-by: Claude Fable 5 <noreply@anthropic.com> (cherry picked from commit 6f7bcd48f4fccfea74fbd2385502b2fa8bb05251) Fix GTP reduce-scatter overlap across local CUDA graphs (#6060) (cherry picked from commit 9fa816bfd9dd994dc5acf6012fb4e841a56cfa4a) Make Bridge Communicator aware of GTP (#6263) (cherry picked from commit 9a3c40b25408f9c5f4ed87500989ebe299f88fea) [feat] Add reduce-scatter-with-fp32-accumulation support for GTP (#6200) (cherry picked from commit 23b2ff22eea5a150fec2b41b4d035afa216c8040) Fix gradient reduction issue when EP=1, EP=TP, and EGTP != GTP (#6080) (cherry picked from commit abf04f46e39a22f8e6f2597ba820a7712aef36c2) Fix gradient-norm undercounting when using EP and TP (#5916) (cherry picked from commit cd4afffa648426a959dc7cb1e24b5ce7d0c3ff54) [fix] Give same-key GTP chain neighbours distinct gather buffers (#6207) Co-authored-by: Jiangfei Duan <jiangfeid@nvidia.com> (cherry picked from commit c2fd8275d9fa13cd7d2e51c9d9370fbb6cdae02d) Fix GTP full-iteration CUDA-graph capture regression (#6077) Co-authored-by: Claude Opus 5 <noreply@anthropic.com> (cherry picked from commit ec53920e605b7d5cd84d182981db3c83439a4fad) [GTP][Feat] Add one-block-ahead prefetch for GTP grouped-expert weights (#6057) (cherry picked from commit 543b8039646f2cdb7a54eacad612849080292bd6) [feat] Generalized Tensor Parallelism (GTP) (#4967) Co-authored-by: Jieming Zhang <jiemingz@nvidia.com> Co-authored-by: Deepak Narayanan <deepakn94@gmail.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Co-authored-by: Jiangfei Duan <jiangfeid@nvidia.com> (cherry picked from commit c5ff22b7f11822e0a36526b0434708d99c8206b0) ci(auth): treat svcnemo-autobot as internal (#6437) [dev] [Deepseek-V4] Fix fused CSA indexer loss normalization and compact attention indices (#6349) fix(ci): redact sensitive environment values from test output (#6408) [Dev] Update moe recipes (#6335) [dev] Optimize MTP contiguous packed-CP rolls (#6246) [dev] Refactor CSA structure: Move CSA implementation helpers into csa_utils dir (#6372) fix(optimizer): AUT-1363 avoid grad threshold check during capture (#6355) Co-authored-by: Hristo Filaretov <hello@hgf.sh> Co-authored-by: Michał Marcinkiewicz <43240942+mmarcinkiewicz@users.noreply.github.com> fix(ci): AUT-1396 stabilize docker cache and external test queue (#6348) cp: chore(ci): bump community workflow to v1.8.8 (#6345) Co-authored-by: oliver könig <okoenig@nvidia.com> build: AUT-1363 verify downloaded yq binaries (#6344) fix(deps): AUT-1323 use released dev dependency sources (#6273) Co-authored-by: NVIDIA NeMo Bot <nemo-bot@nvidia.com> [DSv4] Use the full CSA denominator for unfused indexer loss (#5960) Fix undefined get_data_and_context_parallel_group references (#6072) Add standard MLite local validation (#5888) fix(dsa): skip invalid slots in sparse indexer backward (#6166) [Dev] Keep the mHC mapping computation in fp32 on the fused cuTile path (#6172) [codex] Propagate SBHD padding masks across pipeline stages (#5544) [Dev] Support full activation recompute for EP A2A overlap (VPP-stage recompute) (#5869) Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> Dispatch the CSA Compressor gated pooling to the cudnn-frontend fused kernels (THD) (#5984) Fix/dsv4 cp relax assert offline packed thd (#6158) Co-authored-by: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com> fix(tensor-parallel): AUT-1177 deduplicate expert shards by expert TP (#6165) [dev] Fix Torch checkpoint loading for updated layer-wise optimizer (#6073) [dev] [fix] fix optimizer_cpu_offload with mark_keep_in_fp32 (#6124) [codex] Preserve real and padded THD sequence lengths (#5541) [dev] [fix] propagate LayerWise optimizer flag through model builders (#6082) [Dev] Migrate main-first DSA support to dev (#6020) Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: Pingtian Li <pingtianl@nvidia.com> build: serialize uv dependency installation (#6089) [Dev] Add Triton fused mRoPE for Qwen3.5-VL (#5962) Co-authored-by: Li Tao <lit@nvidia.com> [dev] Align thd padding logic in eager & graph mode (#5724) ci(actions): AUT-977 retry transient log artifact uploads (#6026) fix(deps): AUT-960 use optional-attribution NVRX revision (#6006) test(gdn): AUT-983 quarantine flaky chunkwise CP backward checks (#6040) [dev] moe(perf): Support fused pre-GDR path for chunkwise CP. (#5638) docs(skills): clarify container::lts is the older LTS PyTorch base (#6009) Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> [dev] moe(fix): Restore per-token MTP loss logging (#5967) [Dev] Skip no-op BSHD padding masks in Qwen3.5-VL (#5964) [dev] [fix] fix optimizer state offloading (#5923) [dev]: Fix mHC boundaries in EP overlap schedule (#5471) Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> [Dev] Fix max_seqlen forwarding in Qwen3.5-VL vision RoPE (#5890) Co-authored-by: Li Tao <lit@nvidia.com> chore: nightly sync main into dev (13_07_2026) (#5784) Co-authored-by: Yashaswi Karnati <144376261+yashaswikarnati@users.noreply.github.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> Co-authored-by: Shanmugam Ramasamy <111910568+shanmugamr1992@users.noreply.github.com> Co-authored-by: Philip Petrakian <ppetrakian@nvidia.com> Co-authored-by: wdykas <73254672+wdykas@users.noreply.github.com> Co-authored-by: Keshav Santhanam <ksanthanam@nvidia.com> Co-authored-by: Laura Dang <lauradang.2000@gmail.com> Co-authored-by: oliver könig <okoenig@nvidia.com> Co-authored-by: Deepak Narayanan <dnarayanan@nvidia.com> Co-authored-by: Teodor-Dumitru Ene <34819528+tdene@users.noreply.github.com> Co-authored-by: Chen Cui <chcui@nvidia.com> Co-authored-by: shanmugamr1992 <shanmugamr1992@gmail.com> Co-authored-by: Charlie Truong <chtruong@nvidia.com> Co-authored-by: Jingyue Wu <wujingyue@gmail.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Co-authored-by: ℍ𝕠𝕝𝕝𝕠𝕨 𝕄𝕒𝕟 <hollowman@opensuse.org> Co-authored-by: muyihao <37872457+muyihao@users.noreply.github.com> Co-authored-by: Tom Long <tolong@nvidia.com> Co-authored-by: Siddhartha Raman Sundara Raman <sraman@nvidia.com> Co-authored-by: Aditya Singh <60082699+adityasingh2400@users.noreply.github.com> Co-authored-by: yeyu-nvidia <yeyu@nvidia.com> Co-authored-by: Zhongbo Zhu <42691305+zhongbozhu@users.noreply.github.com> Co-authored-by: Laura Dang <laurad@nvidia.com> Co-authored-by: Antoni-Joan Solergibert <asolergibert@nvidia.com> Co-authored-by: Lawrence McAfee <85179052+lmcafee-nvidia@users.noreply.github.com> Co-authored-by: janEbert <janpabloe@nvidia.com> Co-authored-by: Anil Thomas <anlthms@users.noreply.github.com> Co-authored-by: Ritesh Patel <ripatel@nvidia.com> Co-authored-by: Maanu Grover <maanug@nvidia.com> Co-authored-by: Cory Ye <44509866+cspades@users.noreply.github.com> Co-authored-by: ma-jh <3364870135@qq.com> Co-authored-by: Jenny Chen <jennifchen@nvidia.com> Co-authored-by: realAsma <akuriparambi@nvidia.com> Co-authored-by: Hongbin Liu <lhb8125@users.noreply.github.com> Co-authored-by: Xin Yao <xiny@nvidia.com> Co-authored-by: Jimmy Zhang <133159885+jiemingz@users.noreply.github.com> Co-authored-by: Jingyue Wu <jingyuew@nvidia.com> Co-authored-by: Ajay <abalasa@nvidia.com> Co-authored-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: Dmytro Pykhtar <37850217+dimapihtar@users.noreply.github.com> Co-authored-by: mathemakitten <helenn@nvidia.com> Co-authored-by: Zang Peiyu <166481866+factnn@users.noreply.github.com> Co-authored-by: Zhiyu Li <zhiyul@NVIDIA.com> Co-authored-by: Skand Hurkat <shurkat@nvidia.com> Co-authored-by: Fei Wu <33940270+YangFei1990@users.noreply.github.com> Co-authored-by: Siddharth Singh <sidsingh@nvidia.com> Co-authored-by: Yongqiang Wang <yongqiang.seagull@gmail.com> Co-authored-by: Minh Vu <vuhoangminh97@gmail.com> Co-authored-by: Guihong Li <guihongl@nvidia.com> Co-authored-by: Asadbek Xodjayev <100586658+asadbekXodjayev@users.noreply.github.com> Co-authored-by: asadbekXodjayev <matyoqub18@gmail.com> Co-authored-by: Björn Buschkämper <bjoern.buschkaemper@gmail.com> Co-authored-by: Zhengmao Ye <yezhengmaolove@gmail.com> Co-authored-by: Jon Barker <jbarker@nvidia.com> Co-authored-by: claude[bot] <41898282+claude[bot]@users.noreply.github.com> Co-authored-by: Deyu Fu <deyuf@nvidia.com> [dev] Lite: /bump lite DeepSeek v4 RL to dev (#5862) Co-authored-by: leikaixiang.shyoshyo <leikaixiang.shyoshyo@bytedance.com> Avoid per-segment GPU->CPU syncs in the unfused DSA THD helpers (#5884) enable tms no cpu backup region for param buffers (#5103) Co-authored-by: Yueming Yuan <yym022502@gmail.com> Co-authored-by: lilei <799812479@qq.com> Co-authored-by: Deyu Fu <deyuf@nvidia.com> [Dev] Fix distributed optimizer state save/load for mixed-dtype param groups (#5835) Co-authored-by: Claude Fable 5 <noreply@anthropic.com> fix: preserve static hybrid module specs fix: resolve callable hybrid specs fix: address post-merge unit test regressions Fix zero-loss indexer gradients for DSv4 CP (#5809) Merge remote-tracking branch 'origin/dev' into main2dev/13_07_2026 # Conflicts: # megatron/training/arguments.py # pretrain_hybrid.py # tests/unit_tests/data/test_get_batch.py fix: restore MLA TP fusion and DeepEP v2 SM config [dev] Enable DeepSeek-v4 hybrid_model Part (3/N) (#5762) [dev] Muon fp8 param gather for decoupled layout (#5470) Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> fix: restore dev training-loop behavior fix: align pretrain provider call order fix: forward padding mask in fine-grained MoE routing Keep split A2A schedule routing consistent with the normal forward path so dropless HybridEP excludes padded tokens in both executions. fix: post-CI corrections (unused imports + AbsorbedMLA submodule fields) Fixes for the first CI run's failures: 1. Remove unused imports flagged by the linting job (pylint W0611): - dynamic_engine.py: CUDAGraphBatchDimensionBuilder, InferenceBatchDimensions - gated_delta_net.py: ShardedTensor, ReplicaId, ShardedTensorFactory - transformer_block.py: checkpointed_forward 2. Fix AbsorbedMLASelfAttentionSubmodules construction to match dev's split K/V up-projection API (the merged absorbed_mla.py is dev's version, which splits linear_kv_up_proj into linear_k_up_proj + linear_v_up_proj): - hybrid_layer_specs.py dsa_layer + experimental_attention_variant_module_specs.py now pass linear_k_up_proj/linear_v_up_proj to AbsorbedMLASelfAttentionSubmodules. - The standard MLASelfAttentionSubmodules construction in hybrid_layer_specs.py keeps linear_kv_up_proj (that class was not split). This fixes the install-test ImportError: "AbsorbedMLASelfAttentionSubmodules.__init__() got an unexpected keyword argument 'linear_kv_up_proj'". Co-Authored-By: Claude <noreply@anthropic.com> [codex] Exclude padding tokens from MoE routing (#5542) Co-authored-by: Xin Yao <xiny@nvidia.com> fix: save chained optimizer state in torch checkpoints (#5684) [Dev] Numerical fix for moe single grouped weight with fp8 fp4 primary weight and grad norm spikes (#5464) Merge remote-tracking branch 'origin/main' into main2dev/13_07_2026 Nightly sync of main into dev (13_07_2026). Resolution highlights: - Kept dev's dependency triple (pyproject.toml, uv.lock, docker/Dockerfile.ci.dev), CODEOWNERS, base-image pin (26.04, GitLab pin), and golden values. Main's base-image bump to 26.06 (#5632) belongs in a dedicated bump PR, not the sync. - copy-pr-bot.yaml: union of trustee lists (kept dev-only sanandaraj5597, wplf). - Renamed args.hybrid_context_parallel -> args.dynamic_context_parallel and get_hybrid_data_context_parallel_groups -> get_dynamic_data_context_parallel_groups in main-originated code (training.py, pretrain_hybrid_flex.py) to match dev's canonical naming (the deprecated hybrid arg does not exist in the merged tree). - DSA subsystem: dsa_kernels.py resolved as a UNION of dev's kernel API (indexer_topk, dsa_sparse_attn, ...) and main's backend-neutral hooks (use_fused_dsa_kernels, run_fused_*); dsa.py based on main + re-added dev-only fused_qk_topk_naive_thd/_build_causal_mask_seg for csa.py. - MoE flex dispatcher: unioned backend Literal (deepep/deepepv2/hybridep/ncclep), kept both _DeepepV2Manager (dev) and _NCCLEPManager (main). - transformer_layer.py: kept dev's _forward_mlp_output_with_bias/input_ids plumbing AND main's inter-document-masking MoE reshape (_maybe_unflatten_for_moe/_maybe_reflatten_from_moe). - fine_grained_activation_offload.py + checkpointing.py: took main where the merged non-conflict common regions and callers already required main's interface (group_offload naming; process-group threading params). - MTP inference: kept dev's _decoder_hidden_states_cache path (consumed by text_generation_controller); dropped main's block-scope mtp_decoder_hidden_states approach (infra not present in merged tree). Intentional drops (skill case (a), documented with SHA): - _forward_mlp_postprocess's redundant _restore_token_dispatcher_attrs()-at-top removed by main commit 5e4fe9b3c ("Optimize memory usage of partial CUDA graphs"); replaced by main's weakref tail. Co-Authored-By: Claude <noreply@anthropic.com> Remove use of exec_module (#5744) Assign BERT CODEOWNERS to GPT team (#5746) Set is_first_microbatch when quant_recipe is configured (#5642) Set Bert TE spec q/k_layernorm to None (#5687) Co-authored-by: Philip Petrakian <ppetrakian@nvidia.com> Fix infinite recursion in abstract tokenizer special-id property aliases (#5445) Co-authored-by: asadbekXodjayev <matyoqub18@gmail.com> Co-authored-by: Dmytro Pykhtar <37850217+dimapihtar@users.noreply.github.com> Avoid X11 master port default (#5299) Fix seq_load_balancing loss with inter-document masking and MBS > 1 (#5696) Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Add forward all-gather overlap (#5513) [codex] Restore per-microbatch MTP loss logging (#5543) fix(fsdp): import os in safe_get_rank fallback (#4959) remove deprecated modules from core/dist_checkpointing (#5134) [dev] [fix] fix MTP with contiguous CP partition mode (#5706) fix: add label shifting to mlite THD packing to align SFT loss with M… (#5682) Co-authored-by: Xin Yao <xiny@nvidia.com> [dev] Fix HybridEP token equalization under torch.compile without CUDA graphs (#5668) [dev] Align Megatron Lite MoE RL execution contracts (#5694) Keep DeepSeek V4 CSA compressor and indexer in high precision under FP8 training (#5308) Co-authored-by: Hongxiao Bai <hongxiaob@nvidia.com> Add NeMo Transformer audio encoder model (#5565) Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Triton kernels - avoid recompilation and autotuning in prod (#5608) [2/2] Wiring cuDNN fused DSA kernels support with THD, CP and IndexShare (GLM5.2) (#5099) Refactor RL rollout pipeline (#5491) Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Update base image to nvcr.io/nvidia/pytorch:26.06-py3 (#5632) NCCL EP support (#5129) Add FSDP NVTX annotations (#5704) Remove some barriers in save_checkpoint_and_time (#5557) Co-authored-by: Deepak Narayanan <dnarayanan@nvidia.com> Normalize CRLF in Claude fix commands (#5712) ci: revert unify legacy scope names (#5316) (#5709) Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> deprecate common strategy (#5160) test(determinism): add determinism tests (#5041) Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com> chore: rotate oncall schedule Separate mFSDP v2 unit tests (#5640) Fix Torch FSDP2 crash: add force_all_reduce kwarg to base finish_grad_sync (#4953) Co-authored-by: Cory Ye <44509866+cspades@users.noreply.github.com> ci: Update test configurations to unify legacy scope names (#5316) Add microbatch context helper (#5652) Add cspades to oncall rotation (#5695) MoE routing analysis and metrics capture (#5220) Fuse shared expert MLP with grouped GEMM (#5604) add safe version of numpy.load (#5500) [Dev] Fix FSDP backward hooks for TE-fused experts (#5636) Fix inter-document masking crash and NaNs with TP > 1 and micro_batch_size > 1 (#5635) Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Scatter embeddings for sequence parallelism in standalone LM forwards (#5628) Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> Pre-size the all-gather buffer for inference to max capacity (#5546) Add int4 qat training support (#5108) Co-authored-by: GeLee-Q <leege233@gmail.com> Co-authored-by: fy1214 <282598660@qq.com> Co-authored-by: Gao016 <yngao016@163.com> Co-authored-by: yefei12 <xjtu_yefeichen@163.com> Co-authored-by: yzlnew <yzlnew@gmail.com> Co-authored-by: Deyu Fu <deyuf@nvidia.com> chore(beep boop 🤖): Bump (main) (2026-07-06) Ignore contributor DCO failures in Claude fix (#5625) Update copy-pr-bot.yaml [skip ci] [dev] [DeepSeek-v4] Context Parallel support (#5087) Update golden value files for GPT-3 weekly (#5459) Add smoke test notification functionality and update notify script (#5631) Add NCCL symmetric-memory staging to experimental FSDP (#5440) Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Fix smoke BERT/T5 test failures (#5629) Fix Claude reaction permissions (#5613) Document stacked dependent PR handling in split PR skill (#5496) Optimize memory usage of partial CUDA graphs (#5451) E2E heterogenous non colocated MiMo training (#5602) Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> [dev] moe(feat): support chunkwise context parallelism for GDN. (#3282) [Main][feat] Support CUDA Graph capture offloading modules (#3697) Co-authored-by: Xin Yao <xiny@nvidia.com> Add Auto Quantize in ModelOpt quantize example (#4821) Co-authored-by: realAsma <akuriparambi@nvidia.com> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com> [Dev] fix padding mask docstring (#5598) fix(tensor_parallel): _reduce returns unreduced tensor for non-contig… (#5338) [Megatron-FSDP] MaxPoolAllocator for double-buffering hybrid architectures. (#5462) Update copy-pr-bot.yaml [skip ci] Add /claude fix workflow for on-demand PR fixes (#4862) chore: rotate oncall schedule [dev] Sync Megatron Lite with the latest implementation (#5577) Update mcore skill owners (#5586) Update PR instructions (#5592) Use NVIDIA inference credentials for Claude actions (#5589) Update copy-pr-bot.yaml [skip ci] [dev] Add experimental decoupled compact LayerWise DDP layout for Muon (#5388) Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> [training migration] Finish ModelBuilder integration (#5516) Fix TEGroupedMLP pre-backward unshard in fine-grained FSDP hooks for … (#4990) Thread dp_cp/expt_dp process groups through checkpoint load path (#5579) Add CI duties to oncall (#5510) Co-authored-by: oliver könig <okoenig@nvidia.com> Deduplicate tensor-splitting utility (#5545) Fix PR template typo (#5566) Fix `isort` target Python version (#5567) build: bump transformer-engine to release_v2.16.post (#5517) Implement async scheduling for dynamic inference (#5453) Enable Deepseek-v4 hybrid_model in dev branch Part (2/N) (#5485) [CI] Fix `gpt_dynamic_inference_tp2_pp2_ep2_gptoss_20b_swa` tests (#5527) Update copy-pr-bot.yaml [skip ci] Preserve DSA output across fused inverse RoPE (#5526) Co-authored-by: Kaixiang Lei <5780122+shyoshyo@users.noreply.github.com> Add inter-document attention masking to GPTDataset (#5298) Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> ci: pin HF_HUB_CACHE to bind-mounted cache for gpt-oss-20b inference test (#5512) Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> ci: Use GB300 for Github CI tests (#5520) [dev] moe(perf): Pre-GDR kernel fusion (#5361) Add CUDA graph training iteration test (#5417) ci: cache-from a single coherent buildcache donor (#5509) Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Update copy-pr-bot.yaml [skip ci] test: restore G/G + lag=19 for gpt_grpo_tp4_pp1_dp2_8b throughput tests (#5514) [Main] Generalized fix for mxfp8 param gather (#5236) Improve default dynamic CP packing scheduler (#5154) Add --qad-train-target {base|mtp|both} for QAD / MTP QAT (frozen-base, frozen-MTP, or co-train) (#4785) Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> fix: set DATA_PATH for moe-dynamic-inference recipe (#5506) Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> [dev] [DeepSeek-v4] Packed Sequence (THD) support for DSv4 Hybrid Attention (#5011) Add hybrid FSDP unit module support (#4329) Fix NameError in is_flashinfer_min_version when check_equality=False (#4961) chore: nightly sync main into dev (22_06_2026) (#5430) Co-authored-by: Teodor-Dumitru Ene <34819528+tdene@users.noreply.github.com> Co-authored-by: Asha Anoosheh <aanoosheh@nvidia.com> Co-authored-by: Jorge Albericio <jalbericiola@nvidia.com> Co-authored-by: Pranav Thombre <pthombre@nvidia.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: janEbert <janpabloe@nvidia.com> Co-authored-by: Philip Petrakian <ppetrakian@nvidia.com> Co-authored-by: mathemakitten <helenn@nvidia.com> Co-authored-by: Yashaswi Karnati <144376261+yashaswikarnati@users.noreply.github.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Co-authored-by: Shijie <505749828@qq.com> Co-authored-by: Ajay <abalasa@nvidia.com> Co-authored-by: oliver könig <okoenig@nvidia.com> Co-authored-by: Antoni-Joan Solergibert <asolergibert@nvidia.com> Co-authored-by: Deepak Narayanan <dnarayanan@nvidia.com> Co-authored-by: Tom Long <tolong@nvidia.com> Co-authored-by: Keshav Santhanam <ksanthanam@nvidia.com> Co-authored-by: Teodor-Dumitru Ene <teodord.ene@gmail.com> Co-authored-by: Siddhartha Raman Sundara Raman <sraman@nvidia.com> Co-authored-by: Jingyue Wu <wujingyue@gmail.com> Co-authored-by: ℍ𝕠𝕝𝕝𝕠𝕨 𝕄𝕒𝕟 <hollowman@opensuse.org> Co-authored-by: Hongbin Liu <lhb8125@users.noreply.github.com> Co-authored-by: Charlie Truong <chtruong@nvidia.com> Co-authored-by: Lawrence McAfee <85179052+lmcafee-nvidia@users.noreply.github.com> Co-authored-by: wdykas <73254672+wdykas@users.noreply.github.com> Co-authored-by: claude[bot] <41898282+claude[bot]@users.noreply.github.com> Fix fused MLA down projection with tensor parallelism (#5383) [split 4/4] Enable DSA CP and THD hooks (#5246) Merge cu_seqlens across micro-batch for THD attention (#5454) Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Force RL inference to CP=1 (#5423) fix: restore fused group MLP offload in main2dev sync (#5493) ci: auto-retry test-data download in container-build job (#5498) Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> [dev] Megatron Lite (4/4) shared attention (#5427) fix sequence packing wrapper for eval (#5483) Add Nemotron6-MoE VLM model provider for MIMO example (#5374) Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> Add MIMO forward step and per-token loss for hetero training (#5376) Narrow oncall responsibilities (#5490) fix: post-CI corrections for sync test/impl reconciliation 1) test_optimizer.py: reverted to dev. The sync kept dev's multi_latent_attention.py (split q/kv down-proj, no _synthesize_fused_qkv_down_weight), but auto-merged main's test asserting the fused linear_qkv_down_proj.weight key. 2) training.py: guard the dev-only sequence_packing_scheduler config access with getattr (lines in train_step and train()). main's new MIMO schedule-plumbing test (#5333) passes an empty SimpleNamespace config; the reconciled training.py keeps dev's packing path, so the access must tolerate a config lacking the attribute. Real configs are unaffected (getattr returns the same value). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Thread process groups through training checkpoint paths (#5486) Clean up training.py module header (dedupe + reorganize imports/globals) (#5469) Automated community request assignment (#5147) Fix merges_file kwarg name in HuggingFaceTokenizer (#5406) ci: check megatron.training imports in installation test (#5458) [split 2/4] Scale DSA indexer loss in pipeline schedules (#5244) build: install flash_mla from source in the CI image (#5481) Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> ci: launch GB200 unit tests via launch_on_gb200 marker (#5477) Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Add inference functions to support MCore-/MBridge- training refactor and remove legacy modelbuilder functions (#5169) Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com> chore: rotate oncall schedule test: mark gated_delta_net selective-recompute test flaky_in_dev (#5476) Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> test: mark TestParallelTransformerBlockCudagraphs::test_gpu_cudagraph flaky_in_dev (#5475) Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Add experimental Megatron-FSDP fully_shard implementation (#5387) Support HybridModel feature specs in ModelOpt (#5354) feat(inference): default use_coordinator to True in high-level APIs (#5326) Add hetero grid args and MoE process groups for MIMO example (#5375) Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> Support SWA and sink attention in dynamic inference (#5249) Co-authored-by: shanmugamr1992 <shanmugamr1992@gmail.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Shanmugam Ramasamy <111910568+shanmugamr1992@users.noreply.github.com> Remove DBuffer mesh axis validation (#5441) ci: Set test_save_verify_integrity_manifest_directly as flaky (#5468) Add logprobs_mode (raw/processed) to inference config (#5419) build: point flash_mla at the nv_dev branch (#5448) Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Rename CP batch helpers to describe balancing granularity (#5403) Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Add --functional-test-name to trigger_internal_ci (#5449) Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> test: mark ep_a2a_overlap activation-offloading test flaky_in_dev (#5450) Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> [Fix] Fix MoE router z-loss compatibility with TE CUDA Graph capture. (#5401) Co-authored-by: yangfan.bai <yangfan.bai@shopee.com> Add RL rollout submission and consumption granularity controls (#5306) Add MIMO dual gradient finalization (colocated + non-colocated) (#5286) Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Clean up MTP inference control flow (#5418) Add RADIO vision encoder wrapper for MIMO example (#5397) Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> Disag MR3: Add heterogeneous KV/Mamba reshard planners (#5188) Merge remote-tracking branch 'origin/main' into main2dev/22_06_2026 Nightly sync of main into dev (22_06_2026). Resolves 16 conflicts preserving dev features (pre-push guard: 0 dropped dev lines); brings in main's inference shard-spec API additively. Supersedes #5429. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Consistent oncall schedule (#5404) Stabilize hybrid_2b GB200 perf test against run-to-run noise (#5364) Support the MIMO cross-grid path in training loop (#5373) Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> Disag MR1: Add inference shard specs and pg-collection building (#5186) Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> Disag MR2: Refit into multiple destination pools and tied-embedding + UVM fixes (#5187) Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> Fix Mamba prefix match for chunked prefill (#4758) chore(beep boop 🤖): Bump (main) (2026-06-22) Add --mamba-training-ssm-states-dtype argument (#5309) Co-authored-by: Jorge Albericio <jalbericiola@nvidia.com> Add MIMO runtime setup: per-role RNG seeding and DDP wrapping (#5285) Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> Update goldens for weekly tests after pytorch and TE bumps. (#5399) Revert "Remove checkpoint-time GPU cache reclaim workaround (#5170)" (#5366) Add flaky marker to fine-grained activation offloading test (#5350) (#5368) ci: Remove sync skills workflow (#5091) chore: rotate oncall schedule [dev] moe(perf): Restore fused GDN THD all-to-all on dev (#5389) [split 3/5] Refactor absorbed MLA projection handling (#5245) Add MimoModel.zero_grad_buffer delegating to active DDP submodules (#5372) Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> Thread tensor-parallel group into the RADIO patch embedder (#5371) [feat] Support fine-grained activation offloading in fused group mlp (#5082) Remove unused distributed pytest markers (#5380) [split 1/5] Fix packed THD RoPE under CP (#5243) [Dev] restore DSv4 tflops calc in training and fix the packed seq case (#5358) Add minimal DBuffer implementation (#4835) Document agent PR commit sign-off and signing (#5381) Update copy-pr-bot.yaml [skip ci] Support fused MLA QKV checkpoint reload (#5310) Profiling (#3110) Co-authored-by: Teodor-Dumitru Ene <teodord.ene@gmail.com> Make Megatron RL only materialize last token logit (#4551) Expand the Mamba prefix caching memory safety check to include scratch space buffers (#5348) Add full model cuda graph support for MTP inference (#4950) Clean up pretrain_gpt.py and pretrain_hybrid.py formatting and remove module globals (#5351) Fix memory leak with log_max_attention_logit (#4699) (#5067) Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: Deepak Narayanan <dnarayanan@nvidia.com> [Dev] add cuda graph support for thd format training. (#4359) Co-authored-by: Haochen Yuan <haocheny@login-eos01.eos.clusters.nvidia.com> ci: default functional test time limit to 4h for release/weekly scopes (#5360) [Dev] Add Megatron-FSDP weight prefetch for full recompute (#5175) [Fix] Fix optimizer parameter override bugs. (#5213) Co-authored-by: yangfan.bai <yangfan.bai@shopee.com> Co-authored-by: Xin Yao <xiny@nvidia.com> feat(ssm): whole-module 'gdn' selective recompute for GatedDeltaNet (#5296) Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: Xin Yao <xiny@nvidia.com> chore: nightly sync main into dev (12_06_2026) (#5314) Co-authored-by: megnvidia <mmiranda@nvidia.com> Co-authored-by: Philip Petrakian <pgpetrak@gmail.com> Co-authored-by: Xuanteng Huang <44627253+xuantengh@users.noreply.github.com> Co-authored-by: Ajay <abalasa@nvidia.com> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Co-authored-by: Pingtian Li <158665726+Wohox@users.noreply.github.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Co-authored-by: Dmytro Pykhtar <37850217+dimapihtar@users.noreply.github.com> Co-authored-by: Santosh Bhavani <santosh.bhavani@live.com> Co-authored-by: Cory Ye <44509866+cspades@users.noreply.github.com> Co-authored-by: Yi-Fu Wu <yifu.wu@gmail.com> Co-authored-by: Yan Xu <45385219+Connor-XY@users.noreply.github.com> Co-authored-by: Yashaswi Karnati <144376261+yashaswikarnati@users.noreply.github.com> Co-authored-by: Jingyue Wu <wujingyue@gmail.com> Co-authored-by: lichenlu <lichenlu_8618@163.com> Co-authored-by: peibli <lipeibao@126.com> Co-authored-by: Charlie Truong <chtruong@nvidia.com> Co-authored-by: shurkat-nvidia <shurkat@nvidia.com> Co-authored-by: Philip Petrakian <ppetrakian@nvidia.com> Co-authored-by: oliver könig <okoenig@nvidia.com> Co-authored-by: Lintch <44701395+returnL@users.noreply.github.com> Co-authored-by: janEbert <janpabloe@nvidia.com> Co-authored-by: Siddhartha Raman Sundara Raman <sraman@nvidia.com> Co-authored-by: Siddhartha Raman Sundara Raman <270218152+sraman-rgb@users.noreply.github.com> Co-authored-by: Shanmugam Ramasamy <111910568+shanmugamr1992@users.noreply.github.com> Co-authored-by: Anish Mahishi <amahishi@nvidia.com> Co-authored-by: svcnvidia-nemo-ci <svcnvidia-nemo-ci@users.noreply.github.com> Enable Deepseek-v4 hybrid_model in dev branch Part (1/N) (#5042) Co-authored-by: Guihong Li <guihongl@oci-hsg-cs-001-vscode-02.cm.cluster> Co-authored-by: Yan Xu <yxu1@nvidia.com> Co-authored-by: hx <hongxiaob@nvidia.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Thread MIMO support through the stock training loop (schedule + optimizer) (#5333) Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> fix: preserve seqlen stats in train_step Add zstandard package to Docker LTS requirements. Fix nightly failures (#5347) Merge branch 'dev' into main2dev/12_06_2026 Fix LatentMoE theoretical memory estimate (#5145) [dev] bump emerging optimizers to v0.3.0 (#5320) chore(beep boop 🤖): Bump (main) (2026-06-15) fix tflops calculation when sequence_packing_scheduler is not none (#5342) [Dev] Add DeepEP v2 flex dispatcher backend (#4793) Co-authored-by: Dennis(Zhenhuan) Liu <denliu@nvidia.com> [Dev] Add MoE recipe performance summary (#5289) Co-authored-by: Dennis Liu <denliu@denliu.nvidia.com> Thread pg_collection through wrap_model_chunks_with_ddp (#5328) Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Account for reasoning token stripping (#5313) Inference: Cudagraph-aware admission gating in prefill scheduler (#4870) Fix crash due to tool call at sequence length (#5302) Co-authored-by: Jorge Albericio <jalbericiola@nvidia.com> Fix EP=1 inference by allocating buffers anyway (#5233) Add code owners for optimizer-related files (#5297) Co-authored-by: Philip Petrakian <ppetrakian@nvidia.com> Add moe loss normalization for RL SFT (#3956) Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Allow for pre-bound socket to be passed in server (#5301) Handle None values in sampling parameters (#5300) Co-authored-by: Jorge Albericio <jalbericiola@nvidia.com> Offline Logits-Based Knowledge Distillation (#5019) [dev]: faster implementation of mHC fused kernels (#4624) Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com> fix: restore dev-only training.py features dropped by the main override The merge takes main's training.py wholesale (per the sync skill's override list), which silently reverted four dev-only features whose supporting core code the merge keeps at dev's version. Each surfaced as a CI failure: 1. Dynamic context-parallel API rename. Dev renamed get_hybrid_data_context_parallel_groups -> get_dynamic_data_context_parallel_groups (identical signature) and args.hybrid_context_parallel -> dynamic_context_parallel. Point training.py at dev's names. Fixes the conftest.py ImportError that cascaded to every unit-test bucket. 2. Dynamic-CP / sequence-packing data loading. Replace main's setup-time HybridCPDataLoaderWrapper wrap with dev's per-step wrap_data_iterator (gated on config.sequence_packing_scheduler), which returns a RerunDataIterator-compatible iterator. Fixes the *_cp4_dcp RerunDataIterator assertion. 3. MTP loss logging scale. Dev's MTPLossLoggingHelper stores raw loss sums and token counts and computes the per-token loss after reduction, so the log scale must be 1.0; main's 1/get_num_microbatches() divided the reported mtp_N loss by num_microbatches. Fixes moe gpt3_..._scoped_cudagraph mtp_1 loss (was 16x too small: 0.682 vs golden 10.915). 4. DSA indexer loss cross-PP reduction. dsa.py requires num_layers (and csa_compress_ratios) so first-pipeline-stage ranks lazily initialize the tracker and join the cross-PP all_reduce in reduce_loss_in_tracker; main's call omitted them, so stage 0 returned early and the last stage hung on an unmatched all_reduce. Fixes the gpt3_..._dsv4_hybrid_mhc_mtp NCCL timeout. ci: Allow DCO check in merge queue and add DCO requirement (#5305) Co-authored-by: Charlie Truong <chtruong@nvidia.com> Merge remote-tracking branch 'origin/main' into main2dev/12_06_2026 # Conflicts: # README.md # megatron/core/optimizer/optimizer.py # megatron/core/ssm/gated_delta_net.py # megatron/core/transformer/multi_token_prediction.py # megatron/core/transformer/transformer_config.py # megatron/elastification/pretrain_hybrid_flex.py # megatron/training/training.py # megatron/training/utils/common_utils.py # tests/test_utils/recipes/h100/t5.yaml # tests/test_utils/recipes/moe2.0.yaml # tests/unit_tests/ssm/test_gated_delta_net.py # tests/unit_tests/transformer/test_multi_token_prediction.py # uv.lock Thread pg_collection into train_step reductions (#5259) Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Clip mtp grads separately when mtp_detach_heads=True (#4116) Co-authored-by: Anish Mahishi <amahishi@nvidia.com> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Stabilize hybrid nanov3 gb200 perf (#5295) Enable non-deterministic results in model configuration for nemotron tests (#5239) Thread pg_collection into get_model DDP bucket sizing (#5250) Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> ci: limit retries on unsuccessful test launches (#5275) Fix fused MLA delayed weight grad hooks (#5273) Co-authored-by: Siddhartha Raman Sundara Raman <270218152+sraman-rgb@users.noreply.github.com> Fix Dockerfile warnings (#4856) fix(ci): resolve t5 dataloader stall + GRPO cudagraph-memory regression (CI-validated) (#5280) Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> [dev] Add experimental Megatron Lite as agentic exploration (#4885) Co-authored-by: Deyu Fu <deyuf@nvidia.com> Remove duplicate nccl_allocator import (#5057) [examples] Add dynamic context parallel benchmark example (#5123) [TE] Restore original CP group after dynamic CP forward in TEDotProductAttention (#5215) Co-authored-by: rionawang <rionawang@tencent.com> [Dev] Cherry-pick MTP detach heads (#5223) [Dev] Generalized fix for mxfp8 param gather (#4994) Co-authored-by: Xin Yao <xiny@nvidia.com> [Dev] DeepSeek-V4-Flash recipe 20260610 (#5266) Remove checkpoint-time GPU cache reclaim workaround (#5170) Co-authored-by: Philip Petrakian <ppetrakian@nvidia.com> Co-authored-by: oliver könig <okoenig@nvidia.com> Add MIMO hetero topology + distributed bootstrap (examples/mimo training-loop folder) (#5260) Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Update copy-pr-bot.yaml [skip ci] ci: Allow DCO check in merge queue and add DCO requirement to Contribution guide (#5278) Add optional group= to common_utils model/data-parallel reduction helpers (#5251) Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> chore: rotate oncall schedule Fix test_split_tensor_along_last_dim to actually assert correctness (#4710) Co-authored-by: peibli <lipeibao@126.com> Co-authored-by: Dmytro Pykhtar <37850217+dimapihtar@users.noreply.github.com> Route bridge communicator cross-grid P2P through a dedicated process group (#5234) Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Move uneven DTensor distributed fixture to conftest (#5237) Fix wgrad race condition when using double buffers. (#5222) chore: nightly sync main into dev (06_06_2026) (#5199) Co-authored-by: Nick Schank <nick@reflection.ai> Co-authored-by: Antoni-Joan Solergibert <asolergibert@nvidia.com> Co-authored-by: oliver könig <okoenig@nvidia.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Co-authored-by: janEbert <janpabloe@nvidia.com> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Co-authored-by: Shivanjan Chakravorty <schakravorty846@gmail.com> Co-authored-by: Cory Ye <44509866+cspades@users.noreply.github.com> Co-authored-by: Maanu Grover <maanug@nvidia.com> Co-authored-by: Ajay <abalasa@nvidia.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Chen Cui <chcui@nvidia.com> Co-authored-by: Tuomas Rintamaki <trintamaki@nvidia.com> Co-authored-by: Tyler Poon <tylerpoon@gmail.com> Co-authored-by: Collin McCarthy <cmccarthy@nvidia.com> Co-authored-by: Matthieu Le <matthieul@nvidia.com> Co-authored-by: Piotr Zelasko <pzelasko@nvidia.com> Co-authored-by: Ehsan Hosseini Asl <ehosseiniasl@nvidia.com> Co-authored-by: Teodor-Dumitru Ene <34819528+tdene@users.noreply.github.com> Co-authored-by: Siddharth Singh <sidsingh@nvidia.com> Co-authored-by: Dmytro Pykhtar <37850217+dimapihtar@users.noreply.github.com> Co-authored-by: Jorge Albericio <jalbericiola@nvidia.com> Co-authored-by: Deepak Narayanan <dnarayanan@nvidia.com> Co-authored-by: Xiaowei Ren <103958965+xrennvidia@users.noreply.github.com> Co-authored-by: Philip Petrakian <ppetrakian@nvidia.com> Co-authored-by: wdykas <73254672+wdykas@users.noreply.github.com> Co-authored-by: William Dykas <wdykas@oci-hsg-cs-001-vscode-03.cm.cluster> Co-authored-by: Fei Wu <33940270+YangFei1990@users.noreply.github.com> Co-authored-by: Xuanteng Huang <44627253+xuantengh@users.noreply.github.com> Co-authored-by: Yuzhong Wang <yuzhongw@nvidia.com> Co-authored-by: kunlunl <kunlunl@nvidia.com> Co-authored-by: Xuesong Ye <xuesongyey@gmail.com> Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com> Co-authored-by: Daisy Gao <daisyg@nvidia.com> Co-authored-by: Pingtian Li <158665726+Wohox@users.noreply.github.com> Co-authored-by: gautham-kollu <gkollu@nvidia.com> Co-authored-by: Yashaswi Karnati <144376261+yashaswikarnati@users.noreply.github.com> Co-authored-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: lichenlu <lichenlu_8618@163.com> Co-authored-by: peibli <lipeibao@126.com> Co-authored-by: Yan Xu <45385219+Connor-XY@users.noreply.github.com> Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> Co-authored-by: Xin Yao <xiny@nvidia.com> Co-authored-by: Gao Deng <160076886+gdengk@users.noreply.github.com> Co-authored-by: Gao Deng <gdeng@login-lyris02.lyris.clusters.nvidia.com> Co-authored-by: Zhongbo Zhu <42691305+zhongbozhu@users.noreply.github.com> Co-authored-by: Pavel Gein <pavelgejn@yandex.ru> Co-authored-by: Gao Deng <gdeng@login-lyris01.lyris.clusters.nvidia.com> Co-authored-by: Eric Harper <eharper@nvidia.com> Co-authored-by: Robin Zhang <robinz@nvidia.com> Co-authored-by: Shanmugam Ramasamy <111910568+shanmugamr1992@users.noreply.github.com> Co-authored-by: Charlie Truong <chtruong@nvidia.com> Co-authored-by: Abhishree Thittenamane <47577437+athitten@users.noreply.github.com> Co-authored-by: Abhishree Thittenamane <athittenaman@cw-dfw-cs-001-login-01.cm.cluster> Co-authored-by: root <root@pool0-01849.cm.cluster> Co-authored-by: Jingyue Wu <wujingyue@gmail.com> Co-authored-by: Jianbin Chang <shjwudp@gmail.com> Co-authored-by: Kamran Jafari <kjafarisadeg@nvidia.com> Co-authored-by: Nan Zheng <80790206+nanz-nv@users.noreply.github.com> Co-authored-by: Li Tao <lit@nvidia.com> Co-authored-by: Zhongbo Zhu <zhongboz@nvidia.com> Co-authored-by: Li Jinliang <jinliangl@nvidia.com> Co-authored-by: Robert Kirby <ArEsKay3@users.noreply.github.com> Co-authored-by: Teodor-Dumitru Ene <teodord.ene@gmail.com> Co-authored-by: Cole Hawkins <chawkins@nvidia.com> Co-authored-by: Qi Zhang <qizhang@nvidia.com> Co-authored-by: Vasudevan Rengasamy <vrengasamy@nvidia.com> Co-authored-by: tongliu <tongliu@nvidia.com> Co-authored-by: sraman-rgb <sraman@nvidia.com> Co-authored-by: Siddhartha Raman S <sraman@login-lyris02.lyris.clusters.nvidia.com> Co-authored-by: Siddhartha Raman S <sraman@login-lyris01.lyris.clusters.nvidia.com> Co-authored-by: Qiyu Wan <39144338+WanZzzzzz@users.noreply.github.com> Co-authored-by: Devil1716 <149754374+Devil1716@users.noreply.github.com> Co-authored-by: Zijie Yan <zijiey@nvidia.com> Co-authored-by: Dennis Liu <denliu@nvidia.com> Co-authored-by: Shifang Xu <shifangx@nvidia.com> Co-authored-by: Li Ding <liding@nvidia.com> Co-authored-by: sraman <sraman@users.noreply.github.com> Co-authored-by: sraman-rgb <270218152+sraman-rgb@users.noreply.github.com> Co-authored-by: Yi-Fu Wu <yifu.wu@gmail.com> Co-authored-by: Gerald Shen <geshen@nvidia.com> Co-authored-by: nvskills-svc-account <svc-nvskills-signing@nvidia.com> Co-authored-by: Hongbin Liu <lhb8125@users.noreply.github.com> Co-authored-by: Yuzhong Wang <yuzhongw@computelab-frontend-3.nvidia.com> Co-authored-by: Simiao Zhang <56124251+conver334@users.noreply.github.com> Co-authored-by: Shaurya Singh <sshaurya914@gmail.com> Co-authored-by: Kevin <23258141+kaimo455@users.noreply.github.com> Co-authored-by: mokai <mokai@baidu.com> Co-authored-by: Moozy <108285604+Moozy23232@users.noreply.github.com> Co-authored-by: Aditya Singh <60082699+adityasingh2400@users.noreply.github.com> Co-authored-by: Keshav Santhanam <ksanthanam@nvidia.com> Co-authored-by: Mike Chrzanowski <mchrzanowski@nvidia.com> Co-authored-by: Mike Chrzanowski <mchrzanowski@gcp-nrt-cs-001-login-001.cm.cluster> Co-authored-by: mathemakitten <helenn@nvidia.com> Co-authored-by: Jenny Chen <jennifchen@nvidia.com> Co-authored-by: Yu Yao <54727607+yaoyu-33@users.noreply.github.com> Co-authored-by: kajalj22 <kajalj@nvidia.com> Co-authored-by: claude[bot] <41898282+claude[bot]@users.noreply.github.com> Add named layouts to HyperCommGrid for heterogeneous parallelism (#5148) Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Fix CUDA IMA in fsdp_double_buffer when an FSDP unit's bucket doesn't fit the pool (#4810) Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Apply MIMO SP/CP sharding with explicit groups and enable THD in non-colocated path (#5150) Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Merge remote-tracking branch 'origin/dev' into main2dev/06_06_2026 # Conflicts: # megatron/training/arguments.py Fuse per-sequence AlltoAll into a unified one in GDN forward (#4913) [Dev] Add separate toggle for varlen input padding for HybridEP in THD training (#5048) Co-authored-by: Xin Yao <xiny@nvidia.com> docs: fix install guide NGC container anchor (#5224) fix: wrap mtp logging comment Add mtp_detach_heads config to detach MTP head inputs (#3456) Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> ci: add smoke tests (#5143) varlendataset for thd e2e and benchmark (#4832) Fix bug with Megatron-FSDP zero counter not working with decoupled gradients. (#4802) Merge remote-tracking branch 'origin/dev' into main2dev/06_06_2026 # Conflicts: # megatron/core/models/gpt/gpt_model.py # megatron/core/ssm/mamba_mixer.py # megatron/core/transformer/multi_token_prediction.py # tests/unit_tests/transformer/test_multi_token_prediction.py docs: Update Latest News in README.md (#3790) fix: post-CI corrections Minor improvements for Dynamic-cp (#4226) fix(combined-1f1b): release loss-node input storage after combined backward (#4909) fix(elastification): align with get_batch + utils refactors (#5194) Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Enable selective recompute for `norm_out` in GDN layers (#4715) Avoid stat syscall in rerun result validation (#5107) chore(beep boop 🤖): Bump (main) (2026-06-08) AI aided audit for Nvidia Style guidance (#5141) Co-authored-by: Philip Petrakian <pgpetrak@gmail.com> nvidia style guide audit for getting started folder (#5168) Merge current dev into nightly sync Resolve MoE op-fuser conflicts and carry over #5029 Deyu/Philip fixes. [dev] [DeepSeek-v4] Add ClampedSwiGLU to MoE mlp_op_fuser and add force balance to hash routing (#5130) fix: post-CI corrections (docs dup fields, nvrx version gate, utils re-export) - Remove merge-introduced duplicate definitions of use_transformer_engine_op_fuser…

yashaswikarnati mentioned this pull request Jun 16, 2026

Add hetero MIMO entrypoint, bootstrap, and mock data (integration) #5377

Closed

yashaswikarnati commented Jun 16, 2026

View reviewed changes

yashaswikarnati commented Jun 17, 2026

View reviewed changes

yashaswikarnati and others added 4 commits June 16, 2026 17:33

training: drop schedule-selection comment in MIMO branch

30a54c7

Remove the explanatory comment on the MIMO schedule-selection branch; the branch logic is unchanged. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: ykarnati <ykarnati@nvidia.com>

yashaswikarnati marked this pull request as ready for review June 17, 2026 01:15

yashaswikarnati requested review from a team as code owners June 17, 2026 01:15

yashaswikarnati changed the title ~~Guard MIMO cross-grid edge cases in the stock train loop~~ Support the MIMO cross-grid path in the stock training loop Jun 17, 2026

svcnvidia-nemo-ci added the complexity: low label Jun 17, 2026

Phlip79 approved these changes Jun 17, 2026

View reviewed changes

yaoyu-33 approved these changes Jun 17, 2026

View reviewed changes

yashaswikarnati changed the title ~~Support the MIMO cross-grid path in the stock training loop~~ Support the MIMO cross-grid path in training loop Jun 17, 2026

deepakn94 approved these changes Jun 22, 2026

View reviewed changes

svcnvidia-nemo-ci added the Approved All necessary approvals have been made label Jun 22, 2026

yashaswikarnati enabled auto-merge June 22, 2026 16:32

copy-pr-bot Bot temporarily deployed to public June 22, 2026 16:33 Inactive

copy-pr-bot Bot temporarily deployed to test June 22, 2026 16:33 Inactive

copy-pr-bot Bot temporarily deployed to public June 22, 2026 16:36 Inactive

copy-pr-bot Bot temporarily deployed to public June 22, 2026 16:45 Inactive

yashaswikarnati added this pull request to the merge queue Jun 22, 2026

Merged via the queue into NVIDIA:main with commit e1c4495 Jun 22, 2026
90 of 91 checks passed

yashaswikarnati deleted the p5-training-guards branch June 22, 2026 20:14

ZoRax-A5 pushed a commit to ZoRax-A5/Megatron-LM that referenced this pull request Jun 29, 2026

Support the MIMO cross-grid path in training loop (NVIDIA#5373)

24a097c

Signed-off-by: ykarnati <ykarnati@nvidia.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

terminator123 pushed a commit to 021ai/Megatron-LM that referenced this pull request Aug 3, 2026

Support the MIMO cross-grid path in training loop (NVIDIA#5373)

018648f

Signed-off-by: ykarnati <ykarnati@nvidia.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

Navigation Menu

Uh oh!

Support the MIMO cross-grid path in training loop - #5373

Support the MIMO cross-grid path in training loop#5373
yashaswikarnati merged 5 commits into
NVIDIA:mainfrom
yashaswikarnati:p5-training-guards

yashaswikarnati commented Jun 16, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Jun 16, 2026

Uh oh!

yashaswikarnati Jun 16, 2026

Uh oh!

yashaswikarnati Jun 17, 2026

Uh oh!

yashaswikarnati Jun 17, 2026

Uh oh!

yashaswikarnati commented Jun 22, 2026

Uh oh!

svcnvidia-nemo-ci commented Jun 22, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Uh oh!

Conversation

yashaswikarnati commented Jun 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What

Validation

Uh oh!

copy-pr-bot Bot commented Jun 16, 2026

Uh oh!

yashaswikarnati Jun 16, 2026

Choose a reason for hiding this comment

Uh oh!

yashaswikarnati Jun 17, 2026

Choose a reason for hiding this comment

Uh oh!

yashaswikarnati Jun 17, 2026

Choose a reason for hiding this comment

Uh oh!

yashaswikarnati commented Jun 22, 2026

Uh oh!

svcnvidia-nemo-ci commented Jun 22, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

yashaswikarnati commented Jun 16, 2026 •

edited

Loading