From ff52ccdf2322a834cddd234e85d0158852702556 Mon Sep 17 00:00:00 2001 From: Alexander Kovrigin Date: Thu, 2 Jul 2026 17:30:53 +0200 Subject: [PATCH] feat: per-step inline HF export for Megatron checkpoints Enable megatron-bridge's inline HF export at every checkpoint save (not just the last). _create_checkpoint_config hardcoded CheckpointConfig with no passthrough, so add a minimal opt-in: MegatronConfig gains an optional "checkpoint" dict that _create_checkpoint_config applies onto CheckpointConfig, defaulting hf_source_path to the policy model name. Set policy.megatron_cfg.checkpoint = {also_save_hf_checkpoint: true, hf_distributed_save: true, ckpt_format: torch_dist, non_persistent_ckpt_type: global} to produce iter_*/hf/ on every save. Signed-off-by: Alexander Kovrigin --- nemo_rl/models/megatron/setup.py | 31 +++++++++++++++++++++++++++++-- nemo_rl/models/policy/__init__.py | 7 +++++++ 2 files changed, 36 insertions(+), 2 deletions(-) diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index 626003467b..b40734d382 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -541,6 +541,8 @@ def setup_model_config( weights_path, optimizer_path, load_main_params_from_ckpt, + overrides=config["megatron_cfg"].get("checkpoint"), + hf_model_name=hf_model_name, ) # Validate training configuration @@ -853,9 +855,19 @@ def _create_checkpoint_config( weights_path: Optional[str], optimizer_path: Optional[str], load_main_params_from_ckpt: bool = False, + overrides: Optional[dict[str, Any]] = None, + hf_model_name: Optional[str] = None, ) -> CheckpointConfig: - """Create checkpoint configurations.""" - return CheckpointConfig( + """Create checkpoint configurations. + + `overrides` (from `policy.megatron_cfg.checkpoint`) is applied field-by-field + onto the CheckpointConfig, so callers can opt into megatron-bridge checkpoint + features NeMo-RL doesn't surface explicitly (e.g. inline HF export via + `also_save_hf_checkpoint`). When HF export is requested without an explicit + `hf_source_path`, it defaults to `hf_model_name` (the config/tokenizer + template megatron-bridge needs to build the AutoBridge). + """ + ckpt = CheckpointConfig( save_interval=100, save=weights_path, load=weights_path, @@ -867,6 +879,21 @@ def _create_checkpoint_config( load_rng=False, load_main_params_from_ckpt=load_main_params_from_ckpt, ) + for key, value in (overrides or {}).items(): + if hasattr(ckpt, key): + setattr(ckpt, key, value) + else: + warnings.warn( + f"Ignoring unknown megatron_cfg.checkpoint override '{key}'.", + stacklevel=2, + ) + if ( + getattr(ckpt, "also_save_hf_checkpoint", False) + and not getattr(ckpt, "hf_source_path", None) + and hf_model_name + ): + ckpt.hf_source_path = hf_model_name + return ckpt def _validate_training_config(config: PolicyConfig, model_cfg: Any) -> None: diff --git a/nemo_rl/models/policy/__init__.py b/nemo_rl/models/policy/__init__.py index 16f0f2cc5f..6a0248c080 100644 --- a/nemo_rl/models/policy/__init__.py +++ b/nemo_rl/models/policy/__init__.py @@ -240,6 +240,13 @@ class MegatronConfigDisabled(TypedDict): class MegatronConfig(TypedDict): enabled: Literal[True] env_vars: NotRequired[dict[str, str] | None] + # Optional passthrough of extra fields onto megatron-bridge's + # CheckpointConfig (applied by _create_checkpoint_config). Lets callers + # opt into features NeMo-RL doesn't surface explicitly — e.g. inline HF + # export via {also_save_hf_checkpoint: true, hf_distributed_save: true, + # ckpt_format: torch_dist, non_persistent_ckpt_type: global}. hf_source_path + # defaults to the policy model name when omitted. + checkpoint: NotRequired[dict[str, Any]] # 1 is the minimum recommendation for RL since we almost always need to offload before beginning generation. # Setting to 0 is faster, but you are more likely to run out of GPU memory. In SFT/DPO, the default is 0. empty_unused_memory_level: int