[yaoyu-33]

What does this PR do?

Enables the existing fault-tolerance lifecycle for MegatronMIMO training.

When a supported ConfigContainer enables ft.enable_ft_package, the MIMO
entry point accepted the configuration but never initialized Bridge's
rank-monitor client. Its training loop also skipped the per-step start/end
callbacks. The shared checkpoint, evaluation, shutdown, and abort hooks then
silently remained no-ops because no client existed.

This change initializes fault tolerance before MIMO model, optimizer,
checkpoint, and data setup, then brackets each existing training step with the
same callbacks used by the standard training loop.

Why is this fix minimal?

• Reuses the existing megatron.bridge.training.fault_tolerance integration.
• Adds no configuration, public API, dependency, CI, fallback, or MCore change.
• Keeps behavior unchanged when fault tolerance is disabled.
• Covers only MegatronMIMO setup and step lifecycle wiring.

Validation

The unchanged focused regression failed on main because neither FT setup nor
the step-start callback was invoked:

uv run --active --no-sync python -m pytest --confcutdir=tests/unit_tests/training/megatron_mimo \
  tests/unit_tests/training/megatron_mimo/test_megatron_mimo_checkpointing.py::TestSetupMegatronMIMOCheckpointLoading::test_enabled_fault_tolerance_initializes_monitoring \
  tests/unit_tests/training/megatron_mimo/test_megatron_mimo_checkpointing.py::TestTrainMegatronMIMOCheckpointIntegration::test_calls_checkpoint_and_fault_tolerance_hooks_with_pg_collection \
  -q --tb=short

2 failed

It passed unchanged after the fix. Adjacent verification:

uv run --active --no-sync python -m pytest --confcutdir=tests/unit_tests/training/megatron_mimo \
  tests/unit_tests/training/megatron_mimo/test_megatron_mimo_checkpointing.py \
  -q --tb=short
# 43 passed

uv run --active --no-sync python -m pytest --confcutdir=tests/unit_tests/training/megatron_mimo \
  tests/unit_tests/training/megatron_mimo/test_pretrain_megatron_mimo.py \
  -q --tb=short
# 8 passed

git diff --check
# passed

uv run --no-sync pre-commit run --all-files
# all configured hooks passed

Scope limitations

This verifies lifecycle wiring and ordering for the supported MegatronMIMO
entry point. It does not claim fault-injection recovery, convergence,
performance, or full-suite validation.

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

[claude]

Light review - fault-tolerance hooks for the MegatronMIMO training path

The change wires the existing fault_tolerance API into the MIMO setup and training loop, matching the standard train.py/setup.py integration: setup() + maybe_setup_simulated_fault() in setup, and on_training_step_start/end around the train step. Correctness looks good: the hooks are no-ops when cfg.ft is None, the guard (cfg.ft and cfg.ft.enable_ft_package) mirrors the standard path, and setup() runs before build_megatron_mimo_model (verified by the new ordering assertion). Eval-step FT hooks are already reached via evaluate() in eval.py, and shutdown/abort is handled through _finish_train and _cleanup_after_pretrain_failure, so the FT lifecycle is complete for MIMO.

Observations:

• Unwrapped top-of-loop async finalize (minor): the standard loop wraps the non-blocking finalize_async_saves at the top of the loop in on_checkpointing_start/on_checkpointing_end (train.py:393-395). The MIMO loop (train_megatron_mimo.py:311-314) does not, so async-save finalization time falls into the FT out-of-section window instead of the checkpointing section. With --calc-ft-timeouts this can skew timeout computation. See inline comment. Not merge-blocking.
• Test coverage is good: the new tests cover setup ordering, the enabled path, and assert step-start/step-end hooks fire once each with state. The added cfg.ft = None stubs keep the existing tests green against the new attribute access.

Suggested test cases

• No perf tests impacted.
• test_enabled_fault_tolerance_initializes_monitoring (added - verifies FT setup runs before model build)
• test_calls_checkpoint_and_fault_tolerance_hooks_with_pg_collection (added - verifies step-start/end hooks fire once)
• Consider adding a case asserting the FT hooks are NOT called when cfg.ft.enable_ft_package is False, to lock in the no-op guard for the disabled path.

[yaoyu-33]

/ok to test 330d12c

[yaoyu-33]

/ok to test f3706be

[yaoyu-33]

/ok to test 4cd8f86

[yaoyu-33]

/ok to test 2df3520

[yaoyu-33]

/ok to test 07f17ed

[yaoyu-33]

/ok to test d967754

[yaoyu-33]

/ok to test ad70e79

[yaoyu-33]

/ok to test e679ced

[yaoyu-33]

/ok to test 0967db5

[yaoyu-33]

/ok to test 53e6edb