vmendelev · Oct 27, 2026 · Oct 13, 2026 · Oct 15, 2026 · Oct 15, 2026 · Oct 16, 2026
diff --git a/docs/source/speechlm2/training_and_scaling.rst b/docs/source/speechlm2/training_and_scaling.rst
@@ -3,6 +3,8 @@ Training and Scaling
 
 This page provides detailed information on training speechlm2 models, including setup requirements, running experiments at scale, debugging, and parallelism strategies.
 
+For reward-based transcription fine-tuning, see :ref:`speechlm-gspo`.
+
 Running Experiments
 -------------------
 
@@ -439,3 +441,11 @@ When customizing the training scripts, keep these points in mind:
 1. **Path Overrides**: Override paths in the YAML configuration files with your own, as needed
 2. **W&B Keys**: Update Weights & Biases API keys in configuration files
 3. **Batch Size Tuning**: Adjust batch size based on your GPU memory and model size
+
+.. _speechlm-gspo:
+
+Post-training
+------------------
+
+See
+`SpeechLM GSPO guide <https://github.com/NVIDIA-NeMo/RL/blob/main/docs/guides/gspo-speechlm.md>`_
diff --git a/examples/speechlm2/prepare_rl_export_config.py b/examples/speechlm2/prepare_rl_export_config.py
@@ -0,0 +1,56 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Build canonical native-RL export configuration from the original local HF seed."""
+
+import argparse
+import json
+from pathlib import Path
+
+from nemo.collections.speechlm2.parts.ctc_timestamp_checkpoint import read_ctc_timestamp_config
+
+
+def prepare_config(seed: Path, world_size: int, expert_parallel_size: int) -> dict:
+    seed = seed.resolve()
+    if world_size < 1 or expert_parallel_size < 1 or world_size % expert_parallel_size:
+        raise ValueError('World size must be positive and divisible by expert parallel size')
+    model = json.loads((seed / 'config.json').read_text())
+    if not (seed / 'llm_backbone' / 'config.json').is_file():
+        raise FileNotFoundError(seed / 'llm_backbone' / 'config.json')
+    model['pretrained_llm'] = str(seed / 'llm_backbone')
+    model['tokenizer_path'] = str(seed)
+    ctc_config = read_ctc_timestamp_config(seed / 'model.safetensors')
+    if ctc_config is not None:
+        model['ctc_timestamp_artifact'] = ctc_config
+    strategy = {
+        '_target_': 'nemo.collections.speechlm2.parts.parallel.AutomodelParallelStrategy',
+        'dp_size': world_size,
+        'dp_replicate_size': 1,
+        'tp_size': 1,
+        'pp_size': 1,
+        'cp_size': 1,
+        'ep_size': expert_parallel_size,
+    }
+    return {'model': model, 'trainer': {'strategy': strategy}}
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument('--seed', type=Path, required=True)
+    parser.add_argument('--output', type=Path, required=True)
+    parser.add_argument('--world-size', type=int, default=8)
+    parser.add_argument('--expert-parallel-size', type=int, default=8)
+    parser.add_argument('--strategy-overrides', type=json.loads, default={})
+    parser.add_argument('--backend-config', type=json.loads)
+    args = parser.parse_args()
+    config = prepare_config(args.seed, args.world_size, args.expert_parallel_size)
+    config['trainer']['strategy'].update(args.strategy_overrides)
+    if args.backend_config is not None:
+        config['model']['automodel_backend'] = args.backend_config
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    with args.output.open('x') as output:
+        json.dump(config, output, indent=2)
+        output.write('\n')
+
+
+if __name__ == '__main__':
+    main()
diff --git a/examples/speechlm2/to_hf.py b/examples/speechlm2/to_hf.py
@@ -51,8 +51,33 @@ class HfExportConfig:
     # Dtype used for stored parameters
     dtype: str = "bfloat16"
 
+    # Format of ckpt_path: "lightning" for a Lightning checkpoint, or "nemo_rl" for the
+    # flat DCP directory NeMo-RL saves under weights/model. The output is always HF safetensors.
+    checkpoint_format: str = "lightning"
+
+
+def prepare_native_export_model_config(model_cfg: dict) -> dict:
+    """Keep native MTP topology and enable the canonical dtype adapter."""
+    model_cfg = deepcopy(model_cfg)
+    model_cfg.pop("automodel_freeze_config", None)
+    model_cfg.pop("perception_freeze_config", None)
+    model_cfg["retain_native_mtp"] = True
+    model_cfg["mtp"] = {**(model_cfg.get("mtp") or {}), "enabled": False}
+    model_cfg["automodel_backend"] = {
+        **(model_cfg.get("automodel_backend") or {}),
+        "enable_hf_state_dict_adapter": True,
+    }
+    return model_cfg
+
+
+def load_checkpoint(model: torch.nn.Module, checkpoint_path: str, checkpoint_format: str = "lightning") -> None:
+    if checkpoint_format == "nemo_rl":
+        from nemo.collections.speechlm2.parts.weight_loading import load_nemo_rl_checkpoint
 
-def load_checkpoint(model: torch.nn.Module, checkpoint_path: str) -> None:
+        load_nemo_rl_checkpoint(model, checkpoint_path)
+        return
+    if checkpoint_format != "lightning":
+        raise ValueError(f"Unsupported checkpoint_format={checkpoint_format!r}")
     if Path(checkpoint_path).is_dir():
         from torch.distributed.checkpoint import load
 
@@ -131,6 +156,10 @@ def _hf_export_config(model: torch.nn.Module, dtype: str | torch.dtype) -> dict[
     # Remote-code trust is a runtime security decision. Do not persist a
     # training-time opt-in in checkpoints that may be loaded by another user.
     config.pop("trust_remote_code", None)
+    config.pop("preserve_parameter_dtypes", None)
+    # A new fine-tuning run chooses its own update surface.
+    config.pop("automodel_freeze_config", None)
+    config.pop("perception_freeze_config", None)
     pe_encoder_path = config.get("pe_encoder_path", None)
     pe_encoder_config = config.get("pe_encoder_config", None)
     if pe_encoder_path not in (None, "", False) or pe_encoder_config not in (
@@ -193,6 +222,8 @@ def _hf_export_config(model: torch.nn.Module, dtype: str | torch.dtype) -> dict[
 
     llm = getattr(model, "llm", None)
     text_config = getattr(llm, "config", None)
+    if config.pop("retain_native_mtp", False):
+        config["mtp"] = {**(config.get("mtp") or {}), "enabled": getattr(llm, "mtp", None) is not None}
     explicit_mtp = config.get("mtp")
     mtp_enabled = (
         bool(explicit_mtp.get("enabled", True))
@@ -276,17 +307,33 @@ def save_hf_checkpoint(model: torch.nn.Module, state_dict: dict, cfg: HfExportCo
     if callable(getattr(state_dict_adapter, "forced_hf_dtype_mapping", None)):
         forced_dtypes = state_dict_adapter.forced_hf_dtype_mapping(state_dict)
     state_dict = {
-        key: value.to(torch.float32 if forced_dtypes.get(key) == "F32" else target_dtype)
+        key: (
+            value
+            if key.startswith("ctc_timestamp.")
+            else value.to(torch.float32 if forced_dtypes.get(key) == "F32" else target_dtype)
+        )
         for key, value in state_dict.items()
     }
 
     config = _hf_export_config(model, cfg.dtype)
-    save_file(state_dict, output_dir / "model.safetensors")
+    ctc_timestamp = getattr(model, "ctc_timestamp", None)
+    metadata = ctc_timestamp.artifact_metadata if ctc_timestamp is not None else None
+    save_file(state_dict, output_dir / "model.safetensors", metadata=metadata)
     with open(output_dir / "config.json", "w") as f:
         json.dump(config, f, indent=2)
     save_llm_backbone_config(model, output_dir)
 
 
+def cast_model_for_export(model: torch.nn.Module, dtype: str | torch.dtype) -> torch.nn.Module:
+    """Cast the model to the export dtype; the CTC timestamp artifact keeps its loaded storage."""
+    ctc_timestamp = getattr(model, "ctc_timestamp", None)
+    ctc_state = None if ctc_timestamp is None else ctc_timestamp.state_dict()
+    model = model.to(str_to_dtype(dtype))
+    if ctc_state is not None:
+        ctc_timestamp.load_state_dict(ctc_state, assign=True)
+    return model
+
+
 def save_llm_backbone_config(model: torch.nn.Module, output_dir: str | Path) -> None:
     """Save the original LLM config separately from the NeMo wrapper config."""
     llm_config = getattr(getattr(model, "llm", None), "config", None)
@@ -295,7 +342,18 @@ def save_llm_backbone_config(model: torch.nn.Module, output_dir: str | Path) ->
 
     llm_backbone_dir = Path(output_dir) / LLM_BACKBONE_DIR
     llm_backbone_dir.mkdir(parents=True, exist_ok=True)
-    llm_config.save_pretrained(str(llm_backbone_dir))
+    if getattr(llm_config, "model_type", None) == "nemotron_h":
+        # Transformers >= 5.15 renames these discriminators in memory. Serialize
+        # the original names, which both older and newer serving runtimes accept.
+        # Keep the live model config untouched, including its native MTP topology.
+        config = json.loads(llm_config.to_json_string())
+        aliases = {"linear_attention": "mamba", "full_attention": "attention"}
+        for key in ("layers_block_type", "mtp_layers_block_type"):
+            if config.get(key) is not None:
+                config[key] = [aliases.get(block, block) for block in config[key]]
+        (llm_backbone_dir / "config.json").write_text(json.dumps(config, indent=2) + "\n")
+    else:
+        llm_config.save_pretrained(str(llm_backbone_dir))
 
 
 def _detect_vllm_architecture(model_cfg: dict) -> tuple[str, int]:
@@ -436,12 +494,12 @@ def prepare_for_vllm(output_dir: str, model_cfg: dict) -> None:
     if tok.pad_token_id is not None:
         gen_cfg["pad_token_id"] = tok.pad_token_id
 
-    # Build every tokenizer-side artifact in a sibling staging directory. A
+    # Build every tokenizer-side artifact inside the output filesystem. A
     # caught validation/serialization ``ValueError`` must leave the original
     # HF checkpoint untouched instead of advertising a half-prepared vLLM
     # model. Once staging succeeds, publish each file atomically and write the
     # authoritative root config last.
-    with TemporaryDirectory(prefix=f".{output_dir.name}-vllm-", dir=output_dir.parent) as staging:
+    with TemporaryDirectory(prefix=f".{output_dir.name}-vllm-", dir=output_dir) as staging:
         staging_dir = Path(staging)
         artifacts_dir = staging_dir / "artifacts"
         artifacts_dir.mkdir()
@@ -578,6 +636,8 @@ def main(cfg: HfExportConfig) -> None:
         # The vLLM prompt processor must reserve exactly as many audio
         # placeholders as the checkpoint's encoder emits.
         model_cfg["audio_token_estimator"] = audio_token_estimator
+    if cfg.checkpoint_format == "nemo_rl":
+        model_cfg = prepare_native_export_model_config(model_cfg)
     model_cfg["torch_dtype"] = _canonical_torch_dtype_name(cfg.dtype)
     cls = import_class_by_path(cfg.class_path)
 
@@ -602,7 +662,7 @@ def main(cfg: HfExportConfig) -> None:
         model = cls(model_cfg)
         model.configure_model(distributed_setup=strategy.distributed_setup)
 
-        load_checkpoint(model, cfg.ckpt_path)
+        load_checkpoint(model, cfg.ckpt_path, cfg.checkpoint_format)
 
         # Consolidate DTensors to regular tensors and save on rank 0.
         consolidated = consolidate_state_dict(model)
@@ -616,10 +676,13 @@ def main(cfg: HfExportConfig) -> None:
         model_cfg["init_configure_model"] = True
         model_cfg["pretrained_weights"] = False
         model = cls(model_cfg)
-        load_checkpoint(model, cfg.ckpt_path)
-        model = model.to(str_to_dtype(cfg.dtype))
-        model.save_pretrained(cfg.output_dir, config=_hf_export_config(model, cfg.dtype))
-        save_llm_backbone_config(model, cfg.output_dir)
+        load_checkpoint(model, cfg.ckpt_path, cfg.checkpoint_format)
+        if cfg.checkpoint_format == "nemo_rl":
+            save_hf_checkpoint(model, consolidate_state_dict(model), cfg)
+        else:
+            model = cast_model_for_export(model, cfg.dtype)
+            model.save_pretrained(cfg.output_dir, config=_hf_export_config(model, cfg.dtype))
+            save_llm_backbone_config(model, cfg.output_dir)
         _try_prepare_for_vllm(cfg.output_dir, model_cfg)
 
 

diff --git a/nemo/__init__.py b/nemo/__init__.py
@@ -12,7 +12,11 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from pkgutil import extend_path
 
+# Let other packages that install nemo.* modules (e.g. nemo-lens, used by NeMo RL)
+# import alongside an editable Speech install.
+__path__ = extend_path(__path__, __name__)
 
 from nemo.package_info import (
     __contact_emails__,