fix(training): align Muon distributed setup and checkpointing

[cuichenx]

What does this PR do?

Fix distributed Muon setup, parameter synchronization, and checkpoint save/resume. Changing an Adam recipe to optimizer="muon" can currently retain incompatible distributed settings and silently produce incorrect updates. Explicit layer-wise configurations also miss overlap and checkpoint paths that check only use_distributed_optimizer.

Changelog

• Normalize emerging optimizers to layer-wise distribution when distributed training is requested, including legacy dist_* aliases, and pass the required parameter layout through provider, builder, and component model setup.
• Recognize both distributed optimizer implementations when configuring overlap, synchronizing parameters for evaluation/hash checks/checkpoints, and selecting checkpoint metadata.
• Support parameter synchronization through delegating model wrappers. Disable encoder overlap when empty-input ranks can skip the hooks needed to issue collectives.
• Reject unsupported FSDP, optimizer-step overlap, multiple optimizer instances, and component gradient-buffer reuse combinations early.
• Document distributed Muon configuration and add focused unit coverage plus a two-GPU L1 training/checkpoint regression.

Validation

• Targeted unit tests: 1,008 passed, 7 skipped.
• Two-GPU BF16 GPT, 531,968 parameters: 14 fresh-process cases covering explicit/inherited/legacy Muon settings, provider and builder paths, parameter overlap on/off, replica hashes, two checkpoint sharding formats, and Adam controls. Six updates per uninterrupted run, evaluation at steps 3/6, saves at 2/4/6, and fresh resumes from step 2. Losses, gradient norms, and final tensors match exactly across the intended pairs.
• Four-GPU BF16 GPT + CLIP, 29,228,544 parameters: six Adam/Muon overlap and resume cases with alternating empty encoder ranks. Losses, gradient norms, completed gradient reductions, synchronized checkpoint hashes, and scheduler records match exactly. Resume checks also compare restored model and optimizer tensors before continuing.
• Negative controls: unchanged main silently produces different updates for inherited Muon settings and fails the replica-hash case. Restoring only the old save-time gate causes the checkpoint synchronization assertion to fail at the first save.
• All-file pre-commit checks pass. Strict mypy is not clean: the checked training files and installed dependencies report existing typing errors and missing typed-package markers.

These are short correctness runs; they do not establish long-run convergence or performance.

GitHub Actions CI

The added L1_Launch_training_muon.sh runs the two-GPU matrix with exact comparisons. CI results are pending.

Before your PR is ready for review

• Contributor guidelines followed.
• Necessary regression tests added.
• Documentation updated.
• Independent review completed.

[github-actions]

Automatic Claude reviews have been retired. To request a pull-request review, post a comment containing:

/review

Add model=claude to use a Claude reviewer (the default is model=codex). mode=light|strict selects the review depth; for example, /review model=claude mode=strict. Comment /review help for all options.

[github-actions]

🌿 Preview your docs: https://nvidia-preview-pr-6443-ee73562a10e6.docs.buildwithfern.com/nemo/megatron-bridge

[cuichenx]

/review

[svcnvidia-nemo-ci]

One blocking validation-order issue: runtime setup rejects a Muon configuration before its explicit communication-overlap override can make it valid.

Reviewed all 26 changed files and traced the affected setup, synchronization, and checkpoint paths. This was a static review; tests were not executed.

Suggested test cases

No perf tests impacted.

• test_runtime_config_update_applies_overlap_override_before_layerwise_validation — proposed regression described inline.
• test_distributed_emerging_optimizer_normalization — retain normalization and idempotence assertions.
• test_reject_layerwise_optimizer_step_overlap — retain rejection when optimizer-step overlap remains enabled.
• L1_Launch_training_muon.sh — run the numerical and fresh-process resume matrix.

Note: formal approval is not available for this request, so this review is posted as a comment only.

Review generated by codex

[svcnvidia-nemo-ci]

[P2] Apply communication overrides before validating optimizer-step overlap

Failure: Start with vanilla_gpt_pretrain_config() and set:

cfg.optimizer.optimizer = "muon"
cfg.optimizer.overlap_param_gather_with_optimizer_step = True
cfg.comm_overlap = CommOverlapConfig(
    tp_comm_overlap=False,
    overlap_param_gather_with_optimizer_step=False,
)
runtime_config_update(cfg)

The inherited distributed flag makes this call select layer-wise distribution, then normalize_optimizer_config raises immediately. However, CommOverlapConfig applies explicit user overrides, so the effective optimizer-step-overlap setting should be False and this configuration should proceed.

Correction: Remove this pre-overlap synchronization call; retain the call at the start of ConfigContainer.validate, which runs after overlap setup. The updated CommOverlapConfig.setup gate already recognizes optimizer-side flags and legacy aliases without this early normalization. Suggested patch at this location, for author verification:

-    # Resolve optimizer mode before selecting communication overlap capabilities.
-    _validate_and_sync_distributed_optimizer_settings(cfg)
-

The invariant is to validate effective settings after communication overrides, while still normalizing distribution before model construction.

Preservation and acceptance — proposed checks, not executed: Add test_runtime_config_update_applies_overlap_override_before_layerwise_validation using the setup above, parameterized with muon and dist_muon. Expect successful runtime update, normalized name muon, layer-wise distribution enabled, ordinary optimizer distribution disabled, DDP distribution enabled, and optimizer-step overlap disabled. With no override, or an explicit True override, expect the existing ValueError. Retain the existing normalization test's idempotent explicit/inherited/alias results, replicated Muon's all-distribution-flags-false result, and Adam/SGD's ordinary-distribution result in the distribution tests.

[kamran-nvidia]

This also moves the existing dist_muon recipes (Kimi K2 and its perf recipes, Kimi K2.5-VL) off the legacy LayerWise path, even though they set ddp.use_distributed_optimizer=False, and they now pick up overlap_param_gather from CommOverlapConfig too. Probably the right direction, but it's not in the description or validated on MoE/PP. It also looks like pre-PR Muon checkpoints won't restore optimizer state, and there's no Bridge opt-out like --no-use-layer-wise-param-layout. Could you do a quick before/after on one Kimi config and update the stale recipe comments and docs?