chore: nightly sync main into dev (22_06_2026)

.github/oncall_schedule.json

@@ -1,8 +1,4 @@
 [
-    {
-        "user": "guihong-nv",
-        "date": "2026-06-10"
-    },
     {
         "user": "Phlip79",
         "date": "2026-06-17"
@@ -46,5 +42,9 @@
     {
         "user": "ilml",
         "date": "2026-08-26"
+    },
+    {
+        "user": "janEbert",
+        "date": "2026-09-02"
     }
 ]

AGENTS.md

@@ -24,6 +24,7 @@ skill keyword — infer it from the artifact you read.
 
 - All PRs must be created as **drafts**. Use `gh pr create --draft` or the GitHub UI draft option.
 - Never push branches directly to `https://github.com/NVIDIA/Megatron-LM`. You must push your branch to a personal fork (e.g. `https://github.com/<your-username>/Megatron-LM`), then open a PR from the fork's branch against `NVIDIA/Megatron-LM`.
+- Commit PR changes with both `-s` and `-S`: `-s` adds the required `Signed-off-by` trailer, and `-S` signs the commit so copy-pr-bot and `/ok to test` can verify the pushed commit without manually specifying the SHA.
 - Read @docs/developer/contribute.md for the full contribution policy, including code style, commit message conventions, and issue guidelines.
 
 ### Code Quality

docker/lts/requirements.txt

@@ -21,3 +21,4 @@ fastapi==0.136.3                       # was: fastapi~=0.50 (forces compat with
 datasets==4.8.5
 emerging_optimizers @ git+https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git@v0.2.0
 nvidia-resiliency-ext==0.6.0
+zstandard==0.25.0

docs/user-guide/features/fine_grained_activation_offloading.md

@@ -13,7 +13,7 @@ Contributed in collaboration with RedNote.
 
 Memory is often the limiting factor for very large sparse MoE models such as DeepSeek-V3 and Qwen3-235B. Fine-grained recomputation lowers activation memory at the cost of extra compute. Offloading can use host-device bandwidth so that reload overlaps compute and keeps overhead small in many setups. Fine-grained activation offloading moves activations at module granularity so you can tune how much activation memory leaves the device and adjust training throughput.
 
-Supported offloading modules are `"attn_norm"`, `"core_attn"`, `"attn_proj"`, `"mlp_norm"`, `"expert_fc1"`, and `"moe_act"`. They can be combined with fine-grained recomputation to free almost all activations for a transformer layer on the device.
+Supported offloading modules are `"attn_norm"`, `"qkv_linear"`, `"core_attn"`, `"attn_proj"`, `"mlp_norm"`, `"expert_fc1"`, `"moe_act"`, and `"fused_group_mlp"`. They can be combined with fine-grained recomputation to free almost all activations for a transformer layer on the device. `fused_group_mlp` requires `--use-transformer-engine-op-fuser` and offloads the whole fused grouped MLP, so it cannot be combined with `expert_fc1` or `moe_act`.
 
 ## Features
 
@@ -33,7 +33,7 @@ Supported offloading modules are `"attn_norm"`, `"core_attn"`, `"attn_proj"`, `"
 --fine-grained-activation-offloading
 
 # Modules whose inputs are offloaded (refer to your training script for list or delimiter syntax).
-# Choices: "attn_norm", "core_attn", "attn_proj", "mlp_norm", "expert_fc1", "moe_act".
+# Choices: "attn_norm", "qkv_linear", "core_attn", "attn_proj", "mlp_norm", "expert_fc1", "moe_act", "fused_group_mlp".
 --offload-modules expert_fc1
 ```
 

docs/user-guide/features/paged_stash.md

@@ -21,7 +21,7 @@ Whenever `moe_expert_rank_capacity_factor` is set, a **runner** wraps forward-ba
 
 ## Prerequisites
 
-HybridEP + TE fused grouped experts are required whenever `moe_expert_rank_capacity_factor` is set. With `moe_paged_stash` enabled: capacity factor must be set; no `cpu_offloading`; `offload_modules` must not include `expert_fc1` or `moe_act`. The runner is active whenever capacity factor is set (even without `--moe-paged-stash`) for over-budget reruns; stash overflow is checked only when paged stashing is on.
+HybridEP + TE fused grouped experts are required whenever `moe_expert_rank_capacity_factor` is set. With `moe_paged_stash` enabled: capacity factor must be set; no `cpu_offloading`; `offload_modules` must not include `expert_fc1`, `moe_act`, or `fused_group_mlp`. The runner is active whenever capacity factor is set (even without `--moe-paged-stash`) for over-budget reruns; stash overflow is checked only when paged stashing is on.
 
 ## Configuration
 

examples/mimo/training/runtime.py

@@ -0,0 +1,154 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
+"""Per-rank runtime setup (RNG seeding, freezing, DDP wrapping) for hetero MIMO training."""
+
+from __future__ import annotations
+
+import argparse
+
+import torch
+
+from examples.mimo.training.topology import HeteroTopology
+from megatron.core.distributed import DistributedDataParallel, DistributedDataParallelConfig
+from megatron.core.models.mimo.config.role import MIMO_LANGUAGE_MODULE_KEY
+from megatron.core.models.mimo.model.base import MimoModel
+from megatron.core.process_groups_config import ProcessGroupCollection
+from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
+from megatron.core.transformer.module import Float16Module
+from megatron.core.utils import get_pg_rank, get_pg_size
+from megatron.training.training import resolve_ddp_bucket_size, wrap_model_chunks_with_ddp
+from megatron.training.utils import print_rank_0
+
+
+class _EncoderFloat16Module(Float16Module):
+    """Float16Module that keeps encoder outputs in model precision for the bridge."""
+
+    def forward(self, *inputs, fp32_output=False, **kwargs):  # noqa: D102
+        return super().forward(*inputs, fp32_output=fp32_output, **kwargs)
+
+
+def configure_module_rng(
+    args: argparse.Namespace, pg_collection: ProcessGroupCollection, role_seed_offset: int
+) -> None:
+    """Seed the CUDA RNG tracker for one module role from its tp/pp coordinates plus the offset.
+
+    The seed is shared across a module's DP/CP replicas but distinct across PP stages and roles,
+    so disjoint modules (and stages) get independent RNG state. Caller invokes once per active
+    module on this rank.
+    """
+    for _required in ("pp", "tp", "ep", "expt_tp"):
+        assert (
+            getattr(pg_collection, _required, None) is not None
+        ), f"pg_collection passed to configure_module_rng must define {_required}"
+    pp_rank = get_pg_rank(pg_collection.pp)
+    tp_rank = get_pg_rank(pg_collection.tp)
+    ep_rank = get_pg_rank(pg_collection.ep)
+    expt_tp_rank = get_pg_rank(pg_collection.expt_tp)
+    seed = args.seed + role_seed_offset + (100 * pp_rank)
+    torch.manual_seed(seed)
+    model_parallel_cuda_manual_seed(
+        seed, tp_rank=tp_rank, ep_rank=ep_rank, etp_rank=expt_tp_rank, force_reset_rng=True
+    )
+
+
+def _freeze_modality_submodule(submodule: torch.nn.Module, args: argparse.Namespace) -> None:
+    """Freeze the encoder backbone (--freeze-vit) and/or projector (--freeze-projection)."""
+    if getattr(args, "freeze_vit", False):
+        submodule.encoders.requires_grad_(False)
+    if getattr(args, "freeze_projection", False):
+        submodule.input_projections.requires_grad_(False)
+        submodule.output_projections.requires_grad_(False)
+
+
+def _module_config(module: torch.nn.Module):
+    """Return the module's own config, else the first descendant config (e.g. an encoder)."""
+    config = getattr(module, "config", None)
+    if config is not None:
+        return config
+    for child in module.modules():
+        config = getattr(child, "config", None)
+        if config is not None:
+            return config
+    raise ValueError("Cannot resolve a config for DDP wrapping from module")
+
+
+def _maybe_float16_wrap(module: torch.nn.Module, config, is_encoder: bool) -> torch.nn.Module:
+    """Wrap a submodule in Float16Module when fp16/bf16 is enabled; encoders keep bf16 outputs."""
+    if not (getattr(config, "fp16", False) or getattr(config, "bf16", False)):
+        return module
+    cls = _EncoderFloat16Module if is_encoder else Float16Module
+    return cls(config, module)
+
+
+def wrap_active_modules_with_ddp(
+    args: argparse.Namespace, mimo_model: MimoModel, topology: HeteroTopology
+) -> None:
+    """Freeze (per --freeze-* flags), Float16Module-wrap, and DDP-wrap each active module."""
+    pad_buckets = getattr(args, "ddp_pad_buckets_for_high_nccl_busbw", False)
+    grad_reduce_in_fp32 = getattr(args, "accumulate_allreduce_grads_in_fp32", True)
+
+    ddp_stream = torch.cuda.Stream()
+    ddp_stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(ddp_stream):
+        if mimo_model.language_model is not None:
+            if getattr(args, "freeze_lm", False):
+                mimo_model.language_model.requires_grad_(False)
+            overlap = getattr(args, "overlap_grad_reduce", False)
+            ddp_config = DistributedDataParallelConfig(
+                overlap_grad_reduce=overlap,
+                overlap_param_gather=getattr(args, "overlap_param_gather", False),
+                num_buckets=getattr(args, "ddp_num_buckets", None),
+                bucket_size=getattr(args, "ddp_bucket_size", None),
+                pad_buckets_for_high_nccl_busbw=pad_buckets,
+                use_distributed_optimizer=True,
+                grad_reduce_in_fp32=grad_reduce_in_fp32,
+            )
+            # Resolve the absolute bucket size on the real config, as get_model does.
+            ddp_config.bucket_size = resolve_ddp_bucket_size(
+                ddp_config,
+                topology.module_pgs[MIMO_LANGUAGE_MODULE_KEY].dp_cp,
+                overlap,
+                sum(p.numel() for p in mimo_model.language_model.parameters()),
+            )
+            lm_config = _module_config(mimo_model.language_model)
+            lm_module = _maybe_float16_wrap(mimo_model.language_model, lm_config, is_encoder=False)
+            print_rank_0("wrapping language model in DDP")
+            mimo_model.language_model = wrap_model_chunks_with_ddp(
+                [lm_module],
+                lm_config,
+                ddp_config,
+                DP=DistributedDataParallel,
+                pg_collection=topology.module_pgs[MIMO_LANGUAGE_MODULE_KEY],
+            )[0]
+
+        for name, submodule in mimo_model.modality_submodules.items():
+            if submodule is None or name not in topology.module_pgs:
+                continue
+            _freeze_modality_submodule(submodule, args)
+            ddp_config = DistributedDataParallelConfig(
+                overlap_grad_reduce=False,
+                overlap_param_gather=False,
+                num_buckets=getattr(args, "ddp_num_buckets", None),
+                bucket_size=getattr(args, "ddp_bucket_size", None),
+                pad_buckets_for_high_nccl_busbw=pad_buckets,
+                use_distributed_optimizer=True,
+                grad_reduce_in_fp32=grad_reduce_in_fp32,
+            )
+            # Encoders keep overlap off; resolve_ddp_bucket_size returns None there.
+            ddp_config.bucket_size = resolve_ddp_bucket_size(
+                ddp_config,
+                topology.module_pgs[name].dp_cp,
+                False,
+                sum(p.numel() for p in submodule.parameters()),
+            )
+            enc_config = _module_config(submodule)
+            enc_module = _maybe_float16_wrap(submodule, enc_config, is_encoder=True)
+            print_rank_0(f"wrapping modality submodule {name!r} in DDP")
+            mimo_model.modality_submodules[name] = wrap_model_chunks_with_ddp(
+                [enc_module],
+                enc_config,
+                ddp_config,
+                DP=DistributedDataParallel,
+                pg_collection=topology.module_pgs[name],
+            )[0]
+    torch.cuda.current_stream().wait_stream(ddp_stream)

megatron/core/distributed/finalize_model_grads.py

@@ -353,6 +353,7 @@ def _update_router_expert_bias(
                 hasattr(module, 'expert_bias')
                 and module.training
                 and module.expert_bias is not None
+                and not getattr(module, 'frozen_expert_bias', False)
             ):
                 tokens_per_expert_list.append(module.local_tokens_per_expert)
                 expert_bias_list.append(module.expert_bias)

.github/workflows/sync-skills.yml → megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/__init__.py

@@ -11,19 +11,10 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-name: Sync skills → agent dirs
 
-on:
-  workflow_dispatch:
-  push:
-    branches:
-      - main
-    paths:
-      - "skills/**"
-      - "AGENTS.md"
+"""Experimental Megatron-FSDP implementation."""
 
-jobs:
-  sync:
-    uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_sync_skills.yml@v0.91.0
-    secrets:
-      PAT: ${{ secrets.PAT }}
+from .dbuffer import DBuffer
+from .placement import Flat, Partial, Placement, Replicate
+
+__all__ = ["DBuffer", "Flat", "Partial", "Placement", "Replicate"]