diff --git a/skyrl/backends/skyrl_train/distributed/megatron/megatron_strategy.py b/skyrl/backends/skyrl_train/distributed/megatron/megatron_strategy.py index a030b1b486..61114eb071 100644 --- a/skyrl/backends/skyrl_train/distributed/megatron/megatron_strategy.py +++ b/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. diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py index 75261d4790..83356db7c5 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py +++ b/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): diff --git a/skyrl/train/config/config.py b/skyrl/train/config/config.py index 4806ab1adf..37f195d845 100644 --- a/skyrl/train/config/config.py +++ b/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)