[training] fix: Synchronize MIMO loss scalers

[yaoyu-33]

Problem

Non-colocated MegatronMIMO training supports FP16 dynamic loss scaling, but each disjoint module grid owns a separate inner optimizer and scaler. If only one module grid overflows, its inner prepare_grads() updates that scaler before the outer MimoOptimizer performs world-wide found-inf consensus. The global step is skipped, but the scaler states remain different.

On the next finite step, the terminal module scales the loss with its local value while upstream modules unscale the propagated gradients with different values. This silently changes effective gradient magnitudes across the connected model. With an initial scale of 8 and hysteresis=1, the two-rank reproducer observed [4.0, 8.0] after an encoder-only overflow.

Megatron-LM PR #5331 synchronizes update success after preparation; it does not reconcile the scaler mutation that already occurred.

Fix

Bridge now wraps its MIMO optimizer so the dynamic scaler update uses the world-wide found-inf result:

1. snapshot each deduplicated dynamic scaler, including chained optimizers;
2. run the existing gradient preparation unchanged;
3. reduce found-inf across the world;
4. restore and replay one scaler update with the global result.

BF16 and constant-loss-scale paths delegate directly to MCore and add no extra collective. The wrapper reuses the same module infos and inner optimizers, preserving scheduler and checkpoint ownership.

Validation

Fail before / pass after with the unchanged regression:

uv run --no-sync python -m pytest -p no:cacheprovider tests/unit_tests/training/megatron_mimo/test_megatron_mimo_checkpointing.py::TestPretrainMegatronMIMOSetup::test_setup_megatron_mimo_initializes_checkpoint_manager -q --tb=short

• Before: failed; remote overflow left scale 8.0, expected globally backed-off 4.0.
• After: 1 passed.

Real two-rank FP16 reproducer:

uv run --no-sync python -m torch.distributed.run --standalone --nproc_per_node=2 repro_bridge_mimo_scaler.py

• Before: module scalers [4.0, 8.0].
• After: module scalers [4.0, 4.0].

Adjacent tests:

uv run --no-sync python -m pytest -p no:cacheprovider \
  tests/unit_tests/training/megatron_mimo/test_megatron_mimo_checkpointing.py \
  tests/unit_tests/training/megatron_mimo/test_megatron_mimo_config.py \
  tests/unit_tests/training/megatron_mimo/test_megatron_mimo_parallel_utils.py \
  tests/unit_tests/training/megatron_mimo/test_pretrain_megatron_mimo.py \
  -q --tb=short

• 87 passed.

uv run --no-sync pre-commit run --all-files
git diff --check

• All hooks passed; diff check passed.

Scope

This changes only Bridge-created MIMO optimizers using dynamic FP16 scaling. It does not alter BF16/constant scaling, repair already-divergent scaler values in existing checkpoints, or change direct Megatron-LM users.

[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.

[yaoyu-33]

/ok to test 1088cf5

[claude]

LGTM overall. The world-synced loss-scaler consensus is correct: it snapshots each unique scaler (deduped by id, deep-copied) before super().prepare_grads(), all-reduces found_inf with MAX over the world, then restores and replays a single update(global_found_inf). Returning the global flag keeps every rank skipping/keeping the step in lockstep. device=cuda matches the existing convention in this file.

Two minor, non-blocking observations on how the test covers upstream-API assumptions:

• prepare_grads() reads self._active_optimizers, a private attribute of upstream MimoOptimizer. The unit test assigns result.optimizer._active_optimizers itself, so it would not catch an upstream rename of that attribute; the real population path stays untested.
• The optimizer is rebuilt via _WorldSyncedMimoOptimizer(mcore_optimizer.module_infos, mcore_optimizer.config), a fresh instance from the two constructor args. If get_mimo_optimizer() ever does post-construction setup beyond init, that state is silently dropped. A short comment noting the reconstruction contract would make the intent explicit.

Neither is merge-blocking; both would surface loudly at runtime if the assumption breaks.

Suggested test cases

• No perf tests impacted.