[megatron] add opt-in async distributed checkpoint save

[dinhxuanvu]

Summary

The Megatron torch_dist checkpoint save is fully parallel across ranks but synchronous — training blocks for the entire shard-write duration on every ckpt_interval. For large models checkpointing frequently, that stall is significant.

This adds an opt-in MegatronConfig.async_dist_ckpt_save (default False, so existing behavior is unchanged). When enabled, dist_checkpointing.save stages each rank's shards to host memory and writes them to disk in a background process, letting training resume immediately. This replaces the long-standing TODO(tgriggs): Support configurable async saves in megatron_strategy.py.

Changes

• MegatronConfig.async_dist_ckpt_save knob (off by default).
• MegatronStrategy.save_checkpoint: passes async_sharded_save through and schedules the returned request on the existing persistent AsyncCallsQueue instead of asserting it is None. Blocks on the previous async save before issuing a new one.
• MegatronStrategy.finalize_pending_saves(): drains in-flight writes; called at the start of load_checkpoint and exposed up through Worker → WorkerDispatch → end-of-training in the sync, fully-async, and SFT trainers.

Correctness

• On-disk format is identical to the synchronous save.
• The pending write is finalized before the next save, before any reload, and at end of training.
• Async falls back to synchronous for cloud destinations, where local_work_dir uploads on context exit and would otherwise race a partial checkpoint. Only active for local/shared filesystems.

Test plan

• GPU run with async_dist_ckpt_save=true: verify checkpoints are byte-identical to sync saves and reload correctly.
• Confirm training resumes before the disk write completes (timing/save_checkpoint drops).
• Verify final checkpoint at end of training is fully written before exit.

[gemini-code-assist]

Code Review

This pull request introduces support for asynchronous distributed checkpoint saving (async_dist_ckpt_save) in the Megatron strategy, allowing training to resume immediately while checkpoint writing occurs in the background. It adds mechanism to block on previous async saves, fall back to synchronous saving for cloud paths, and introduces finalize_pending_saves across trainers and workers to ensure pending writes are completed before teardown or reload. The review feedback highlights a potential issue in trainer.py and fully_async_trainer.py where finalize_pending_saves is only called if ckpt_interval > 0. It is recommended to call this method unconditionally on teardown to prevent potential checkpoint corruption if checkpoints are triggered via other means (such as callbacks) when the interval is disabled.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.