Add Nemotron 3.5 speech prompt support

Author: pzelasko

nemo/collections/common/prompts/__init__.py

@@ -20,6 +20,7 @@
 from nemo.collections.common.prompts.mistral import MistralPromptFormatter
 from nemo.collections.common.prompts.nemotron_h import NemotronHPromptFormatter
 from nemo.collections.common.prompts.nemotron_nano_v3 import NemotronNanoV3PromptFormatter
+from nemo.collections.common.prompts.nemotron3p5 import Nemotron3p5PromptFormatter
 from nemo.collections.common.prompts.phi2 import (
     Phi2ChatPromptFormatter,
     Phi2CodePromptFormatter,

nemo/collections/common/prompts/nemotron3p5.py

@@ -0,0 +1,42 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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 lhotse.cut import Cut
+
+from nemo.collections.common.data.prompt_fn import registered_prompt_format_fn
+from nemo.collections.common.prompts.nemotron_nano_v3 import (
+    NemotronNanoV3PromptFormatter,
+    nemotron_nano_v3,
+)
+
+
+class Nemotron3p5PromptFormatter(NemotronNanoV3PromptFormatter):
+    """Speech prompt formatter for NVIDIA Nemotron 3.5 chat checkpoints.
+
+    Nemotron 3.5 and Nemotron 3 Nano use the same wire format for the
+    system/user/plain-assistant turns supported by SpeechLM. They have
+    separate upstream chat templates, however, so keep a distinct registered
+    name rather than making recipes claim to use the Nano model family.
+
+    The upstream templates differ for structured ``reasoning_content`` and
+    tool-call messages. Those fields are outside the SpeechLM formatter's
+    current ``message``-slot schema and are intentionally not approximated
+    here.
+    """
+
+    NAME = "nemotron3p5"
+
+
+@registered_prompt_format_fn(Cut, Nemotron3p5PromptFormatter)
+def nemotron3p5(cut: Cut, prompt: Nemotron3p5PromptFormatter):
+    return nemotron_nano_v3(cut, prompt)

nemo/collections/speechlm2/models/salm_automodel.py

@@ -60,7 +60,10 @@ def __init__(self, cfg) -> None:
 
         tokenizer_src = self.cfg.get("tokenizer_path", None) or self.cfg.pretrained_llm
         self.tokenizer = AutoTokenizer(
-            tokenizer_src, use_fast=True, trust_remote_code=self.cfg.get("trust_remote_code", False)
+            tokenizer_src,
+            use_fast=True,
+            trust_remote_code=self.cfg.get("trust_remote_code", False),
+            pad_token=self.cfg.get("pad_token", None),
         )
         self.tokenizer.add_special_tokens({"additional_special_tokens": [self.audio_locator_tag]})
         self.speaker_token_ids = build_speaker_tokens(self.cfg.get("speaker_tokens", None), self.tokenizer)

tests/collections/common/prompt_formatters/test_nemotron3p5_prompt_formatter.py

@@ -0,0 +1,58 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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 nemo.collections.common.prompts import PromptFormatter
+from nemo.collections.common.prompts.nemotron3p5 import Nemotron3p5PromptFormatter
+
+
+def test_nemotron3p5_is_registered():
+    assert PromptFormatter.resolve("nemotron3p5") is Nemotron3p5PromptFormatter
+
+
+def test_nemotron3p5_training_basic(bpe_tokenizer_with_think):
+    formatter = Nemotron3p5PromptFormatter(bpe_tokenizer_with_think)
+    ans = formatter.encode_dialog(
+        [
+            {"role": "user", "slots": {"message": "TEST"}},
+            {"role": "assistant", "slots": {"message": "TEST"}},
+        ]
+    )
+
+    assert set(ans) == {"input_ids", "context_ids", "answer_ids", "mask"}
+    assert (
+        bpe_tokenizer_with_think.ids_to_text(ans["input_ids"].tolist())
+        == "<|im_start|>system\n<|im_end|>\n <|im_start|>user\nTEST<|im_end|>\n "
+        "<|im_start|>assistant\n<think></think>TEST<|im_end|>\n"
+    )
+    assert ans["mask"].tolist() == [False] * len(ans["context_ids"]) + [True] * len(ans["answer_ids"])
+
+
+def test_nemotron3p5_inference_generation_prompt(bpe_tokenizer_with_think):
+    formatter = Nemotron3p5PromptFormatter(bpe_tokenizer_with_think)
+    thinking = formatter.encode_dialog(
+        [{"role": "user", "slots": {"message": "TEST"}}], enable_thinking=True
+    )
+    no_thinking = formatter.encode_dialog(
+        [{"role": "user", "slots": {"message": "TEST"}}], enable_thinking=False
+    )
+
+    assert (
+        bpe_tokenizer_with_think.ids_to_text(thinking["input_ids"].tolist())
+        == "<|im_start|>system\n<|im_end|>\n <|im_start|>user\nTEST<|im_end|>\n "
+        "<|im_start|>assistant\n<think>\n"
+    )
+    assert (
+        bpe_tokenizer_with_think.ids_to_text(no_thinking["input_ids"].tolist())
+        == "<|im_start|>system\n<|im_end|>\n <|im_start|>user\nTEST<|im_end|>\n "
+        "<|im_start|>assistant\n<think></think>"
+    )

tests/collections/speechlm2/test_salm_automodel.py

@@ -206,6 +206,48 @@ def test_salm_automodel_training_step_uses_dataloader_iter_signature():
     assert list(inspect.signature(SALMAutomodel.training_step).parameters) == ["self", "dataloader_iter"]
 
 
+def test_salm_automodel_pad_token_override_preserves_eot_labels(monkeypatch):
+    seen = {}
+
+    class FakeTokenizer:
+        def __init__(self, _src, *, use_fast, trust_remote_code, pad_token):
+            seen["pad_token"] = pad_token
+            self.pad = 0 if pad_token == "<unk>" else 11
+            self.unk_id = 0
+
+        def add_special_tokens(self, _tokens):
+            return 0
+
+    salm_module = __import__(
+        "nemo.collections.speechlm2.models.salm_automodel", fromlist=["AutoTokenizer"]
+    )
+    monkeypatch.setattr(salm_module, "AutoTokenizer", FakeTokenizer)
+    model = SALMAutomodel(
+        {
+            "pretrained_llm": "unused",
+            "audio_locator_tag": "<|audio|>",
+            "pad_token": "<unk>",
+        }
+    )
+
+    assert seen["pad_token"] == "<unk>"
+    assert model.text_pad_id == 0
+
+    from nemo.collections.speechlm2.parts.packed_sequences import (
+        prepare_packed_llm_inputs,
+    )
+
+    packed = prepare_packed_llm_inputs(
+        input_ids=torch.tensor([[0, 10, 11, 10, 42, 11]]),
+        text_embs=torch.randn(1, 6, 2),
+        audio_embs=[],
+        target_ids=torch.tensor([[-100, -100, -100, -100, 42, 11]]),
+        padding_id=model.text_pad_id,
+        placeholder_id=999,
+    )
+    assert packed["target_ids"].tolist() == [-100, -100, 42, 11, -100]
+
+
 def test_salm_automodel_fused_linear_forward_keeps_hidden_states_without_logits():
     class FakeLLM(torch.nn.Module):
         def __init__(self):