feat: per-step inline HF export for Megatron checkpoints

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:

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