[megatron] add opt-in distributed HF checkpoint export#1831
Conversation
The Megatron save path calls AutoBridge.save_hf_weights with the bridge default distributed_save=False, so a single rank writes the entire HF checkpoint serially. For large models this dominates the export and stalls every rank on each periodic save (a ~1 TB 550B export takes 40+ minutes). Add policy.megatron_config.hf_export_config with distributed_save and save_every_n_ranks, threaded through the Megatron strategy's save_hf_model into save_hf_weights. distributed_save fans shard writes across ranks (one saver per save_every_n_ranks) while producing the same standard HF sharded safetensors + index layout, so from_pretrained loads it unchanged. Defaults keep the serial behavior (distributed_save=False); strictly opt-in and Megatron-only. Signed-off-by: Vu Dinh <vudinh@outlook.com>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces parallelized Hugging Face model exporting by adding a new configuration class MegatronHFExportConfig with options for distributed_save and save_every_n_ranks. These configurations are propagated from the worker down to the Megatron strategy to allow fanning out the safetensors export across multiple ranks. Feedback suggests adding validation to ensure save_every_n_ranks is at least 1, both in the configuration's post-initialization and defensively at the start of the saving process, to prevent potential division-by-zero errors.
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.
A value < 1 raises ZeroDivisionError in the bridge's distributed save (modulo/floor-div over ranks). Fail fast at config construction. Signed-off-by: Vu Dinh <vudinh@outlook.com>
1 participant
What
Add an opt-in
policy.megatron_config.hf_export_configwithdistributed_saveandsave_every_n_ranks, threaded through the Megatron strategy'ssave_hf_modelintoAutoBridge.save_hf_weights.Why
The Megatron path currently calls
save_hf_weightswith the bridge defaultdistributed_save=False, so a single rank writes the entire HF checkpoint serially. For large models this dominates the export and stalls every rank on each periodic save — a ~1 TB 550B checkpoint takes 40+ minutes to write.distributed_save=Truefans the shard writes across ranks (one saver persave_every_n_ranks, e.g. 8 = one writer per 8-GPU node). The on-disk result is the same standard HF shardedsafetensors+model.safetensors.index.jsonlayout, sofrom_pretrainedloads it unchanged — only the write is parallelized, not the format.Behavior
distributed_save=False,save_every_n_ranks=1— no change for existing users; strictly opt-in.Notes
distributed_saveonly partitions which rank writes which standard shard file; rank 0 still writes a single consolidated index. It is not a Megatron distributed (DCP) checkpoint and needs no conversion to reload.