[yashaswikarnati]

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

What does this PR do?

Allows multimodal models to use DDP-wrapped submodules with different
parameter-gather overlap policies while participating in the standard
Megatron training, checkpoint, rerun, and paged-stash lifecycles.

Issue tracking

Linked issue: N/A — compatibility gaps found during MIMO training integration.

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code
• I have added relevant documentation
• I have run tools/autoformat.sh on my PR

Validation

• CMH job 3188397: 20 focused tests passed; 3 expected 8-GPU tests skipped (1 node / 4 allocated GPUs / 1 pytest process).
• Black 26.3.0, isort 5.13.2, Ruff 0.9.10, Pylint 3.2.6, Python compilation, and git diff --check pass on the changed scope. training.py imports are unchanged and retain pre-existing full-file isort debt.

[copy-pr-bot]

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]

/claude strict-review

[claude]

[SUGGESTION Simplification] copy(config) on a dataclass does not re-run __post_init__, so the derived config skips every validation in OptimizerConfig.__post_init__ (megatron/core/optimizer/optimizer_config.py:398-431) — including the reuse_grad_buf_for_mxfp8_param_ag + overlap_param_gather_with_optimizer_step ValueError.

For the two fields actually mutated here that's benign today (neither feeds use_precision_aware_optimizer_no_fp8_or_ds_fp8 or any other derived field), and skipping re-validation is arguably what you want since you're relaxing the flags. But it means a future field added to this function bypasses validation silently. dataclasses.replace re-runs __post_init__ and expresses the intent as "derive a variant" rather than "shallow-copy then poke fields":

def _optimizer_config_for_module(
    config: OptimizerConfig, module: torch.nn.Module
) -> OptimizerConfig:
    """Match optimizer overlap behavior to the module's DDP lifecycle."""
    ddp_config = getattr(module, 'ddp_config', None)
    if ddp_config is None:
        return config
    overlap_param_gather = ddp_config.overlap_param_gather
    return replace(
        config,
        overlap_param_gather=overlap_param_gather,
        overlap_param_gather_with_optimizer_step=(
            config.overlap_param_gather_with_optimizer_step and overlap_param_gather
        ),
    )

If __post_init__ re-validation is deliberately unwanted, keep copy but add a one-line comment saying so — otherwise the next reader will assume the derived config was validated. Note also that OptimizerConfig.overlap_param_gather's own docstring (optimizer_config.py:325-341) says it is "intended to have the same value as" the DDP-level flag; this function is the mechanism that now enforces that per module, which is worth stating in the docstring.

[claude]

[SUGGESTION Naming] Setting this to 0 is correct for RNG parity with a standalone LM run (offsets 0 and 10_000 cannot collide with the 100 * pp_rank + 10 * dp_rank terms in _set_random_seed), but it leaves two problems:

1. The constant is now a no-op addend, so role_seed_offset = _LANGUAGE_SEED_OFFSET at line 147 reads as "apply an offset" while applying none. That is fine as documentation of intent, but the comment should say the value is intentionally zero because it is the baseline, not just that encoders need offsets.
2. examples/mimo/training/data.py still defines its own _LANGUAGE_SEED_OFFSET = 20_000 (and _ENCODER_SEED_OFFSET = 10_000). Two module-level constants with the same name in the same package now mean different things, and a reader grepping _LANGUAGE_SEED_OFFSET gets contradictory answers. If the data-loader offsets are a separate namespace on purpose, rename one pair to say which seed stream it offsets (e.g. _LANGUAGE_DATA_SEED_OFFSET in data.py); if they were meant to be shared, hoist them to one module and import.

Suggested change

	# Match the standalone language-model RNG; only encoders need role offsets.
	_LANGUAGE_SEED_OFFSET = 0
	# Zero: the language role keeps the baseline seed so its RNG matches a standalone
	# language-model run bit-for-bit. Only encoders need a role offset.
	_LANGUAGE_SEED_OFFSET = 0

[claude]

[SUGGESTION] ddp_average_in_collective=True asserts a configuration MIMO can never reach, which makes these two assertions misleading rather than protective.

MIMO requires calculate_per_token_loss=True (examples/mimo/model_providers/nemotron_moe_vlm.py:133-134, enforced again in examples/mimo/training/grad_sync.py:143-146), and DistributedDataParallel.__init__ asserts not average_in_collective whenever calculate_per_token_loss is set (megatron/core/distributed/distributed_data_parallel.py:244-247). So a MIMO DDP wrap with average_in_collective=True fails at construction — the test proves the flag is forwarded, but a reader will take it as proof the combination is supported.

Either drop the flag from this test and leave a comment explaining why it can't be exercised, or turn it into an explicit negative test that the wrap rejects it. The use_distributed_optimizer / fp8_param_gather / reuse_grad_buf_for_mxfp8_param_ag assertions are the valuable ones here and should stay.

[claude]

[SUGGESTION Simplification] This helper hand-mirrors get_megatron_ddp_config's argument contract: vars(DistributedDataParallelConfig()) covers the same-named fields, and the ten explicit entries cover the renamed ones. That works today, but get_megatron_ddp_config reads bare args.X (no getattr default) at megatron/training/training.py:2617-2658, so any new bare read added there breaks these MIMO tests with an AttributeError that names a field unrelated to MIMO — a maintenance trap for whoever touches that function.

Two options, either is fine:

• Add a short comment above base saying this mirrors the bare-args contract of get_megatron_ddp_config and must be extended when that function grows a field, so the next person hitting the AttributeError knows where to look.
• Better: have get_megatron_ddp_config (or a small helper next to it) expose the arg-name mapping, and build the namespace from that, so the mapping lives in one place instead of being duplicated in a MIMO test.

Also use_precision_aware_optimizer=ddp_defaults["megatron_fsdp_use_decoupled_grad"] is confusing enough to warrant a word — it reads like a typo until you check that get_megatron_ddp_config maps megatron_fsdp_use_decoupled_grad=args.use_precision_aware_optimizer.

[deepakn94]

Add comment on why this changed?

[yashaswikarnati]

this is the seed passed to _set_random_seed(),and which derives the data/tensor/expert parallel seeds for the rng tracker. I changed this to match pretrain_hrybrid, which passes seed 0. and the actual historic reason for having differnt encoder and language model seed offset is that in mimo encoders and LLM can be colocated or non colocated with different paralleisms

[yashaswikarnati]

the outer loop uses args to check for overlap settings which are global, but we want to have encoders and LLM independently choose different overlap settings. the below checks uses the module's optimizer instance's ddp config to indepdently derive the overlap and reuse grad buffer settings

[yashaswikarnati]

/ok to test 1670ebb

[yashaswikarnati]

/ok to test de0ca4e

[yashaswikarnati]

/ok to test de0ca4e

[nemo-automation-bot]

🔄 Merge queue validation started!

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