[yaoyu-33]

Problem

When the public pretrain() path initializes distributed and Megatron global state, an exception from setup, training, evaluation, checkpoint finalization, or a documented callback bypasses the success-only cleanup. If the caller catches that exception, the framework-created default process group and Megatron globals remain live, so a retry can reuse stale state, report an already-initialized group, leak resources, or hang.

Caller-created process groups are not affected by this ownership bug.

Root cause

_finish_train() and _maybe_destroy_process_group() ran only after the full pretrain body completed successfully. The ordinary exception path had no corresponding cleanup, while the in-process-restart path has separate NVRx lifecycle ownership.

Fix

• Clean Megatron global state when ordinary pretrain execution fails.
• Destroy the default process group only when this training session owned its initialization.
• Skip the distributed barrier during exceptional cleanup so a rank-local failure is not converted into a barrier hang.
• Preserve the original exception and leave in-process-restart cleanup under NVRx ownership.
• Add focused coverage for framework-owned and caller-owned groups, global cleanup, barrier omission, and the injected restart wrapper.

Validation

Fail before the production fix:

uv run --active --no-sync python -m pytest tests/unit_tests/training/test_pretrain.py -q
1 failed, 5 passed
AssertionError: Expected 'destroy_process_group' to be called once. Called 0 times.

The same deterministic one-rank Gloo reproducer also failed on the unmodified base because the framework-owned real process group remained initialized, then passed with the fix.

Pass after the fix and expanded regression coverage:

uv run --no-sync python -m pytest tests/unit_tests/training/test_pretrain.py -q
7 passed

Adjacent lifecycle tests:

uv run --no-sync python -m pytest tests/unit_tests/training/test_inprocess_restart.py tests/unit_tests/training/test_initialize.py -q
47 passed

Static checks:

git diff --check
passed

uv run --active --no-sync pre-commit run --all-files
all hooks passed

Scope

This changes only exceptional cleanup for ordinary pretrain() execution. Successful cleanup behavior, caller-owned process groups, and NVRx-managed in-process restart cleanup remain unchanged. No full-suite, convergence, performance, or multi-rank training claim is made.

[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 4af184b

[claude]

LGTM

Clean, well-scoped change that adds failure cleanup to _pretrain without altering the success path. Control flow is correct across all three ownership scenarios:

• Framework-owned PG (created during a failed setup): destroyed on failure, with the barrier correctly skipped (synchronize=False) so inconsistent ranks cannot deadlock during cleanup.
• Caller-owned PG (dist already initialized before entry): preserved on failure.
• Inprocess-restart path: cleanup deferred to NVRx, PG left intact for recovery.

Minor notes (non-blocking):

• destroy_global_state() is not called on the success path, preserving prior behavior. Worth confirming this asymmetry is intended.
• _cleanup_after_pretrain_failure catches Exception (not BaseException), so a KeyboardInterrupt raised during cleanup would propagate before the PG is torn down. Likely acceptable.

Suggested test cases:

• test_framework_owned_process_group_is_destroyed_when_setup_raises
• test_caller_owned_process_group_is_preserved_when_setup_raises
• test_inprocess_restart_wrapper_retains_cleanup_ownership_when_setup_raises
• (gap) failure during train()/_finish_train() rather than setup(), confirming cleanup still runs
• (gap) assert destroy_global_state() is NOT called on the normal success path

No perf tests impacted.