[megatron] add opt-in distributed HF checkpoint export

skyrl/backends/skyrl_train/distributed/megatron/megatron_strategy.py

@@ -518,7 +518,17 @@ def _load_lora_adapters(self, model, ckpt_dir):
                 raise ValueError(f"Unexpected keys in LoRA adapter state dict: {unexpected}")
             self.print(f"Loaded {len(state_dict['model_state_dict'])} LoRA adapters from {adapter_path}.")
 
-    def save_hf_model(self, bridge, model: MegatronModelWrapper, output_dir: str, tokenizer=None, **kwargs) -> None:
+    def save_hf_model(
+        self,
+        bridge,
+        model: MegatronModelWrapper,
+        output_dir: str,
+        tokenizer=None,
+        *,
+        distributed_save: bool = False,
+        save_every_n_ranks: int = 1,
+        **kwargs,
+    ) -> None:
         # Create checkpoint directory if it doesn't exist.
         if self.is_rank_0():
             io.makedirs(output_dir, exist_ok=True)
@@ -530,7 +540,17 @@ def save_hf_model(self, bridge, model: MegatronModelWrapper, output_dir: str, to
             # on a Qwen3.5 VL checkpoint, whose shards co-mingle vision and text
             # weights): the bridge writes a shard only once all its keys are yielded,
             # so strict=True silently writes zero weights. No-op for complete exports.
-            bridge.save_hf_weights(model.actor_module, work_dir, strict=False)
+            #
+            # distributed_save fans the shard writes across ranks (one saver per
+            # save_every_n_ranks) instead of serializing them on rank 0 -- the same
+            # standard HF sharded layout, just written in parallel.
+            bridge.save_hf_weights(
+                model.actor_module,
+                work_dir,
+                strict=False,
+                distributed_save=distributed_save,
+                save_every_n_ranks=save_every_n_ranks,
+            )
             self.print(f"Successfully saved HF safetensors model to {output_dir}")
 
             # Only rank 0 saves the Huggingface config and tokenizer.

skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py

@@ -427,6 +427,7 @@ def init_configs(
 
         self.provider = provider
         self.bridge = bridge
+        self.megatron_config = megatron_config
 
         # strategy.hf_config is the on-disk source-of-truth used by
         # save_hf_configs and must NOT carry runtime overrides like
@@ -685,11 +686,14 @@ def _pad_microbatch_to_size(self, micro_dict: dict, target_batch_size: int) -> d
 
     def save_hf_model(self, export_dir: str, tokenizer):
         # Save model in HuggingFace safetensors format
+        hf_export = self.megatron_config.hf_export_config
         self.strategy.save_hf_model(
             self.bridge,
             self.model,
             export_dir,
             tokenizer=tokenizer,
+            distributed_save=hf_export.distributed_save,
+            save_every_n_ranks=hf_export.save_every_n_ranks,
         )
 
     def _get_module_for_offload(self):

skyrl/train/config/config.py

@@ -164,6 +164,24 @@ class MegatronTorchProfilerConfig(BaseConfig):
     save_path: Optional[str] = None
 
 
+@dataclass
+class MegatronHFExportConfig(BaseConfig):
+    distributed_save: bool = False
+    """Fan the Megatron->HF safetensors export across ranks instead of writing the
+    whole checkpoint from rank 0. The on-disk result is the standard HF sharded
+    format either way; this only parallelizes the write, which is decisive for
+    multi-hundred-GB checkpoints whose serial rank-0 write stalls every rank."""
+    save_every_n_ranks: int = 1
+    """In distributed save, only ranks 0, N, 2N, ... write shards (e.g. 8 = one
+    writer per 8-GPU node). Ignored when ``distributed_save`` is False."""
+
+    def __post_init__(self) -> None:
+        # save_every_n_ranks indexes ranks via modulo/floor-div in the bridge's
+        # distributed save; < 1 raises ZeroDivisionError there. Fail fast instead.
+        if self.save_every_n_ranks < 1:
+            raise ValueError(f"save_every_n_ranks must be >= 1, got {self.save_every_n_ranks}")
+
+
 @dataclass
 class MegatronLoraConfig(BaseConfig):
     lora_type: str = "lora"
@@ -209,6 +227,7 @@ class MegatronConfig(BaseConfig):
     """Pass through to Megatron-Bridge - can be set to 'fp64' for additional numerical stability."""
     ddp_config: MegatronDDPConfig = field(default_factory=MegatronDDPConfig)
     torch_profiler_config: MegatronTorchProfilerConfig = field(default_factory=MegatronTorchProfilerConfig)
+    hf_export_config: MegatronHFExportConfig = field(default_factory=MegatronHFExportConfig)
     lora_config: MegatronLoraConfig = field(default_factory=MegatronLoraConfig)
     optimizer_config_kwargs: Dict[str, Any] = field(
         default_factory=lambda: copy.deepcopy(DEFAULT_MEGATRON_OPTIMIZER_KWARGS)