[ckpt] fix: Preserve async checkpoint logger iteration

[yaoyu-33]

Problem

When a supported async distributed checkpoint remains in flight while training advances, its completion callbacks receive the fixed checkpoint path for iteration N but read the later mutable training step N+k.

This silently creates conflicting checkpoint provenance:

• W&B records iteration N+k for the iter_N artifact.
• MLflow uploads iter_N contents beneath an iter_N+k artifact namespace.
• Comet records last_saved_iteration N+k for the iter_N checkpoint.

Root cause

The async logger finalizers close over TrainState and read train_state.step only when the queued save completes. The checkpoint path and an immutable checkpoint_step snapshot are created when the save is scheduled.

Fix

Pass the existing checkpoint_step snapshot to the W&B, MLflow, and Comet completion helpers. This is limited to checkpoint logger provenance; callback state and other checkpoint lifecycle behavior are unchanged.

Regression evidence

The focused regression schedules iter_0001000, advances TrainState.step to 1001, then executes the captured async finalizers.

Before the fix:

Expected: on_save_checkpoint_success('/checkpoints/iter_0001000', '/checkpoints', 1000, ...)
Actual:   on_save_checkpoint_success('/checkpoints/iter_0001000', '/checkpoints', 1001, ...)
1 failed

After the fix:

uv run python -m pytest tests/unit_tests/training/test_checkpointing.py::TestSaveCheckpoint::test_async_checkpoint_loggers_use_scheduled_step -q

Result: 1 passed.

Focused and adjacent validation:

uv run python -m pytest \
  tests/unit_tests/training/test_checkpointing.py::TestSaveCheckpoint::test_async_checkpoint_loggers_use_scheduled_step \
  tests/unit_tests/training/test_checkpointing.py::TestSaveCheckpoint::test_async_retention_keeps_tracker_checkpoint_until_finalize \
  tests/unit_tests/training/test_checkpointing.py::TestSaveCheckpoint::test_async_global_non_persistent_overlap_preserves_tracker_checkpoint \
  tests/unit_tests/training/test_checkpointing.py::TestSaveCheckpoint::test_save_checkpoint_async_hf_schedules_megatron_before_hf_export \
  tests/unit_tests/training/utils/test_mlflow_utils.py \
  tests/unit_tests/training/utils/test_comet_utils.py -q

Result: 34 passed.

Additional checks:

• git diff --check: passed
• uv run pre-commit run --all-files: passed

The regression ran on CPU against the exact audited Bridge revision and its pinned Megatron-Core source. No GPU or distributed runtime behavior was mocked away; the test controls only the async completion timing and external logger calls.

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

[claude]

LGTM - clean, well-targeted bug fix.

The change routes the async W&B / MLflow / Comet completion finalizers through the immutable checkpoint_step snapshot (captured at checkpoint_step = train_state.step, checkpointing.py:1149) instead of re-reading the mutable train_state.step when the queued save later completes. This matches the sibling async cleanup finalizers (max_iteration=checkpoint_step at lines 1603/1622/1644), which already used the snapshot, so the three logger finalizers were the last outliers. Correctly fixes the provenance mismatch where an in-flight iter_N artifact was tagged with a later iteration N+k.

Verified:

• checkpoint_step is assigned exactly once and never reassigned, safe to close over.
• checkpoint_name / save_dir were already immutable snapshots, so the finalizers now consistently reference iteration N.
• The regression test advances train_state.step to 1001 after scheduling iter_0001000, executes the captured finalizers, and asserts each logger receives step 1000, exercising exactly the closure-timing bug (would fail on pre-fix code).

Minor test-robustness note (non-blocking): test_async_checkpoint_loggers_use_scheduled_step asserts on the hardcoded path /checkpoints/iter_0001000, coupling the test to the fixture save dir and the iter naming convention. Deriving it from the fixture + get_checkpoint_name would be more resilient, but this is optional.

Suggested test cases:

• No perf tests impacted.
• tests/unit_tests/training/test_checkpointing.py::TestSaveCheckpoint::test_async_checkpoint_loggers_use_scheduled_step
• tests/unit_tests/training/test_checkpointing.py::TestSaveCheckpoint::test_async_retention_keeps_tracker_checkpoint_until_finalize
• tests/unit_tests/training/test_checkpointing.py::TestSaveCheckpoint::test_async_global_non_persistent_overlap_preserves_tracker_checkpoint
• tests/unit_tests/training/utils/test_mlflow_utils.py
• tests/unit_tests/training/utils/test_comet_utils.py