chore: import upstream snapshot with attribution

tests/collections/speechlm2/_chunking_helpers.py

@@ -0,0 +1,63 @@
+# Copyright (c) 2025, 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.
+"""Shared fakes for SALM / SALMAutomodel encoder-chunking tests."""
+from types import SimpleNamespace
+
+import torch
+
+
+def chunking_test_devices():
+    devices = [torch.device("cpu")]
+    if torch.cuda.is_available():
+        devices.append(torch.device("cuda"))
+    return devices
+
+
+class ChunkingTestTokenizer:
+    pad = 0
+    unk_id = None
+    bos_id = 1
+    eos_id = 2
+
+    def __init__(self, audio_locator_tag):
+        self._audio_locator_tag = audio_locator_tag
+
+    def token_to_id(self, token):
+        assert token == self._audio_locator_tag
+        return 99
+
+
+class ChunkingTestPerception(torch.nn.Module):
+    def __init__(self, sampling_rate, hop_length):
+        super().__init__()
+        self.preprocessor = SimpleNamespace(featurizer=_ChunkingTestFeaturizer(sampling_rate, hop_length))
+        self.calls = []
+        self.time_offsets = []
+        self.spk_targets_calls = []
+
+    def forward(self, input_signal=None, input_signal_length=None, time_offset=None, spk_targets=None):
+        self.calls.append((input_signal.detach().clone(), input_signal_length.detach().clone()))
+        self.time_offsets.append(None if time_offset is None else time_offset.detach().clone())
+        self.spk_targets_calls.append(None if spk_targets is None else spk_targets.detach().clone())
+        max_len = int(input_signal_length.max().item()) if input_signal_length.numel() > 0 else 0
+        return input_signal[:, :max_len].unsqueeze(-1), input_signal_length.clone()
+
+
+class _ChunkingTestFeaturizer:
+    def __init__(self, sampling_rate, hop_length):
+        self.sample_rate = sampling_rate
+        self.hop_length = hop_length
+
+    def get_seq_len(self, seq_len):
+        return torch.floor_divide(seq_len, self.hop_length).to(dtype=torch.long)

tests/collections/speechlm2/distributed_oomptimizer_model.py

@@ -0,0 +1,97 @@
+# Copyright (c) 2025, 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.
+
+import lightning.pytorch as pl
+import torch
+import torch.nn.functional as F
+
+from nemo.core.neural_types import AudioSignal, LabelsType, LengthsType, NeuralType
+
+
+class SingleBlockDistributedOOMptimizerModel(pl.LightningModule):
+    """Single-transformer-block model used by the distributed OOMptimizer functional test."""
+
+    def __init__(self, cfg: dict):
+        super().__init__()
+        self.vocab_size = int(cfg.get("vocab_size", 64))
+        self.sample_rate = int(cfg.get("sample_rate", 32))
+        self.frame_stride = int(cfg.get("frame_stride", 4))
+        hidden_size = int(cfg.get("hidden_size", 128))
+        num_heads = int(cfg.get("num_heads", 4))
+        ffn_hidden_size = int(cfg.get("ffn_hidden_size", hidden_size * 4))
+        dropout = float(cfg.get("dropout", 0.0))
+        self.activation_reserve_elements_per_frame = int(
+            float(cfg.get("activation_reserve_mb_per_frame", 0.0)) * 1024 * 1024 // 4
+        )
+        self.max_activation_reserve_frames = int(cfg.get("max_activation_reserve_frames", 160))
+
+        self.input_projection = torch.nn.Linear(self.frame_stride, hidden_size)
+        self.encoder = torch.nn.TransformerEncoderLayer(
+            d_model=hidden_size,
+            nhead=num_heads,
+            dim_feedforward=ffn_hidden_size,
+            dropout=dropout,
+            batch_first=True,
+            norm_first=True,
+        )
+        self.classifier = torch.nn.Linear(hidden_size, self.vocab_size)
+
+    @property
+    def oomptimizer_schema(self) -> dict:
+        return {
+            "cls": dict,
+            "inputs": [
+                {"name": "audio", "type": NeuralType(("B", "T"), AudioSignal()), "seq_length": "input"},
+                {"name": "audio_lens", "type": NeuralType(("B",), LengthsType()), "seq_length": "input"},
+                {
+                    "name": "tokens",
+                    "type": NeuralType(("B", "T"), LabelsType()),
+                    "seq_length": "output",
+                    "vocab_size": self.vocab_size,
+                },
+            ],
+        }
+
+    def training_step(self, batch: dict, batch_idx: int) -> dict:
+        audio = batch["audio"].float()
+        tokens = batch["tokens"].long()
+        pad = (-audio.shape[1]) % self.frame_stride
+        if pad:
+            audio = F.pad(audio, (0, pad))
+
+        frames = audio.reshape(audio.shape[0], -1, self.frame_stride)
+        hidden = self.input_projection(frames)
+        hidden = self.encoder(hidden)
+        logits = self.classifier(hidden.mean(dim=1))
+        target = tokens[:, 0].remainder(self.vocab_size)
+        loss = F.cross_entropy(logits, target)
+
+        self._reserve_peak_memory(hidden)
+        return {"loss": loss}
+
+    def _reserve_peak_memory(self, hidden: torch.Tensor) -> None:
+        if self.activation_reserve_elements_per_frame <= 0:
+            return
+        reserve_frames = min(int(hidden.shape[1]), self.max_activation_reserve_frames)
+        if reserve_frames <= 0:
+            return
+
+        # Keep transformer compute small while making memory pressure scale with sequence length.
+        reserve = hidden.new_empty(
+            (int(hidden.shape[0]), reserve_frames, self.activation_reserve_elements_per_frame), dtype=torch.float32
+        )
+        reserve[:, :, :1].zero_()
+
+    def configure_optimizers(self) -> dict:
+        return {"optimizer": torch.optim.SGD(self.parameters(), lr=1e-3)}

tests/collections/speechlm2/test_audio_placeholders.py

@@ -0,0 +1,137 @@
+# Copyright (c) 2025, 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.
+import torch
+
+from nemo.collections.speechlm2.models.salm import replace_placeholders_and_build_targets
+
+
+def test_replace_placeholders():
+    # fmt: off
+    PAD = 0
+    AUDIO = 100
+    input_ids = torch.tensor([
+        [7  , AUDIO, 1, 2    , AUDIO, 1],
+        [PAD,   PAD, 3, AUDIO, 4    , 5]  # note: left padding required
+    ])
+    loss_mask = torch.tensor([
+        [False, False, False, False, False, True],  # predict last token
+        [False, False, False, False, True , True]  # predict last two tokens
+    ])
+    embeds = torch.ones(2, 6, 2)
+    embeds[1, :2] = 0  # note: indicate left padding
+    # 3 embedding sequences with varying shapes, corresponding to 3 AUDIO tokens
+    replacements = [
+        torch.full((4, 2), fill_value=2.0),
+        torch.full((3, 2), fill_value=3.0),
+        torch.full((2, 2), fill_value=4.0),
+    ]
+
+    embeds_r, targets_r, attention_mask_r = replace_placeholders_and_build_targets(
+        input_ids=input_ids,
+        embeds=embeds,
+        padding_id=PAD,
+        placeholder_id=AUDIO,
+        replacements=replacements,
+        target_ids=input_ids.where(loss_mask, -100)
+    )
+
+    assert embeds_r.shape == (2, 11, 2)
+    # batch item 0
+    assert (embeds_r[0, 0]    == 1.0).all()  # 1=orig
+    assert (embeds_r[0, 1:5]  == 2.0).all()  # 2=repl
+    assert (embeds_r[0, 5:7]  == 1.0).all()  # 1=orig
+    assert (embeds_r[0, 7:10] == 3.0).all()  # 3=repl
+    assert (embeds_r[0, 10]   == 1.0).all()  # 1=orig
+    # batch item 1
+    assert (embeds_r[1, :6]   == 0.0).all()  # 0=pad
+    assert (embeds_r[1, 6:7]  == 1.0).all()  # 1=orig
+    assert (embeds_r[1, 7:9]  == 4.0).all()  # 4=repl
+    assert (embeds_r[1, 9:]   == 1.0).all()  # 1=orig
+
+    assert targets_r.shape == (2, 11)
+    torch.testing.assert_close(
+        targets_r,
+        torch.tensor([
+            [-100, -100, -100, -100, -100, -100, -100, -100, -100, -100, 1],
+            [-100, -100, -100, -100, -100, -100, -100, -100, -100, 4   , 5],
+        ])
+    )
+
+    assert attention_mask_r.shape == (2, 11)
+    torch.testing.assert_close(
+        attention_mask_r,
+        torch.tensor([
+            [True, True, True, True, True,  True,  True,  True,  True,  True,  True],
+            [False, False, False, False, False, False, True, True, True, True, True],
+        ])
+    )
+    # fmt: on
+
+
+def test_replace_placeholders_removes_excessive_left_padding():
+    # fmt: off
+    PAD = 0
+    AUDIO = 100
+    input_ids = torch.tensor([
+        [7  , AUDIO, 1    , 2],
+        [PAD,   PAD, AUDIO, 3]  # note: left padding required
+    ])
+    loss_mask = torch.tensor([
+        [False, False, True , True],  # predict last two tokens
+        [False, False, False, True]   # predict last token
+    ])
+    embeds = torch.ones(2, 4, 2)
+    embeds[1, :2] = 0  # note: indicate left padding
+    # 3 embedding sequences with varying shapes, corresponding to 3 AUDIO tokens
+    replacements = [
+        torch.full((3, 2), fill_value=2.0),
+        torch.full((5, 2), fill_value=4.0),
+    ]
+
+    embeds_r, targets_r, attention_mask_r = replace_placeholders_and_build_targets(
+        input_ids=input_ids,
+        embeds=embeds,
+        padding_id=PAD,
+        placeholder_id=AUDIO,
+        replacements=replacements,
+        target_ids=input_ids.where(loss_mask, -100)
+    )
+
+    assert embeds_r.shape == (2, 6, 2)
+    # batch item 0
+    assert (embeds_r[0, 0  ]  == 1.0).all()  # 1=orig
+    assert (embeds_r[0, 1:4]  == 2.0).all()  # 2=repl
+    assert (embeds_r[0, 4: ]  == 1.0).all()  # 1=orig
+    # batch item 1
+    assert (embeds_r[1, :5]  == 4.0).all()  # 4=repl
+    assert (embeds_r[1, 5 ]  == 1.0).all()  # 1=orig
+
+    assert targets_r.shape == (2, 6)
+    torch.testing.assert_close(
+        targets_r,
+        torch.tensor([
+            [-100, -100, -100, -100,    1, 2],
+            [-100, -100, -100, -100, -100, 3],
+        ])
+    )
+
+    assert attention_mask_r.shape == (2, 6)
+    torch.testing.assert_close(
+        attention_mask_r,
+        torch.tensor([
+            [True, True, True, True, True, True],
+            [True, True, True, True, True, True],
+        ])
+    )
+    # fmt: on

tests/collections/speechlm2/test_datamodule.py

@@ -0,0 +1,115 @@
+# Copyright (c) 2025, 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.
+import pytest
+import torch
+from lhotse import CutSet
+from lhotse.testing.dummies import DummyManifest
+from lightning.pytorch.utilities import CombinedLoader
+from omegaconf import DictConfig
+
+from nemo.collections.common.data.lhotse.broadcasting import BroadcastingDataLoader
+from nemo.collections.common.tokenizers.sentencepiece_tokenizer import SentencePieceTokenizer, create_spt_model
+from nemo.collections.speechlm2.data import DataModule
+
+
+@pytest.fixture
+def data_config(tmp_path):
+    ap, cp = tmp_path / "audio", str(tmp_path) + "/{tag}_cuts.jsonl.gz"
+
+    def _assign(k, v):
+        def _inner(obj):
+            setattr(obj, k, v)
+            return obj
+
+        return _inner
+
+    for tag in ("train", "val_set_0", "val_set_1"):
+        (
+            DummyManifest(CutSet, begin_id=0, end_id=2, with_data=True)
+            .map(_assign("tag", tag))
+            .save_audios(ap)
+            .drop_in_memory_data()
+            .to_file(cp.format(tag=tag))
+        )
+
+    return DictConfig(
+        {
+            "train_ds": {
+                "input_cfg": [
+                    {
+                        "type": "lhotse",
+                        "cuts_path": cp.format(tag="train"),
+                    }
+                ],
+                "batch_size": 2,
+            },
+            "validation_ds": {
+                "datasets": {
+                    "val_set_0": {"cuts_path": cp.format(tag="val_set_0")},
+                    "val_set_1": {"cuts_path": cp.format(tag="val_set_1")},
+                },
+                "batch_size": 2,
+            },
+        }
+    )
+
+
+@pytest.fixture
+def tokenizer(tmp_path_factory):
+    tmpdir = tmp_path_factory.mktemp("tok")
+    text_path = tmpdir / "text.txt"
+    text_path.write_text("\n".join(chr(i) for i in range(256)))
+    create_spt_model(
+        text_path,
+        vocab_size=512,
+        sample_size=-1,
+        do_lower_case=False,
+        output_dir=str(tmpdir),
+        bos=True,
+        eos=True,
+        remove_extra_whitespaces=True,
+    )
+    return SentencePieceTokenizer(str(tmpdir / "tokenizer.model"))
+
+
+class Identity(torch.utils.data.Dataset):
+    def __getitem__(self, item):
+        return item
+
+
+def test_datamodule_train_dataloader(data_config, tokenizer):
+    data = DataModule(data_config, tokenizer=tokenizer, dataset=Identity())
+    dl = data.train_dataloader()
+    assert isinstance(dl, (BroadcastingDataLoader, torch.utils.data.DataLoader))
+    dli = iter(dl)
+
+    batch = next(dli)
+    assert isinstance(batch, CutSet)
+    assert len(batch) == 2
+    assert all(c.tag == "train" for c in batch)
+
+
+def test_datamodule_validation_dataloader(data_config, tokenizer):
+    val_sets = {"val_set_0", "val_set_1"}
+    data = DataModule(data_config, tokenizer=tokenizer, dataset=Identity())
+    dl = data.val_dataloader()
+    assert isinstance(dl, CombinedLoader)
+    dli = iter(dl)
+
+    batch, batch_idx, dataloader_idx = next(dli)
+    assert isinstance(batch, dict)
+    assert batch.keys() == val_sets
+    for vs in val_sets:
+        assert len(batch[vs]) == 2
+        assert all(c.tag == vs for c in batch[vs])

tests/collections/speechlm2/test_datamodule_parallel.py

@@ -0,0 +1,402 @@
+# Copyright (c) 2025, 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.
+import os
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.distributed as dist
+from lhotse import CutSet
+from lhotse.testing.dummies import DummyManifest
+from omegaconf import DictConfig
+
+try:
+    from torch.distributed._local_tensor import LocalTensorMode  # noqa: E402
+except ImportError:
+    pytest.skip("Local tensor mode requires PyTorch >= 2.10", allow_module_level=True)
+
+from lightning.pytorch.strategies.model_parallel import _setup_device_mesh
+
+from nemo.collections.common.data.lhotse import get_lhotse_dataloader_from_config
+from nemo.collections.common.tokenizers.sentencepiece_tokenizer import SentencePieceTokenizer, create_spt_model
+from nemo.collections.speechlm2.data import DataModule
+
+nemo_automodel = pytest.importorskip("nemo_automodel")
+from nemo_automodel.components.distributed.config import FSDP2Config  # noqa: E402
+
+from nemo.collections.speechlm2.parts.parallel import AutomodelParallelStrategy  # noqa: E402
+
+# ---------------------------------------------------------------------------
+# Fixtures
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def fake_dist():
+    """Set up minimal env for fake distributed backend, tear down after test."""
+    os.environ.setdefault("MASTER_ADDR", "localhost")
+    os.environ.setdefault("MASTER_PORT", "0")
+    yield
+    if dist.is_initialized():
+        dist.destroy_process_group()
+
+
+@pytest.fixture
+def tokenizer(tmp_path_factory):
+    tmpdir = tmp_path_factory.mktemp("tok")
+    text_path = tmpdir / "text.txt"
+    text_path.write_text("\n".join(chr(i) for i in range(256)))
+    create_spt_model(
+        text_path,
+        vocab_size=512,
+        sample_size=-1,
+        do_lower_case=False,
+        output_dir=str(tmpdir),
+        bos=True,
+        eos=True,
+        remove_extra_whitespaces=True,
+    )
+    return SentencePieceTokenizer(str(tmpdir / "tokenizer.model"))
+
+
+class Identity(torch.utils.data.Dataset):
+    def __getitem__(self, item):
+        return item
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _extract_rank_mapping(sym_val, world_size):
+    """Extract ``{global_rank: value}`` mapping from a SymInt produced by LocalTensorMode.
+
+    Inside ``LocalTensorMode``, ``DeviceMesh.get_local_rank()`` returns a
+    ``SymInt`` backed by a ``LocalIntNode`` (per-rank varying values) or a
+    ``ConstantIntNode`` (same value on every rank).  Arithmetic on these
+    ``SymInt`` values preserves the per-rank structure.
+
+    This helper unwraps the result into a plain dict for easy assertions.
+    """
+    if isinstance(sym_val, int):
+        return {r: sym_val for r in range(world_size)}
+    node = sym_val.node
+    if hasattr(node, "_local_ints"):
+        return dict(node._local_ints)
+    # ConstantIntNode – same value for every rank.
+    val = node.maybe_as_int() if hasattr(node, "maybe_as_int") else node.val
+    return {r: val for r in range(world_size)}
+
+
+# ---------------------------------------------------------------------------
+# Tests
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    "pp,dp_rep,dp_shard,cp,tp",
+    [
+        (1, 1, 1, 1, 1),
+        (1, 2, 2, 1, 1),
+        (1, 1, 2, 1, 1),
+        (2, 3, 2, 1, 2),
+        (2, 2, 4, 2, 2),
+    ],
+)
+def test_dp_rank_via_strategy(fake_dist, pp, dp_rep, dp_shard, cp, tp):
+    """Verify DataModule DP rank using the real AutomodelParallelStrategy.
+
+    Uses ``LocalTensorMode`` to obtain per-rank symbolic results for every
+    global rank in a single process and checks that:
+
+    * The DataModule rank agrees with the strategy's flattened ``"dp"`` submesh.
+    * The correct number of unique DP ranks exists.
+    * All DP ranks are in ``[0, dp_size)``.
+    * Each DP rank covers exactly the non-DP ranks assigned to it.
+    """
+    dp_size = dp_rep * dp_shard
+    world_size = pp * dp_size * cp * tp
+    dist.init_process_group(backend="fake", rank=0, world_size=world_size)
+
+    strategy = AutomodelParallelStrategy(
+        pp_size=pp,
+        tp_size=tp,
+        cp_size=cp,
+        dp_size=dp_size,
+        dp_replicate_size=dp_rep,
+        distributed_config=FSDP2Config(backend="gloo"),
+    )
+
+    with LocalTensorMode(world_size):
+        device_mesh, _ = strategy.create_device_mesh()
+
+        data = DataModule(DictConfig({"train_ds": {"batch_size": 2}}), tokenizer=None, dataset=Identity())
+        data.trainer = SimpleNamespace(model=SimpleNamespace(device_mesh=device_mesh))
+
+        dp_rank_data = data._get_dp_rank()
+        dp_world_size = data._get_world_size()
+        sampler_kwargs = strategy.distributed_sampler_kwargs
+
+    assert dp_world_size == dp_size
+    assert sampler_kwargs["num_replicas"] == dp_size
+
+    data_rank = _extract_rank_mapping(dp_rank_data, world_size)
+    sampler_rank = _extract_rank_mapping(sampler_kwargs["rank"], world_size)
+
+    # DataModule and strategy sampler rank must agree for every global rank.
+    assert data_rank == sampler_rank
+
+    # Correct number of unique dp ranks.
+    assert len(set(data_rank.values())) == dp_size
+
+    # All dp_ranks in valid range.
+    assert all(0 <= v < dp_size for v in data_rank.values())
+
+    ranks_per_dp = {dp_rank: [] for dp_rank in range(dp_size)}
+    for rank, dp_rank in data_rank.items():
+        ranks_per_dp[dp_rank].append(rank)
+    assert all(len(ranks) == pp * cp * tp for ranks in ranks_per_dp.values())
+
+
+def test_non_dp_dims_share_dp_rank(fake_dist):
+    """Ranks that differ only in pp / tp / cp get the same dp_rank."""
+    pp, dp_rep, dp_shard, cp, tp = 2, 3, 2, 1, 2
+    dp_size = dp_rep * dp_shard
+    world_size = pp * dp_size * cp * tp
+    dist.init_process_group(backend="fake", rank=0, world_size=world_size)
+
+    with LocalTensorMode(world_size):
+        strategy = AutomodelParallelStrategy(
+            pp_size=pp,
+            tp_size=tp,
+            cp_size=cp,
+            dp_size=dp_size,
+            dp_replicate_size=dp_rep,
+            distributed_config=FSDP2Config(backend="gloo"),
+        )
+        device_mesh, _ = strategy.create_device_mesh()
+
+        data = DataModule(DictConfig({"train_ds": {"batch_size": 2}}), tokenizer=None, dataset=Identity())
+        data.trainer = SimpleNamespace(model=SimpleNamespace(device_mesh=device_mesh))
+        dp_rank_sym = data._get_dp_rank()
+
+    rank_to_dp = _extract_rank_mapping(dp_rank_sym, world_size)
+
+    # Mesh layout: (pp=2, dp_rep=3, dp_shard=2, cp=1, tp=2)
+    # Rank 0: (pp=0, dp_rep=0, dp_shard=0, cp=0, tp=0)
+    # Rank 1: (pp=0, dp_rep=0, dp_shard=0, cp=0, tp=1)  -- differs in TP
+    # Rank 12: (pp=1, dp_rep=0, dp_shard=0, cp=0, tp=0) -- differs in PP
+    # Rank 4:  (pp=0, dp_rep=1, dp_shard=0, cp=0, tp=0) -- differs in dp_replicate
+    # Rank 2:  (pp=0, dp_rep=0, dp_shard=1, cp=0, tp=0) -- differs in dp_shard
+
+    # Same DP coordinates, different non-DP dims → same dp_rank
+    assert rank_to_dp[0] == rank_to_dp[1], "TP variation should not change dp_rank"
+    assert rank_to_dp[0] == rank_to_dp[12], "PP variation should not change dp_rank"
+
+    # Different DP coordinates → different dp_rank
+    assert rank_to_dp[0] != rank_to_dp[4], "dp_replicate variation must change dp_rank"
+    assert rank_to_dp[0] != rank_to_dp[2], "dp_shard variation must change dp_rank"
+
+
+def test_datamodule_get_dp_rank_automodel(fake_dist, tokenizer):
+    """DataModule._get_dp_rank() / _get_world_size() return correct values."""
+    pp, dp_rep, dp_shard, cp, tp = 2, 3, 2, 1, 2
+    dp_size = dp_rep * dp_shard
+    world_size = pp * dp_size * cp * tp
+    dist.init_process_group(backend="fake", rank=0, world_size=world_size)
+
+    strategy = AutomodelParallelStrategy(
+        pp_size=pp,
+        tp_size=tp,
+        cp_size=cp,
+        dp_size=dp_size,
+        dp_replicate_size=dp_rep,
+        distributed_config=FSDP2Config(backend="gloo"),
+    )
+    # Create mesh outside LocalTensorMode → real int values for fake rank 0
+    device_mesh, _ = strategy.create_device_mesh()
+
+    cfg = DictConfig({"train_ds": {"batch_size": 2}})
+    data = DataModule(cfg, tokenizer=tokenizer, dataset=Identity())
+
+    # Wire up a mock trainer so _get_dp_rank() can find the device mesh.
+    data.trainer = SimpleNamespace(model=SimpleNamespace(device_mesh=device_mesh))
+
+    dp_rank = data._get_dp_rank()
+    dp_ws = data._get_world_size()
+
+    assert dp_rank is not None, "_get_dp_rank() returned None (missing return?)"
+    assert dp_ws is not None, "_get_world_size() returned None (missing return?)"
+    assert dp_rank == 0, f"Fake rank 0 should map to dp_rank 0, got {dp_rank}"
+    assert dp_ws == dp_size, f"Expected dp_world_size={dp_size}, got {dp_ws}"
+
+
+def test_dataloader_data_partitioning(tmp_path, tokenizer):
+    """Different dp_ranks get disjoint data; same dp_rank is deterministic."""
+    dp_world_size = 6
+    num_cuts = 24
+    cuts_path = str(tmp_path / "cuts.jsonl.gz")
+    (
+        DummyManifest(CutSet, begin_id=0, end_id=num_cuts, with_data=True)
+        .save_audios(tmp_path / "audio")
+        .drop_in_memory_data()
+        .to_file(cuts_path)
+    )
+
+    cfg = DictConfig(
+        {
+            "input_cfg": [{"type": "lhotse", "cuts_path": cuts_path}],
+            "batch_size": 2,
+            "force_finite": True,
+            "force_map_dataset": True,
+            "seed": 0,
+            "num_workers": 0,
+        }
+    )
+
+    # Collect cut IDs from each dp_rank.
+    ids_per_rank = {}
+    for dp_rank in range(dp_world_size):
+        dl = get_lhotse_dataloader_from_config(
+            config=cfg,
+            global_rank=dp_rank,
+            world_size=dp_world_size,
+            dataset=Identity(),
+            tokenizer=tokenizer,
+        )
+        ids_per_rank[dp_rank] = [c.id for batch in dl for c in batch]
+
+    # Different DP ranks must receive disjoint cuts.
+    for r1 in range(dp_world_size):
+        for r2 in range(r1 + 1, dp_world_size):
+            overlap = set(ids_per_rank[r1]) & set(ids_per_rank[r2])
+            assert not overlap, f"Ranks {r1} and {r2} share cut IDs: {overlap}"
+
+    # Union of all ranks covers the full dataset.
+    all_ids = set()
+    for ids in ids_per_rank.values():
+        all_ids.update(ids)
+    assert len(all_ids) == num_cuts
+
+    # Same dp_rank called again → identical sequence (deterministic).
+    for dp_rank in (0, dp_world_size - 1):
+        dl_again = get_lhotse_dataloader_from_config(
+            config=cfg,
+            global_rank=dp_rank,
+            world_size=dp_world_size,
+            dataset=Identity(),
+            tokenizer=tokenizer,
+        )
+        ids_again = [c.id for batch in dl_again for c in batch]
+        assert ids_per_rank[dp_rank] == ids_again, f"dp_rank={dp_rank} was not deterministic"
+
+
+# ---------------------------------------------------------------------------
+# Lightning ModelParallelStrategy tests
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    "dp_size,tp_size",
+    [
+        (1, 1),
+        (1, 4),
+        (4, 1),
+        (3, 4),
+        (6, 2),
+    ],
+)
+def test_dp_rank_via_lightning_model_parallel(fake_dist, dp_size, tp_size):
+    """Verify DP rank using Lightning's ``_setup_device_mesh``.
+
+    Lightning's ``ModelParallelStrategy`` creates a 2D mesh with dimensions
+    ``("data_parallel", "tensor_parallel")``.  The DataModule reads DP rank
+    via ``dm["data_parallel"].get_local_rank()``.
+
+    Uses ``LocalTensorMode`` to verify every global rank in one process.
+    """
+    world_size = dp_size * tp_size
+    dist.init_process_group(backend="fake", rank=0, world_size=world_size)
+
+    with LocalTensorMode(world_size):
+        device_mesh = _setup_device_mesh(dp_size, tp_size, world_size, torch.device("cpu"))
+
+        dp_rank_sym = device_mesh["data_parallel"].get_local_rank()
+        dp_world_size = device_mesh["data_parallel"].size()
+
+    assert dp_world_size == dp_size
+
+    rank_to_dp = _extract_rank_mapping(dp_rank_sym, world_size)
+
+    # Correct number of unique dp ranks.
+    assert len(set(rank_to_dp.values())) == dp_size
+
+    # All dp_ranks in valid range.
+    assert all(0 <= v < dp_size for v in rank_to_dp.values())
+
+    # Ranks sharing the same data_parallel coordinate must have the same
+    # dp_rank; different coordinates must differ.
+    mesh_tensor = torch.arange(world_size).reshape(dp_size, tp_size)
+    for dp_coord in range(dp_size):
+        ranks_in_group = mesh_tensor[dp_coord, :].flatten().tolist()
+        dp_ranks_in_group = {rank_to_dp[r] for r in ranks_in_group}
+        assert (
+            len(dp_ranks_in_group) == 1
+        ), f"DP group (data_parallel={dp_coord}) maps to multiple dp_ranks: {dp_ranks_in_group}"
+
+
+def test_lightning_tp_variation_does_not_change_dp_rank(fake_dist):
+    """Ranks that differ only in tensor_parallel get the same dp_rank."""
+    dp_size, tp_size = 3, 4
+    world_size = dp_size * tp_size
+    dist.init_process_group(backend="fake", rank=0, world_size=world_size)
+
+    with LocalTensorMode(world_size):
+        device_mesh = _setup_device_mesh(dp_size, tp_size, world_size, torch.device("cpu"))
+        dp_rank_sym = device_mesh["data_parallel"].get_local_rank()
+
+    rank_to_dp = _extract_rank_mapping(dp_rank_sym, world_size)
+
+    # Mesh layout: (dp=3, tp=4)
+    # Rank 0: (dp=0, tp=0)
+    # Rank 1: (dp=0, tp=1)  -- differs only in TP
+    # Rank 4: (dp=1, tp=0)  -- differs in DP
+
+    assert rank_to_dp[0] == rank_to_dp[1], "TP variation should not change dp_rank"
+    assert rank_to_dp[0] == rank_to_dp[2], "TP variation should not change dp_rank"
+    assert rank_to_dp[0] == rank_to_dp[3], "TP variation should not change dp_rank"
+    assert rank_to_dp[0] != rank_to_dp[4], "DP variation must change dp_rank"
+
+
+def test_datamodule_get_dp_rank_lightning_model_parallel(fake_dist, tokenizer):
+    """DataModule._get_dp_rank() / _get_world_size() with Lightning's 2D mesh."""
+    dp_size, tp_size = 3, 4
+    world_size = dp_size * tp_size
+    dist.init_process_group(backend="fake", rank=0, world_size=world_size)
+
+    # Create mesh outside LocalTensorMode → real int values for fake rank 0
+    device_mesh = _setup_device_mesh(dp_size, tp_size, world_size, torch.device("cpu"))
+
+    cfg = DictConfig({"train_ds": {"batch_size": 2}})
+    data = DataModule(cfg, tokenizer=tokenizer, dataset=Identity())
+    data.trainer = SimpleNamespace(model=SimpleNamespace(device_mesh=device_mesh))
+
+    dp_rank = data._get_dp_rank()
+    dp_ws = data._get_world_size()
+
+    assert dp_rank is not None, "_get_dp_rank() returned None (missing return?)"
+    assert dp_ws is not None, "_get_world_size() returned None (missing return?)"
+    assert dp_rank == 0, f"Fake rank 0 should map to dp_rank 0, got {dp_rank}"
+    assert dp_ws == dp_size, f"Expected dp_world_size={dp_size}, got {dp_ws}"

tests/collections/speechlm2/test_duplex_eartts.py

@@ -0,0 +1,583 @@
+# Copyright (c) 2025, 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.
+
+import os
+import pytest
+import torch
+from lhotse import CutSet, SupervisionSegment
+from lhotse.testing.dummies import dummy_cut, dummy_recording
+
+from nemo.collections.common.data.utils import move_data_to_device
+from nemo.collections.speechlm2.data.duplex_ear_tts_dataset import (
+    DuplexEARTTSDataset,
+    add_speech_delay,
+    sample_audio_segments_repeat,
+)
+from nemo.collections.speechlm2.models import DuplexEARTTS
+
+
+if torch.cuda.is_available():
+    torch.set_default_device('cuda')
+
+
+test_eartts_config = {
+    "model": {
+        "pretrained_lm_name": "nvidia/NVIDIA-Nemotron-Nano-9B-v2",
+        "pretrained_ae_dir": None,
+        "pretrained_tts_model": None,
+        "trust_remote_code": True,
+        "scoring_asr": "stt_en_fastconformer_transducer_large",
+        "freeze_params": [
+            r"^audio_codec\..+$",  # Keep audio codec frozen as it only provides supervision for training.
+            r"^embed_tokens\..+$",  # Keep embed_tokens frozen as done in eartts
+        ],
+        "bos_token": "<s>",
+        "eos_token": "</s>",
+        "pad_token": "<SPECIAL_12>",
+        "audio_codec_run_dtype": "float32",
+        "prevent_freeze_params": [],
+        "audio_save_path": "",
+        "inference_guidance_scale": 0.5,
+        "inference_noise_scale": 0.8,
+        "inference_top_p_or_k": 0.8,
+        "inference_guidance_enabled": False,
+        "subword_mask_exactly_as_eartts": False,
+        "context_hidden_mask_exactly_as_eartts": False,
+        "exclude_norm_from_wd": True,
+        "optimizer": {
+            "_target_": "torch.optim.AdamW",
+            "lr": 4e-5,
+            "betas": [0.9, 0.98],
+            "weight_decay": 0,
+            "foreach": True,
+        },
+        "lr_scheduler": {
+            "_target_": "nemo.core.optim.lr_scheduler.InverseSquareRootAnnealing",
+            "warmup_steps": 2500,
+            "min_lr": 1e-6,
+            "max_steps": 100_000_000,
+        },
+        "codec_config": {
+            "latent_size": 512,
+            "n_fft": 16,
+            "hop_length": 4,
+            "base_hidden_size": 384,
+            "channel_mult": [1, 2, 4],
+            "rates": [7, 7, 9],
+            "num_blocks": 3,
+            "kernel_size": 7,
+            "groups": 1,
+            "codebook_size": 1024,
+            "num_quantizers": 31,
+            "wav_to_token_ratio": 1764,
+        },
+        "tts_config": {
+            "use_gated_fusion_for_text_audio": True,
+            "disable_eos_prediction": True,
+            "use_bos_eos_emb": True,
+            "use_subword_flag_emb": True,
+            "num_delay_speech_tokens": 2,
+            "backbone_type": "gemma3_text",
+            "backbone_model_class": None,
+            "backbone_config_class": None,
+            "backbone_config": {
+                "hidden_size": 1152,
+                "intermediate_size": 4608,
+                "num_hidden_layers": 1,
+                "num_attention_heads": 16,
+                "num_key_value_heads": 16,
+                "head_dim": 72,
+                "attention_dropout": 0.1,
+                "use_cache": False,
+            },
+            "latent_size": 512,
+            "codebook_size": 1024,
+            "num_quantizers": 31,
+            "context_hidden_size": None,
+            "cas_config": {
+                "backbone_type": "t5gemma",
+                "backbone_model_class": None,
+                "backbone_config_class": None,
+                "backbone_config": {
+                    "is_encoder_decoder": False,
+                    "encoder": {
+                        "hidden_size": 1152,
+                        "intermediate_size": 4608,
+                        "num_hidden_layers": 1,
+                        "num_attention_heads": 16,
+                        "num_key_value_heads": 16,
+                        "head_dim": 72,
+                        "use_cache": False,
+                        "attention_dropout": 0.1,
+                    },
+                },
+            },
+            "mog_head_config": {
+                "intermediate_size": 4608,
+                "num_layers": 3,
+                "low_rank": 64,
+                "num_predictions": 1024,
+                "min_log_std": -4.0,
+                "eps": 1e-6,
+            },
+            "p_uncond": 0.1,
+            "label_smoothing": 0.01,
+            "max_training_rate": 0.8,
+            "quantizer_dropout": 0.5,
+            "random_target_masking": False,
+            "exponent": 3.0,
+        },
+    },
+    "trainer": {
+        "devices": -1,
+        "accelerator": "gpu",
+        "num_nodes": 1,
+        "precision": 32,
+        "logger": False,
+        "enable_checkpointing": False,
+        "use_distributed_sampler": False,
+        "max_steps": 100_000_000,
+        "val_check_interval": 1000,
+        "limit_train_batches": "${trainer.val_check_interval}",
+        "limit_val_batches": 2,
+        "log_every_n_steps": 20,
+        "num_sanity_val_steps": 0,
+        "gradient_clip_val": 1.0,
+        "accumulate_grad_batches": 1,
+        "strategy": {
+            "_target_": "lightning.pytorch.strategies.DDPStrategy",
+            "gradient_as_bucket_view": True,
+            "find_unused_parameters": True,
+        },
+    },
+    "data": {
+        "add_text_bos_and_eos_in_each_turn": True,
+        "add_audio_prompt": True,
+        "audio_prompt_duration": 3.0,
+        "frame_length": 0.08,
+        "source_sample_rate": 22050,
+        "target_sample_rate": 22050,
+        "input_roles": ["user", "User"],
+        "output_roles": ["agent", "Assistant", "assistant", "Agent"],
+    },
+    "exp_manager": {
+        "exp_dir": None,
+        "explicit_log_dir": "",
+        "name": "eartts",
+        "create_tensorboard_logger": False,
+        "create_checkpoint_callback": True,
+        "use_datetime_version": True,
+        "max_time_per_run": "00:03:50:00",
+        "resume_from_checkpoint": None,
+        "resume_if_exists": True,
+        "resume_ignore_no_checkpoint": True,
+        "create_wandb_logger": True,
+        "wandb_logger_kwargs": {
+            "name": "duplex_eartts_test",
+            "project": "duplex_eartts",
+            "resume": True,
+        },
+    },
+}
+
+# set CI cached path
+if os.path.exists("/home/TestData/nvidia--NVIDIA-Nemotron-Nano-9B-v2/"):
+    test_eartts_config["model"]["pretrained_lm_name"] = "/home/TestData/nvidia--NVIDIA-Nemotron-Nano-9B-v2/"
+
+
+@pytest.fixture(scope="session")
+def model():
+    model = DuplexEARTTS(test_eartts_config)
+    if torch.cuda.is_available():
+        model.to("cuda")
+    return model
+
+
+@pytest.fixture(scope="session")
+def dataset(model):
+    return DuplexEARTTSDataset(
+        model.tokenizer,
+        add_text_bos_and_eos_in_each_turn=True,
+        add_audio_prompt=True,
+        audio_prompt_duration=3.0,
+        frame_length=0.08,
+        source_sample_rate=22050,
+        target_sample_rate=22050,
+        input_roles=["user", "User"],
+        output_roles=["agent", "Assistant", "assistant", "Agent"],
+    )
+
+
+@pytest.fixture(scope="session")
+def training_cutset_batch():
+    cut = dummy_cut(0, recording=dummy_recording(0, with_data=True, duration=1.0, sampling_rate=22050))
+    cut.target_audio = dummy_recording(1, with_data=True, duration=1.0, sampling_rate=22050)
+    cut.supervisions = [
+        SupervisionSegment(
+            id=cut.id,
+            recording_id=cut.recording_id,
+            start=0,
+            duration=0.1,
+            text='hi',
+            speaker="user",
+        ),
+        SupervisionSegment(
+            id=cut.id,
+            recording_id=cut.recording_id,
+            start=0.3,
+            duration=0.1,
+            text='hello',
+            speaker="assistant",
+        ),
+        SupervisionSegment(
+            id=cut.id,
+            recording_id=cut.recording_id,
+            start=0.5,
+            duration=0.1,
+            text='ok',
+            speaker="user",
+        ),
+        SupervisionSegment(
+            id=cut.id,
+            recording_id=cut.recording_id,
+            start=0.6,
+            duration=0.1,
+            text='okay',
+            speaker="assistant",
+        ),
+    ]
+    return CutSet([cut])
+
+
+def test_eartts_dataset(dataset, training_cutset_batch):
+    batch = dataset[training_cutset_batch]
+    expected_keys = {
+        "sample_id",
+        "non_prompt_mask",
+        "prompt_lens",
+        "aligned_attention_mask",
+        "aligned_position_ids",
+        "source_audio",
+        "source_audio_lens",
+        "target_audio",
+        "target_audio_lens",
+        "target_text_tokens",
+        "target_token_lens",
+        "source_tokens",
+        "source_token_lens",
+        "target_texts",
+        "audio_prompt",
+        "audio_prompt_lens",
+        "task",
+    }
+
+    for key in expected_keys:
+        assert key in batch, f"Missing key: {key}"
+
+    tensor_keys = [
+        "non_prompt_mask",
+        "aligned_attention_mask",
+        "aligned_position_ids",
+        "source_audio",
+        "source_audio_lens",
+        "target_audio",
+        "target_audio_lens",
+        "target_text_tokens",
+        "target_token_lens",
+        "source_tokens",
+        "source_token_lens",
+        "audio_prompt",
+        "audio_prompt_lens",
+    ]
+
+    for key in tensor_keys:
+        assert torch.is_tensor(batch[key]), f"{key} must be a tensor"
+
+    # Check target text consistency
+    assert batch["target_texts"] == ["hello okay"]
+    assert batch["source_tokens"].tolist() == [
+        [
+            2,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            2,
+            1,
+            2,
+            12,
+            12,
+            12,
+            12,
+            1,
+            1662,
+            2,
+            12,
+            12,
+            12,
+            12,
+        ]
+    ]
+
+    assert batch["target_text_tokens"].tolist() == [
+        [
+            2,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            12,
+            2,
+            12,
+            12,
+            12,
+            12,
+            1,
+            2,
+            12,
+            12,
+            1,
+            2,
+            1417,
+            12,
+            12,
+        ]
+    ]
+
+    # Check task
+    assert batch["task"] == ["s2s_duplex"]
+
+
+# test extra functions inside of eartts dataset
+def test_add_speech_delay():
+    source_audio = torch.ones(1, 16000)
+    target_audio = torch.ones(1, 22050)
+
+    source_lens = torch.tensor([16000])
+    target_lens = torch.tensor([22050])
+
+    num_delays = 2
+
+    # samples per frame (float → int handled explicitly)
+    target_samples_per_frame = source_audio.size(1) / 12.5
+    source_samples_per_frame = target_audio.size(1) / 12.5
+
+    expected_extra_src_size = int(source_samples_per_frame * num_delays)
+    expected_extra_tgt_size = int(target_samples_per_frame * num_delays)
+
+    out_src, out_src_lens, out_tgt, out_tgt_lens = add_speech_delay(
+        source_audio=source_audio,
+        source_audio_lens=source_lens,
+        target_audio=target_audio,
+        target_audio_lens=target_lens,
+        num_delay_speech_tokens=num_delays,
+        target_samples_per_frame=target_samples_per_frame,
+        source_samples_per_frame=source_samples_per_frame,
+    )
+
+    # --------------------------------------------------
+    # Shape & length bookkeeping
+    # --------------------------------------------------
+    assert out_src.shape == (1, source_audio.size(1) + expected_extra_src_size)
+    assert out_tgt.shape == (1, target_audio.size(1) + expected_extra_tgt_size)
+    assert out_src_lens.item() == source_lens.item() + expected_extra_src_size
+    assert out_tgt_lens.item() == target_lens.item() + expected_extra_tgt_size
+
+    # --------------------------------------------------
+    # Padding direction & content
+    # --------------------------------------------------
+    # Target audio is left-padded
+    assert torch.all(out_tgt[:, :expected_extra_tgt_size] == 0)
+    assert torch.all(out_tgt[:, expected_extra_tgt_size:] == 1)
+
+    # Source audio is right-padded
+    assert torch.all(out_src[:, : source_audio.size(1)] == 1)
+    assert torch.all(out_src[:, source_audio.size(1) :] == 0)
+
+
+def test_sample_audio_segments_repeat():
+    cases = [
+        # (audio, lens, n_sample, expected_when_sample_false)
+        (
+            torch.tensor([[1.0, 2.0, 3.0, 4.0, 5.0]]),
+            torch.tensor([5]),
+            3,
+            torch.tensor([[1.0, 2.0, 3.0]]),
+        ),
+        (
+            torch.tensor([[1.0, 2.0]]),
+            torch.tensor([2]),
+            5,
+            torch.tensor([[1.0, 2.0, 1.0, 2.0, 1.0]]),
+        ),
+        (
+            torch.zeros(1, 10),
+            torch.tensor([0]),
+            4,
+            torch.zeros(1, 4),
+        ),
+    ]
+
+    for prompt_audio, prompt_audio_lens, n_sample, expected in cases:
+        # --------------------------------------------------
+        # sample=False → deterministic + sequence check
+        # --------------------------------------------------
+        out = sample_audio_segments_repeat(
+            prompt_audio,
+            prompt_audio_lens,
+            n_sample=n_sample,
+            sample=False,
+        )
+
+        assert out.shape == expected.shape
+        assert torch.equal(out, expected)
+
+        # --------------------------------------------------
+        # sample=True → stochastic, shape only
+        # --------------------------------------------------
+        out = sample_audio_segments_repeat(
+            prompt_audio,
+            prompt_audio_lens,
+            n_sample=n_sample,
+            sample=True,
+        )
+
+        assert out.shape == expected.shape
+
+
+def test_eartts_training_step(model, dataset, training_cutset_batch):
+    model.train()
+    model.on_train_epoch_start()
+    batch = dataset[training_cutset_batch]
+    batch = move_data_to_device(batch, device=model.device)
+    results = model.training_step(batch, batch_idx=0)
+    assert torch.is_tensor(results["loss"])
+    assert not torch.isnan(results["loss"])
+    assert results["loss"] > 0
+
+
+def test_eartts_validation_step(model, dataset, training_cutset_batch):
+    model.eval()
+    model.on_validation_epoch_start()
+    batch = dataset[training_cutset_batch]
+    batch = move_data_to_device(batch, device=model.device)
+    results = model.validation_step({"dummy_val_set": batch}, batch_idx=0)
+    assert results is None  # no return value
+
+
+def test_eartts_offline_generation(model):
+    model.eval()
+    # generate random subword_ids
+    subword_ids = torch.ones(2, 10).long()
+
+    # set init inputs and get it
+    model.set_init_inputs(
+        speaker_audio=torch.randn(1, 22050),
+        speaker_audio_lens=torch.tensor([22050]),
+    )
+    init_inputs = model.get_init_inputs(B=subword_ids.size(0))
+    torch.manual_seed(42)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(42)
+
+    gen_audio, gen_audio_len = model.offline_inference(
+        next_subword_ids=subword_ids,
+        init_inputs=init_inputs,
+    )
+    torch.manual_seed(42)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(42)
+
+    gen_audio_inc, gen_audio_len_inc = model.offline_inference(
+        next_subword_ids=subword_ids, init_inputs=init_inputs, incremental_audio_decoding=True
+    )
+
+    assert torch.equal(
+        gen_audio_len, gen_audio_len_inc
+    ), "Audio lengths differ between incremental and non-incremental decoding."
+
+    # compare waveform
+    torch.testing.assert_close(
+        gen_audio,
+        gen_audio_inc,
+        atol=1e-1,
+        rtol=0,
+    )
+
+    assert gen_audio.shape == (2, 17640)
+    assert gen_audio_len[0] == gen_audio.size(-1)
+    assert gen_audio.dtype == torch.float32

tests/collections/speechlm2/test_duplex_s2s.py

@@ -0,0 +1,227 @@
+# Copyright (c) 2025, 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.
+import os
+
+import pytest
+import torch
+from lhotse import CutSet, SupervisionSegment
+from lhotse.testing.dummies import dummy_cut, dummy_recording
+
+from nemo.collections.common.data.utils import move_data_to_device
+from nemo.collections.speechlm2.data import DuplexS2SDataset
+from nemo.collections.speechlm2.models import DuplexS2SModel
+
+if torch.cuda.is_available():
+    torch.set_default_device('cuda')
+
+
+def resolve_pretrained_models():
+    if os.path.exists("/home/TestData/speechlm/pretrained_models"):
+        # CI pre-cached paths:
+        return {
+            "pretrained_llm": "/home/TestData/speechlm/pretrained_models/TinyLlama--TinyLlama_v1.1",
+            "pretrained_audio_codec": "/home/TestData/speechlm/pretrained_models/low-frame-rate-speech-codec-22khz.nemo",
+            "pretrained_asr": "/home/TestData/speechlm/pretrained_models/stt_en_fastconformer_hybrid_large_streaming_80ms.nemo",
+            "scoring_asr": "/home/TestData/speechlm/pretrained_models/stt_en_fastconformer_transducer_large.nemo",
+        }
+    else:
+        # HF URLs:
+        return {
+            "pretrained_asr": "stt_en_fastconformer_hybrid_large_streaming_80ms",
+            "scoring_asr": "stt_en_fastconformer_transducer_large",
+            "pretrained_llm": "TinyLlama/TinyLlama_v1.1",
+            "pretrained_audio_codec": "nvidia/low-frame-rate-speech-codec-22khz",
+        }
+
+
+@pytest.fixture(scope="session")
+def model():
+    cfg = {
+        **resolve_pretrained_models(),
+        "pretrained_weights": False,
+        "freeze_params": ["^audio_codec\\..+$"],
+        "audio_loss_weight": 1,
+        "text_loss_weight": 3,
+        "perception": {
+            "target": "nemo.collections.speechlm2.modules.perception.AudioPerceptionModule",
+            "output_dim": 2048,
+            "encoder": {
+                "_target_": "nemo.collections.asr.modules.ConformerEncoder",
+                "att_context_size": [-1, -1],
+                "causal_downsampling": False,
+                "conv_context_size": None,
+                "conv_kernel_size": 9,
+                "conv_norm_type": "batch_norm",
+                "d_model": 1024,
+                "dropout": 0.1,
+                "dropout_att": 0.1,
+                "dropout_emb": 0.0,
+                "dropout_pre_encoder": 0.1,
+                "feat_in": 128,
+                "feat_out": -1,
+                "ff_expansion_factor": 4,
+                "n_heads": 8,
+                "n_layers": 2,
+                "pos_emb_max_len": 5000,
+                "self_attention_model": "rel_pos",
+                "subsampling": "dw_striding",
+                "subsampling_conv_channels": 256,
+                "subsampling_factor": 8,
+            },
+            "modality_adapter": {
+                "_target_": "nemo.collections.speechlm2.modules.perception.IdentityConnector",
+                "d_model": 1024,
+            },
+            "preprocessor": {
+                "_target_": "nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor",
+                "dither": 1e-05,
+                "features": 128,
+                "frame_splicing": 1,
+                "log": True,
+                "n_fft": 512,
+                "normalize": "per_feature",
+                "pad_to": 0,
+                "pad_value": 0.0,
+                "sample_rate": 16000,
+                "window": "hann",
+                "window_size": 0.025,
+                "window_stride": 0.01,
+            },
+        },
+        "optimizer": {"_target_": "torch.optim.AdamW"},
+    }
+    model = DuplexS2SModel(cfg)
+    if torch.cuda.is_available():
+        model.to("cuda")
+    return model
+
+
+@pytest.fixture(scope="session")
+def dataset(model):
+    return DuplexS2SDataset(
+        model.tokenizer,
+        frame_length=0.08,
+        source_sample_rate=16000,
+        target_sample_rate=22050,
+        input_roles=["user"],
+        output_roles=["assistant"],
+    )
+
+
+@pytest.fixture(scope="session")
+def training_cutset_batch():
+    cut = dummy_cut(0, recording=dummy_recording(0, with_data=True))
+    cut.target_audio = dummy_recording(1, with_data=True)
+    cut.supervisions = [
+        SupervisionSegment(
+            id=cut.id,
+            recording_id=cut.recording_id,
+            start=0,
+            duration=0.1,
+            text='hi',
+            speaker="user",
+        ),
+        SupervisionSegment(
+            id=cut.id,
+            recording_id=cut.recording_id,
+            start=0.3,
+            duration=0.1,
+            text='hello',
+            speaker="assistant",
+        ),
+        SupervisionSegment(
+            id=cut.id,
+            recording_id=cut.recording_id,
+            start=0.5,
+            duration=0.1,
+            text='ok',
+            speaker="user",
+        ),
+        SupervisionSegment(
+            id=cut.id,
+            recording_id=cut.recording_id,
+            start=0.6,
+            duration=0.4,
+            text='okay',
+            speaker="assistant",
+        ),
+    ]
+    return CutSet([cut])
+
+
+def test_s2s_dataset(dataset, training_cutset_batch):
+    batch = dataset[training_cutset_batch]
+    for key in (
+        "source_audio",
+        "target_audio",
+        "source_audio_lens",
+        "target_audio_lens",
+        "target_tokens",
+        "target_token_lens",
+        "source_tokens",
+        "source_token_lens",
+    ):
+        assert key in batch
+        assert torch.is_tensor(batch[key])
+    assert batch["source_audio"].shape == (1, 16000)
+    assert batch["target_audio"].shape == (1, 22050)
+    assert batch["target_texts"] == ["hello okay"]
+    assert batch["target_tokens"].tolist() == [[0, 0, 0, 0, 1, 2, 0, 0, 1, 20759, 0, 0, 0]]
+    assert batch["source_tokens"].tolist() == [[1, 2, 0, 0, 0, 0, 1, 3431, 2, 0, 0, 0, 0]]
+
+
+def test_s2s_training_step(model, dataset, training_cutset_batch):
+    model.on_train_epoch_start()
+    batch = dataset[training_cutset_batch]
+    batch = move_data_to_device(batch, device=model.device)
+    results = model.training_step(batch, batch_idx=0)
+    assert torch.is_tensor(results["loss"])
+    assert not torch.isnan(results["loss"])
+    assert results["loss"] > 0
+
+
+def test_s2s_validation_step(model, dataset, training_cutset_batch):
+    model.on_validation_epoch_start()
+    batch = dataset[training_cutset_batch]
+    batch = move_data_to_device(batch, device=model.device)
+    results = model.validation_step({"dummy_val_set": batch}, batch_idx=0)
+    assert results is None  # no return value
+
+
+def test_s2s_offline_generation(model):
+    # 16000 samples == 1 second == 12.5 frames ~= 14 frames after encoder padding
+    ans = model.offline_inference(
+        input_signal=torch.randn(1, 16000),
+        input_signal_lens=torch.tensor([16000]),
+    )
+
+    assert ans.keys() == {"text", "tokens_text", "tokens_audio", "audio", "audio_len", "tokens_len"}
+
+    assert isinstance(ans["text"], list)
+    assert isinstance(ans["text"][0], str)
+
+    gen_text = ans["tokens_text"]
+    assert gen_text.shape == (1, 13)
+    assert gen_text.dtype == torch.long
+    assert (gen_text >= 0).all()
+    assert (gen_text < model.text_vocab_size).all()
+
+    gen_audio_codes = ans["tokens_audio"]
+    assert gen_audio_codes.shape == (1, 13, 8)
+    assert gen_audio_codes.dtype == torch.long
+    assert (gen_audio_codes >= 0).all()
+    assert (gen_audio_codes < model.speech_vocab_size).all()
+
+    gen_audio = ans["audio"]
+    assert gen_audio.dtype == torch.float32

tests/collections/speechlm2/test_duplex_s2s_speech_decoder.py

@@ -0,0 +1,214 @@
+# Copyright (c) 2025, 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.
+import os
+
+import pytest
+import torch
+from lhotse import CutSet, SupervisionSegment
+from lhotse.testing.dummies import dummy_cut, dummy_recording
+
+from nemo.collections.common.data.utils import move_data_to_device
+from nemo.collections.speechlm2.data import DuplexS2SDataset
+from nemo.collections.speechlm2.models import DuplexS2SSpeechDecoderModel
+
+if torch.cuda.is_available():
+    torch.set_default_device('cuda')
+
+
+def resolve_pretrained_models():
+    if os.path.exists("/home/TestData/speechlm/pretrained_models"):
+        # CI pre-cached paths:
+        return {
+            "pretrained_llm": "/home/TestData/speechlm/pretrained_models/TinyLlama--TinyLlama_v1.1",
+            "pretrained_audio_codec": "/home/TestData/speechlm/pretrained_models/low-frame-rate-speech-codec-22khz.nemo",
+            "pretrained_asr": "/home/TestData/speechlm/pretrained_models/stt_en_fastconformer_hybrid_large_streaming_80ms.nemo",
+            "scoring_asr": "/home/TestData/speechlm/pretrained_models/stt_en_fastconformer_transducer_large.nemo",
+        }
+    else:
+        # HF URLs:
+        return {
+            "pretrained_asr": "stt_en_fastconformer_hybrid_large_streaming_80ms",
+            "scoring_asr": "stt_en_fastconformer_transducer_large",
+            "pretrained_llm": "TinyLlama/TinyLlama_v1.1",
+            "pretrained_audio_codec": "nvidia/low-frame-rate-speech-codec-22khz",
+        }
+
+
+@pytest.fixture(scope="session")
+def model():
+    cfg = {
+        **resolve_pretrained_models(),
+        "pretrained_weights": False,
+        "freeze_params": ["^audio_codec\\..+$"],
+        "audio_loss_weight": 1,
+        "text_loss_weight": 3,
+        "perception": {
+            "target": "nemo.collections.speechlm2.modules.perception.AudioPerceptionModule",
+            "output_dim": 2048,
+            "encoder": {
+                "_target_": "nemo.collections.asr.modules.ConformerEncoder",
+                "att_context_size": [-1, -1],
+                "causal_downsampling": False,
+                "conv_context_size": None,
+                "conv_kernel_size": 9,
+                "conv_norm_type": "batch_norm",
+                "d_model": 1024,
+                "dropout": 0.1,
+                "dropout_att": 0.1,
+                "dropout_emb": 0.0,
+                "dropout_pre_encoder": 0.1,
+                "feat_in": 128,
+                "feat_out": -1,
+                "ff_expansion_factor": 4,
+                "n_heads": 8,
+                "n_layers": 2,
+                "pos_emb_max_len": 5000,
+                "self_attention_model": "rel_pos",
+                "subsampling": "dw_striding",
+                "subsampling_conv_channels": 256,
+                "subsampling_factor": 8,
+            },
+            "modality_adapter": {
+                "_target_": "nemo.collections.speechlm2.modules.perception.IdentityConnector",
+                "d_model": 1024,
+            },
+            "preprocessor": {
+                "_target_": "nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor",
+                "dither": 1e-05,
+                "features": 128,
+                "frame_splicing": 1,
+                "log": True,
+                "n_fft": 512,
+                "normalize": "per_feature",
+                "pad_to": 0,
+                "pad_value": 0.0,
+                "sample_rate": 16000,
+                "window": "hann",
+                "window_size": 0.025,
+                "window_stride": 0.01,
+            },
+        },
+        "speech_decoder": {
+            "n_layers": 1,
+            "d_model": 768,
+            "d_ffn": 3072,
+            "sa_n_heads": 12,
+            "kernel_size": 3,
+            "is_causal": True,
+        },
+        "optimizer": {"_target_": "torch.optim.AdamW"},
+    }
+    model = DuplexS2SSpeechDecoderModel(cfg)
+    if torch.cuda.is_available():
+        model.to("cuda")
+    return model
+
+
+@pytest.fixture(scope="session")
+def dataset(model):
+    return DuplexS2SDataset(
+        model.tokenizer,
+        frame_length=0.08,
+        source_sample_rate=16000,
+        target_sample_rate=22050,
+        input_roles=["user"],
+        output_roles=["assistant"],
+    )
+
+
+@pytest.fixture(scope="session")
+def training_cutset_batch():
+    cut = dummy_cut(0, recording=dummy_recording(0, with_data=True))
+    cut.target_audio = dummy_recording(1, with_data=True)
+    cut.supervisions = [
+        SupervisionSegment(
+            id=cut.id,
+            recording_id=cut.recording_id,
+            start=0,
+            duration=0.1,
+            text='hi',
+            speaker="user",
+        ),
+        SupervisionSegment(
+            id=cut.id,
+            recording_id=cut.recording_id,
+            start=0.3,
+            duration=0.1,
+            text='hello',
+            speaker="assistant",
+        ),
+        SupervisionSegment(
+            id=cut.id,
+            recording_id=cut.recording_id,
+            start=0.5,
+            duration=0.1,
+            text='ok',
+            speaker="user",
+        ),
+        SupervisionSegment(
+            id=cut.id,
+            recording_id=cut.recording_id,
+            start=0.6,
+            duration=0.4,
+            text='okay',
+            speaker="assistant",
+        ),
+    ]
+    return CutSet([cut])
+
+
+def test_s2s_speech_decoder_training_step(model, dataset, training_cutset_batch):
+    model.on_train_epoch_start()
+    batch = dataset[training_cutset_batch]
+    batch = move_data_to_device(batch, device=model.device)
+    results = model.training_step(batch, batch_idx=0)
+    assert torch.is_tensor(results["loss"])
+    assert not torch.isnan(results["loss"])
+    assert results["loss"] > 0
+
+
+def test_s2s_speech_decoder_validation_step(model, dataset, training_cutset_batch):
+    model.on_validation_epoch_start()
+    batch = dataset[training_cutset_batch]
+    batch = move_data_to_device(batch, device=model.device)
+    results = model.validation_step({"dummy_val_set": batch}, batch_idx=0)
+    assert results is None  # no return value
+
+
+def test_s2s_speech_decoder_offline_generation(model):
+    # 16000 samples == 1 second == 12.5 frames ~= 14 frames after encoder padding
+    ans = model.offline_inference(
+        input_signal=torch.randn(1, 16000, device=model.device),
+        input_signal_lens=torch.tensor([16000], device=model.device),
+    )
+
+    assert ans.keys() == {"text", "tokens_text", "tokens_audio", "audio", "audio_len", "tokens_len"}
+
+    assert isinstance(ans["text"], list)
+    assert isinstance(ans["text"][0], str)
+
+    gen_text = ans["tokens_text"]
+    assert gen_text.shape == (1, 13)
+    assert gen_text.dtype == torch.long
+    assert (gen_text >= 0).all()
+    assert (gen_text < model.text_vocab_size).all()
+
+    gen_audio_codes = ans["tokens_audio"]
+    assert gen_audio_codes.shape == (1, 13, 8)
+    assert gen_audio_codes.dtype == torch.long
+    assert (gen_audio_codes >= 0).all()
+    assert (gen_audio_codes < model.speech_vocab_size).all()
+
+    gen_audio = ans["audio"]
+    assert gen_audio.dtype == torch.float32

tests/collections/speechlm2/test_duplex_stt.py

@@ -0,0 +1,327 @@
+# Copyright (c) 2025, 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.
+
+import os
+
+import pytest
+import torch
+from lhotse import CutSet, SupervisionSegment
+from lhotse.testing.dummies import dummy_cut, dummy_recording
+
+from nemo.collections.common.data.utils import move_data_to_device
+from nemo.collections.speechlm2.data import DuplexSTTDataset
+from nemo.collections.speechlm2.models import DuplexSTTModel
+
+if torch.cuda.is_available():
+    torch.set_default_device('cuda')
+
+
+def resolve_pretrained_models():
+    if os.path.exists("/home/TestData/speechlm/pretrained_models"):
+        return {
+            "pretrained_llm": "/home/TestData/speechlm/pretrained_models/TinyLlama--TinyLlama_v1.1",
+        }
+    return {
+        "pretrained_llm": "TinyLlama/TinyLlama_v1.1",
+    }
+
+
+def create_model(
+    predict_user_text=False,
+    force_use_noise_augmentation=False,
+    old_noise_prob=0.0,
+    old_noise_min_snr=0.0,
+    old_noise_max_snr=0.0,
+):
+    """Helper function to create a model with configurable settings."""
+    cfg = {
+        "model": {
+            **resolve_pretrained_models(),
+            "pretrained_weights": False,
+            "trust_remote_code": True,
+            "audio_loss_weight": 1,
+            "text_loss_weight": 3,
+            "source_sample_rate": 16000,
+            "validation_save_path": "/tmp/test_duplex_stt_logs",
+            "perception": {
+                "_target_": "nemo.collections.speechlm2.modules.perception.AudioPerceptionModule",
+                "preprocessor": {
+                    "_target_": "nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor",
+                    "features": 80,
+                },
+                "encoder": {
+                    "_target_": "nemo.collections.asr.modules.ConformerEncoder",
+                    "feat_in": 80,
+                    "d_model": 512,
+                    "n_heads": 8,
+                    "n_layers": 1,
+                    "subsampling_factor": 8,
+                },
+                "modality_adapter": {
+                    "_target_": "nemo.collections.speechlm2.modules.perception.IdentityConnector",
+                    "d_model": 512,
+                },
+                "output_dim": 2048,
+            },
+            "predict_user_text": predict_user_text,
+            "force_use_noise_augmentation": force_use_noise_augmentation,
+            "old_noise_prob": old_noise_prob,
+            "old_noise_min_snr": old_noise_min_snr,
+            "old_noise_max_snr": old_noise_max_snr,
+            "optimizer": {"_target_": "torch.optim.AdamW"},
+        },
+        "data": {
+            "source_sample_rate": 16000,
+        },
+        "exp_manager": {
+            "explicit_log_dir": "/tmp/test_duplex_stt_logs",
+        },
+    }
+    model = DuplexSTTModel(cfg["model"])
+    if torch.cuda.is_available():
+        model.to("cuda")
+    return model
+
+
+@pytest.fixture(scope="session")
+def model():
+    return create_model(predict_user_text=False)
+
+
+@pytest.fixture(scope="session")
+def dataset(model):
+    return DuplexSTTDataset(
+        model.tokenizer,
+        frame_length=0.08,
+        source_sample_rate=16000,
+        input_roles=["user"],
+        output_roles=["assistant"],
+    )
+
+
+@pytest.fixture(scope="session")
+def training_cutset_batch():
+    cut = dummy_cut(0, recording=dummy_recording(0, with_data=True))
+    cut.supervisions = [
+        SupervisionSegment(
+            id=cut.id,
+            recording_id=cut.recording_id,
+            start=0,
+            duration=0.1,
+            text='hi',
+            speaker="user",
+        ),
+        SupervisionSegment(
+            id=cut.id,
+            recording_id=cut.recording_id,
+            start=0.3,
+            duration=0.1,
+            text='hello',
+            speaker="assistant",
+        ),
+        SupervisionSegment(
+            id=cut.id,
+            recording_id=cut.recording_id,
+            start=0.5,
+            duration=0.1,
+            text='ok',
+            speaker="user",
+        ),
+        SupervisionSegment(
+            id=cut.id,
+            recording_id=cut.recording_id,
+            start=0.6,
+            duration=0.4,
+            text='okay',
+            speaker="assistant",
+        ),
+    ]
+    return CutSet([cut])
+
+
+def test_stt_training_step(model, dataset, training_cutset_batch):
+    model.on_train_epoch_start()
+    batch = dataset[training_cutset_batch]
+    batch = move_data_to_device(batch, device=model.device)
+    results = model.training_step(batch, batch_idx=0)
+    assert torch.is_tensor(results["loss"])
+    assert not torch.isnan(results["loss"])
+    assert results["loss"] > 0
+
+
+@pytest.fixture(scope="function")
+def model_with_asr():
+    """Model fixture with ASR head enabled."""
+    return create_model(predict_user_text=True)
+
+
+@pytest.fixture(scope="function")
+def model_with_noise():
+    """Model fixture with noise augmentation enabled."""
+    model = create_model(
+        force_use_noise_augmentation=True,
+        old_noise_prob=0.9,
+        old_noise_min_snr=5.0,
+        old_noise_max_snr=15.0,
+    )
+    return model
+
+
+@pytest.fixture(scope="function")
+def model_with_asr_and_noise():
+    """Model fixture with both ASR head and noise augmentation enabled."""
+    model = create_model(
+        predict_user_text=True,
+        force_use_noise_augmentation=True,
+        old_noise_prob=0.9,
+        old_noise_min_snr=5.0,
+        old_noise_max_snr=15.0,
+    )
+    return model
+
+
+def test_stt_training_step_with_asr(model_with_asr, dataset, training_cutset_batch):
+    model_with_asr.on_train_epoch_start()
+    batch = dataset[training_cutset_batch]
+    batch = move_data_to_device(batch, device=model_with_asr.device)
+    results = model_with_asr.training_step(batch, batch_idx=0)
+    assert torch.is_tensor(results["loss"])
+    assert not torch.isnan(results["loss"])
+    assert results["loss"] > 0
+
+    assert "asr_loss" in results
+    assert torch.is_tensor(results["asr_loss"])
+    assert not torch.isnan(results["asr_loss"])
+    assert results["asr_loss"] >= 0
+
+
+def test_stt_training_step_with_noise(model_with_asr_and_noise, dataset, training_cutset_batch):
+    model_with_asr_and_noise.on_train_epoch_start()
+    batch = dataset[training_cutset_batch]
+    batch = move_data_to_device(batch, device=model_with_asr_and_noise.device)
+    results = model_with_asr_and_noise.training_step(batch, batch_idx=0)
+    assert torch.is_tensor(results["loss"])
+    assert not torch.isnan(results["loss"])
+    assert results["loss"] > 0
+
+    assert "asr_loss" in results
+    assert torch.is_tensor(results["asr_loss"])
+    assert not torch.isnan(results["asr_loss"])
+    assert results["asr_loss"] >= 0
+
+
+def test_stt_validation_step(model, dataset, training_cutset_batch):
+    model.on_validation_epoch_start()
+    batch = dataset[training_cutset_batch]
+    batch = move_data_to_device(batch, device=model.device)
+    results = model.validation_step({"dummy_val_set": batch}, batch_idx=0)
+    assert results is None  # no return value
+
+
+def test_stt_offline_generation(model):
+    # 16000 samples == 1 second == 12.5 frames ~= 14 frames after encoder padding
+    ans = model.streaming_inference.offline_inference(
+        input_signal=torch.randn(1, 16000, device=model.device),
+        input_signal_lens=torch.tensor([16000], device=model.device),
+    )
+
+    assert ans.keys() == {
+        'text',
+        'src_text',
+        'tokens_text_src',
+        'tokens_text',
+        'tokens_len',
+        'source_audio',
+        'source_audio_len',
+    }
+
+    assert isinstance(ans["text"], list)
+    assert isinstance(ans["text"][0], str)
+
+    gen_text = ans["tokens_text"]
+    assert gen_text.shape == (1, 14)
+    assert gen_text.dtype == torch.long
+    assert (gen_text >= 0).all()
+    assert (gen_text < model.text_vocab_size).all()
+
+
+def test_stt_offline_generation_with_asr(model_with_asr):
+    """Test offline generation with ASR head enabled for user text prediction."""
+    # 16000 samples == 1 second == 12.5 frames ~= 14 frames after encoder padding
+    ans = model_with_asr.streaming_inference.offline_inference(
+        input_signal=torch.randn(1, 16000, device=model_with_asr.device),
+        input_signal_lens=torch.tensor([16000], device=model_with_asr.device),
+    )
+
+    # Verify all expected output keys are present
+    assert ans.keys() == {
+        'text',
+        'src_text',
+        'tokens_text_src',
+        'tokens_text',
+        'tokens_len',
+        'source_audio',
+        'source_audio_len',
+    }
+
+    # Verify agent text output
+    assert isinstance(ans["text"], list)
+    assert isinstance(ans["text"][0], str)
+
+    # Verify user text (ASR) output
+    assert isinstance(ans["src_text"], list)
+    assert isinstance(ans["src_text"][0], str)
+
+    # Verify generated text tokens
+    gen_text = ans["tokens_text"]
+    assert gen_text.shape == (1, 14)
+    assert gen_text.dtype == torch.long
+    assert (gen_text >= 0).all()
+    assert (gen_text < model_with_asr.text_vocab_size).all()
+
+    # Verify ASR tokens
+    asr_tokens = ans["tokens_text_src"]
+    assert asr_tokens.shape[0] == 1  # batch size
+    assert asr_tokens.dtype == torch.long
+    assert (asr_tokens >= 0).all()
+    assert (asr_tokens < model_with_asr.text_vocab_size).all()
+
+
+def test_trailing_pad_loss_scale_is_masked(dataset, training_cutset_batch):
+    """Test that trailing pad positions (from batching) have loss_scale=0 when token_loss_weight is set."""
+    model = create_model(predict_user_text=True)
+    # Enable token_loss_weight with non-zero pad weight
+    model.cfg["token_loss_weight"] = {"pad": 1.0, "bos": 10.0, "eos": 5.0, "text": 5.0}
+    model.cfg["mask_sequence_loss"] = True
+    if torch.cuda.is_available():
+        model.to("cuda")
+
+    batch = dataset[training_cutset_batch]
+    batch = move_data_to_device(batch, device=model.device)
+    inputs = model.prepare_inputs(batch["audio_data"])
+
+    loss_scale = inputs["loss_scale"]  # (B, T, 1)
+    asr_loss_scale = inputs["asr_loss_scale"]  # (B, T, 1)
+    seq_mask = inputs["seq_mask"]  # (B, T, 1)
+    target_token_lens = batch["audio_data"]["target_token_lens"]
+
+    for i in range(target_token_lens.size(0)):
+        end_idx = target_token_lens[i]
+        # Trailing positions (after target_token_lens) must have loss_scale=0
+        assert (loss_scale[i, end_idx:, :] == 0).all(), f"Batch {i}: loss_scale not zero after position {end_idx}"
+        assert (
+            asr_loss_scale[i, end_idx:, :] == 0
+        ).all(), f"Batch {i}: asr_loss_scale not zero after position {end_idx}"
+        # In-sequence positions should have non-zero loss_scale
+        assert (loss_scale[i, :end_idx, :] > 0).any(), f"Batch {i}: loss_scale all zero before position {end_idx}"

tests/collections/speechlm2/test_duplex_stt_dataset.py

@@ -0,0 +1,415 @@
+# Copyright (c) 2025, 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.
+
+import os
+import re
+
+import pytest
+import torch
+from lhotse import CutSet, SupervisionSegment, compute_num_frames
+from lhotse.dataset.collation import collate_audio, collate_vectors
+from lhotse.testing.dummies import dummy_cut, dummy_recording
+
+from nemo.collections.common.tokenizers import AutoTokenizer
+from nemo.collections.speechlm2.data.duplex_stt_dataset import (
+    DuplexSTTDataset,
+    collate_system_prompt,
+    collate_token_channel,
+)
+from nemo.collections.speechlm2.data.utils import get_pad_id
+
+SR = 16000
+FL = 0.08
+
+
+def _clean_text(text):
+    """Strip timestamp tokens and normalize whitespace, matching _text_to_ids(remove_timestamps=True)."""
+    text = re.sub(r'<\|\d+\|>', '', text)
+    return ' '.join(text.strip().split())
+
+
+def _verify_supervision_tokens(tokens_1d, start, duration, raw_text, tokenizer, pad, bos, eos, total_frames):
+    """Verify BOS/EOS placement and that decoded token IDs match the original supervision text.
+
+    Checks:
+        1. BOS is placed at the frame corresponding to supervision start.
+        2. EOS is placed at the frame corresponding to supervision end (if within cut).
+        3. Text token IDs between BOS and EOS are a prefix of tokenizer.text_to_ids(clean_text).
+        4. Decoding those IDs back to text matches (or is a prefix of) the original text.
+    """
+    pos = compute_num_frames(start, FL, SR)
+    eospos = compute_num_frames(start + duration, FL, SR)
+    clean = _clean_text(raw_text)
+
+    # 1. BOS at turn start
+    assert tokens_1d[pos].item() == bos, f"Expected BOS={bos} at frame {pos}, got {tokens_1d[pos].item()}"
+
+    # 2. EOS at turn end (only if within cut bounds)
+    if eospos < total_frames:
+        assert tokens_1d[eospos].item() == eos, f"Expected EOS={eos} at frame {eospos}, got {tokens_1d[eospos].item()}"
+
+    # 3. Extract text token IDs between BOS and EOS, filtering out pad
+    end = min(eospos, total_frames)
+    actual_ids = [t for t in tokens_1d[pos + 1 : end].tolist() if t != pad]
+    expected_ids = tokenizer.text_to_ids(clean)
+
+    # Actual IDs should be a prefix of expected (truncation may occur when text is long)
+    assert actual_ids == expected_ids[: len(actual_ids)], (
+        f"Token ID mismatch for '{clean}':\n"
+        f"  actual_ids      = {actual_ids}\n"
+        f"  expected_prefix = {expected_ids[: len(actual_ids)]}\n"
+        f"  full_expected   = {expected_ids}"
+    )
+
+    # 4. Decode IDs back to text and verify against original
+    if actual_ids:
+        decoded = tokenizer.ids_to_text(actual_ids).strip()
+        if len(actual_ids) == len(expected_ids):
+            assert decoded == clean, f"Full decode mismatch: '{decoded}' != '{clean}', ids={actual_ids}"
+        else:
+            assert clean.startswith(decoded), (
+                f"Truncated decode '{decoded}' is not a prefix of '{clean}'\n"
+                f"  ids={actual_ids} (truncated {len(expected_ids)} → {len(actual_ids)})"
+            )
+
+
+@pytest.fixture(scope="session")
+def tokenizer():
+    if os.path.exists("/home/TestData/speechlm/pretrained_models"):
+        model_path = "/home/TestData/speechlm/pretrained_models/TinyLlama--TinyLlama_v1.1"
+    else:
+        model_path = "TinyLlama/TinyLlama_v1.1"
+    return AutoTokenizer(model_path, use_fast=True)
+
+
+@pytest.fixture(scope="session")
+def cuts():
+    """Two cuts: cut1 with plain text, cut2 with timestamped text and a system prompt."""
+    cut1 = dummy_cut(0, duration=1.0, recording=dummy_recording(0, duration=1.0, with_data=True))
+    cut1.supervisions = [
+        SupervisionSegment(
+            id="s0-user", recording_id=cut1.recording_id, start=0, duration=0.3, text="hi", speaker="user"
+        ),
+        SupervisionSegment(
+            id="s0-agent", recording_id=cut1.recording_id, start=0.4, duration=0.3, text="hello", speaker="assistant"
+        ),
+    ]
+
+    cut2 = dummy_cut(1, duration=2.0, recording=dummy_recording(1, duration=2.0, with_data=True))
+    cut2.supervisions = [
+        SupervisionSegment(
+            id="s1-user",
+            recording_id=cut2.recording_id,
+            start=0,
+            duration=0.5,
+            text="<|0|> good <|1|> <|3|> morning <|5|>",
+            speaker="user",
+        ),
+        SupervisionSegment(
+            id="s1-agent",
+            recording_id=cut2.recording_id,
+            start=0.6,
+            duration=0.5,
+            text="<|0|> good <|2|> <|2|> morning <|4|> <|4|> to <|5|> <|5|> you <|6|>",
+            speaker="assistant",
+        ),
+        SupervisionSegment(
+            id="s1-user2",
+            recording_id=cut2.recording_id,
+            start=1.2,
+            duration=0.3,
+            text="<|0|> thanks <|3|>",
+            speaker="user",
+        ),
+        SupervisionSegment(
+            id="s1-agent2",
+            recording_id=cut2.recording_id,
+            start=1.6,
+            duration=0.4,
+            text="<|0|> welcome <|4|>",
+            speaker="assistant",
+        ),
+    ]
+    cut2.custom = {"system_prompt": "be helpful"}
+
+    return CutSet([cut1, cut2])
+
+
+def test_collate_audio(cuts):
+    """Test collate_audio: shapes, lengths, and zero-padding for shorter cuts."""
+    audio, audio_lens = collate_audio(cuts.resample(SR))
+
+    assert audio.shape == (2, 32000)
+    assert audio_lens.tolist() == [16000, 32000]
+    # Padding region for the shorter cut must be zero
+    assert (audio[0, 16000:] == 0).all(), "Audio padding should be zero"
+    # Non-padding region should have non-zero data (random audio from dummy_recording)
+    assert (audio[0, :16000] != 0).any(), "Audio data should be non-zero"
+    assert (audio[1, :32000] != 0).any(), "Audio data should be non-zero"
+
+
+def test_collate_token_channel_target(cuts, tokenizer):
+    """Test collate_token_channel for target (assistant) role: BOS/EOS placement, token decode."""
+    pad = get_pad_id(tokenizer)
+    bos = tokenizer.bos
+    eos = tokenizer.eos
+    total1 = compute_num_frames(1.0, FL, SR)  # 13
+    total2 = compute_num_frames(2.0, FL, SR)  # 25
+
+    target_tokens, target_token_lens = collate_token_channel(
+        cuts,
+        tokenizer,
+        frame_length=FL,
+        roles={"assistant"},
+        bos_id=bos,
+        eos_id=eos,
+        remove_timestamps=True,
+    )
+
+    assert target_token_lens.tolist() == [total1, total2]
+
+    # fmt: off
+    # Cut 1: "hello"(22172) at frames 5–9, padded to 25
+    # Cut 2: "good morning to you"(1781,7250,304,366) at frames 8–14, "welcome"(12853) at frames 20–24
+    expected_target = torch.tensor([
+        [0, 0, 0, 0, 0, 1, 22172, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+        [0, 0, 0, 0, 0, 0, 0, 0, 1, 1781, 7250, 304, 366, 0, 2, 0, 0, 0, 0, 0, 1, 12853, 0, 0, 0],
+    ])
+    # fmt: on
+    assert torch.equal(target_tokens, expected_target)
+
+    # Decode verification
+    _verify_supervision_tokens(target_tokens[0], 0.4, 0.3, "hello", tokenizer, pad, bos, eos, total1)
+    _verify_supervision_tokens(
+        target_tokens[1],
+        0.6,
+        0.5,
+        "<|0|> good <|2|> <|2|> morning <|4|> <|4|> to <|5|> <|5|> you <|6|>",
+        tokenizer,
+        pad,
+        bos,
+        eos,
+        total2,
+    )
+    _verify_supervision_tokens(target_tokens[1], 1.6, 0.4, "<|0|> welcome <|4|>", tokenizer, pad, bos, eos, total2)
+
+
+def test_collate_token_channel_source(cuts, tokenizer):
+    """Test collate_token_channel for source (user) role with timestamp-based word alignment."""
+    pad = get_pad_id(tokenizer)
+    bos = tokenizer.bos
+    eos = tokenizer.eos
+    total1 = compute_num_frames(1.0, FL, SR)  # 13
+    total2 = compute_num_frames(2.0, FL, SR)  # 25
+
+    source_tokens, source_token_lens = collate_token_channel(
+        cuts,
+        tokenizer,
+        frame_length=FL,
+        roles={"user"},
+        bos_id=bos,
+        eos_id=eos,
+        remove_timestamps=False,
+        prepend_word_space=False,
+    )
+
+    assert source_tokens.shape == (2, total2)
+    assert source_token_lens.tolist() == [total1, total2]
+
+    # fmt: off
+    # Cut 1: "hi"(7251) plain text, placed contiguously at frames 0–4
+    # Cut 2: "good"(1781) at ts 0–1, pad gap, "morning"(7250) at ts 3–5 → frames 0–6
+    #         "thanks"(3969) at ts 0–3 → frames 15–19
+    expected_source = torch.tensor([
+        [1, 7251, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+        [1, 1781, 0, 0, 7250, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0, 1, 3969, 0, 0, 2, 0, 0, 0, 0, 0],
+    ])
+    # fmt: on
+    assert torch.equal(source_tokens, expected_source)
+
+    # Decode verification
+    # ── Cut 1: no timestamps, sentence-level tokenization ──
+    _verify_supervision_tokens(source_tokens[0], 0.0, 0.3, "hi", tokenizer, pad, bos, eos, total1)
+    assert (source_tokens[0, total1:] == pad).all(), "Batch padding should be pad"
+
+    # ── Cut 2: two user turns with timestamp alignment ──
+    _verify_supervision_tokens(
+        source_tokens[1],
+        0.0,
+        0.5,
+        "<|0|> good <|1|> <|3|> morning <|5|>",
+        tokenizer,
+        pad,
+        bos,
+        eos,
+        total2,
+    )
+    _verify_supervision_tokens(source_tokens[1], 1.2, 0.3, "<|0|> thanks <|3|>", tokenizer, pad, bos, eos, total2)
+
+
+def test_collate_system_prompt(cuts, tokenizer):
+    """Test collate_system_prompt: cut1 has no prompt, cut2 has 'be helpful'."""
+
+    prompt_tokens, prompt_token_lens = collate_system_prompt(cuts, tokenizer)
+
+    # fmt: off
+    # cut1: no system_prompt → all pad, len=0
+    # cut2: "be helpful" → [BOS=1, "be"=367, "helpful"=8444, EOS=2], len=4
+    expected_prompt = torch.tensor([
+        [0, 0, 0, 0],
+        [1, 367, 8444, 2],
+    ])
+    # fmt: on
+    assert prompt_token_lens.tolist() == [0, 4]
+    assert torch.equal(prompt_tokens, expected_prompt)
+
+    # Decode prompt tokens back to text
+    prompt_ids = tokenizer.text_to_ids("be helpful")
+    decoded_prompt = tokenizer.ids_to_text(prompt_ids).strip()
+    assert decoded_prompt == "be helpful", f"Prompt decode: '{decoded_prompt}'"
+
+
+def test_collate_text_data(tokenizer):
+    """Test collate_vectors for text token inputs: padding and lengths."""
+    pad = get_pad_id(tokenizer)
+
+    # "hi" → [7251], "good morning" → [1781, 7250]
+    text_tokens_list = [
+        torch.tensor([7251], dtype=torch.long),
+        torch.tensor([1781, 7250], dtype=torch.long),
+    ]
+    text_token_lens = torch.tensor([t.shape[0] for t in text_tokens_list], dtype=torch.long)
+    text_tokens = collate_vectors(text_tokens_list, padding_value=pad)
+
+    assert text_token_lens.tolist() == [1, 2]
+    assert text_tokens.shape == (2, 2)
+    # Shorter sequence is right-padded
+    assert text_tokens[0].tolist() == [7251, pad]
+    assert text_tokens[1].tolist() == [1781, 7250]
+    # Decode back to verify
+    assert tokenizer.ids_to_text([7251]).strip() == "hi"
+    assert tokenizer.ids_to_text([1781, 7250]).strip() == "good morning"
+
+
+def test_duplex_stt_dataset(cuts, tokenizer):
+    """End-to-end test of DuplexSTTDataset.__getitem__: covers all collate outputs including timestamps."""
+    dataset = DuplexSTTDataset(
+        tokenizer=tokenizer,
+        frame_length=FL,
+        source_sample_rate=SR,
+        input_roles=["user"],
+        output_roles=["assistant"],
+        cfg={"prepend_word_space": False},
+        model_cfg={"predict_user_text": True},
+    )
+    batch = dataset[cuts]
+    ad = batch["audio_data"]
+
+    total1 = compute_num_frames(1.0, FL, SR)  # 13
+    total2 = compute_num_frames(2.0, FL, SR)  # 25
+
+    # sample_id
+    assert len(ad["sample_id"]) == 2
+
+    # source_audio
+    assert ad["source_audio"].shape == (2, 32000)
+    assert ad["source_audio_lens"].tolist() == [16000, 32000]
+
+    # target_tokens (remove_timestamps=True, same as unit test)
+    assert ad["target_token_lens"].tolist() == [total1, total2]
+    # fmt: off
+    assert torch.equal(ad["target_tokens"], torch.tensor([
+        [0, 0, 0, 0, 0, 1, 22172, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+        [0, 0, 0, 0, 0, 0, 0, 0, 1, 1781, 7250, 304, 366, 0, 2, 0, 0, 0, 0, 0, 1, 12853, 0, 0, 0],
+    ]))
+    # fmt: on
+
+    # source_tokens (remove_timestamps=False via predict_user_text=True, timestamp-aligned)
+    assert ad["source_token_lens"].tolist() == [total1, total2]
+    # fmt: off
+    assert torch.equal(ad["source_tokens"], torch.tensor([
+        [1, 7251, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+        [1, 1781, 0, 0, 7250, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0, 1, 3969, 0, 0, 2, 0, 0, 0, 0, 0],
+    ]))
+    # fmt: on
+
+    # system prompt (cut2 has "be helpful")
+    assert "prompt_tokens" in ad
+    assert ad["prompt_token_lens"].tolist() == [0, 4]
+    # fmt: off
+    assert torch.equal(ad["prompt_tokens"], torch.tensor([
+        [0, 0, 0, 0],
+        [1, 367, 8444, 2],
+    ]))
+    # fmt: on
+
+    # source/target texts
+    assert ad["source_texts"] == ["hi", "good morning thanks"]
+    assert ad["target_texts"] == [
+        "hello",
+        "<|0|> good <|2|> <|2|> morning <|4|> <|4|> to <|5|> <|5|> you <|6|> <|0|> welcome <|4|>",
+    ]
+
+    # task
+    assert ad["task"] == ["s2s_duplex", "s2s_duplex"]
+
+    # no text data (no Formattable cuts)
+    assert batch["text_data"] is None
+
+    # source_audio is not augmented when augmenter is not configured (audio is unchanged from collate_audio)
+
+
+def test_duplex_stt_dataset_augmentation(cuts, tokenizer, tmp_path):
+    """Test that audio augmentation modifies source_audio in-place when configured."""
+    import numpy as np
+    import soundfile as sf
+
+    # Create dummy noise files for the augmenter
+    noise_dir = tmp_path / "noise" / "all"
+    noise_dir.mkdir(parents=True)
+    for i in range(3):
+        noise = np.random.randn(SR).astype(np.float32) * 0.01
+        sf.write(str(noise_dir / f"noise_{i}.wav"), noise, SR)
+
+    cfg = {
+        "prepend_word_space": False,
+        "use_noise_aug": True,
+        "noise_prob": 1.0,
+        "noise_aug_path": str(tmp_path / "noise"),
+        "noise_min_snr": 20,
+        "noise_max_snr": 20,
+    }
+
+    dataset = DuplexSTTDataset(
+        tokenizer=tokenizer,
+        frame_length=FL,
+        source_sample_rate=SR,
+        input_roles=["user"],
+        output_roles=["assistant"],
+        cfg=cfg,
+        model_cfg={"predict_user_text": True},
+    )
+
+    assert dataset.audio_augmenter is not None
+
+    # Get original audio for comparison
+    original_audio, _ = collate_audio(cuts.resample(SR))
+
+    batch = dataset[cuts]
+    ad = batch["audio_data"]
+
+    # source_audio should be augmented (different from original)
+    assert "source_audio_aug" not in ad
+    assert ad["source_audio"].shape == original_audio.shape
+    assert not torch.equal(ad["source_audio"], original_audio)

tests/collections/speechlm2/test_early_interruption.py

@@ -0,0 +1,339 @@
+# Copyright (c) 2025, 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 unittest.mock import MagicMock, PropertyMock
+
+import pytest
+import torch
+
+from nemo.collections.speechlm2.data import DuplexSTTDataset
+
+
+@pytest.fixture
+def mock_tokenizer():
+    """Create a mock tokenizer for testing."""
+    tokenizer = MagicMock()
+    type(tokenizer).bos = PropertyMock(return_value=1)
+    type(tokenizer).eos = PropertyMock(return_value=2)
+    type(tokenizer).pad = PropertyMock(return_value=0)
+    type(tokenizer).pad_id = PropertyMock(return_value=0)
+    type(tokenizer).unk_id = PropertyMock(return_value=None)
+    tokenizer.text_to_ids = MagicMock(return_value=[1])
+    return tokenizer
+
+
+@pytest.fixture
+def dataset_with_early_interruption(mock_tokenizer):
+    """Create a dataset with early interruption enabled."""
+    cfg = {"early_interruption_prob": 1.0, "early_interruption_overlap_tokens": 5}
+    model_cfg = {"predict_user_text": False, "force_align_user_text": False}
+
+    dataset = DuplexSTTDataset(
+        tokenizer=mock_tokenizer,
+        frame_length=0.08,
+        source_sample_rate=16000,
+        input_roles=["user"],
+        output_roles=["assistant"],
+        cfg=cfg,
+        model_cfg=model_cfg,
+    )
+    return dataset
+
+
+def test_early_interruption_basic_truncation(dataset_with_early_interruption):
+    """Test that early interruption truncates an agent turn correctly."""
+    # Setup: Create mock tensors with realistic structure:
+    # BOS, non-pad tokens, pad tokens (silence), then EOS
+    # Early interruption should move EOS into the non-pad token region
+    batch_size = 1
+    seq_len = 25
+
+    target_tokens = torch.full((batch_size, seq_len), 0, dtype=torch.long)
+    target_tokens[0, 0] = 1  # BOS
+    target_tokens[0, 1:13] = torch.arange(10, 22)  # 12 content tokens (positions 1-12)
+    # Positions 13-17 are PAD (0) - representing silence/gap
+    target_tokens[0, 18] = 2  # Original EOS at position 18 (after the padding)
+
+    source_tokens = torch.full((batch_size, seq_len), 0, dtype=torch.long)
+
+    target_audio = torch.zeros((batch_size, 20000), dtype=torch.float32)
+    source_audio = torch.zeros((batch_size, 16000), dtype=torch.float32)
+
+    target_token_lens = torch.tensor([19], dtype=torch.long)
+    source_token_lens = torch.tensor([1], dtype=torch.long)
+    target_audio_lens = torch.tensor([20000], dtype=torch.long)
+    source_audio_lens = torch.tensor([16000], dtype=torch.long)
+
+    # Apply early interruption
+    dataset_with_early_interruption._apply_early_interruption_augmentation(
+        target_tokens=target_tokens,
+        source_tokens=source_tokens,
+        source_audio=source_audio,
+        source_audio_lens=source_audio_lens,
+        batch_idx=0,
+    )
+
+    # Verify: EOS should be moved earlier (within the non-pad token region)
+    eos_positions = (target_tokens[0] == 2).nonzero(as_tuple=True)[0]
+    assert len(eos_positions) > 0, "EOS should still exist after early interruption"
+
+    new_eos_pos = eos_positions[0].item()
+    # The new EOS position should be before the original position (18)
+    assert new_eos_pos < 18, f"New EOS position {new_eos_pos} should be before original position 18"
+
+    # CRITICAL: New EOS should be at cutoff_pos + overlap_tokens
+    # cutoff_pos is in non-pad region (1-12), overlap_tokens is 5
+    # So new_eos_pos should be in range (1+5) to (12+5) = 6 to 17
+    overlap_tokens = 5  # from fixture cfg
+    assert 6 <= new_eos_pos <= 17, (
+        f"New EOS at {new_eos_pos} should be within cutoff + overlap range (6-17), "
+        f"meaning agent continues for {overlap_tokens} tokens after user interruption"
+    )
+
+    # Check that tokens after new EOS are shifted or padded
+    assert target_tokens[0, -1] == 0, "Last token should be PAD after early interruption"
+
+
+def test_early_interruption_with_multiple_turns(dataset_with_early_interruption):
+    """Test early interruption with multiple agent turns."""
+    batch_size = 1
+    seq_len = 50
+
+    target_tokens = torch.full((batch_size, seq_len), 0, dtype=torch.long)
+    # First turn: BOS at 0, tokens 1-10, PAD 11-13, EOS at 14
+    target_tokens[0, 0] = 1
+    target_tokens[0, 1:11] = torch.arange(10, 20)  # 10 content tokens
+    # Positions 11-13 are PAD
+    target_tokens[0, 14] = 2  # EOS after padding
+
+    # Second turn: BOS at 18, tokens 19-28, PAD 29-31, EOS at 32
+    target_tokens[0, 18] = 1
+    target_tokens[0, 19:29] = torch.arange(20, 30)  # 10 content tokens
+    # Positions 29-31 are PAD
+    target_tokens[0, 32] = 2  # EOS after padding
+
+    source_tokens = torch.full((batch_size, seq_len), 0, dtype=torch.long)
+    target_audio = torch.zeros((batch_size, 40000), dtype=torch.float32)
+    source_audio = torch.zeros((batch_size, 16000), dtype=torch.float32)
+
+    target_token_lens = torch.tensor([33], dtype=torch.long)
+    source_token_lens = torch.tensor([1], dtype=torch.long)
+    target_audio_lens = torch.tensor([40000], dtype=torch.long)
+    source_audio_lens = torch.tensor([16000], dtype=torch.long)
+
+    # Apply early interruption
+    dataset_with_early_interruption._apply_early_interruption_augmentation(
+        target_tokens=target_tokens,
+        source_tokens=source_tokens,
+        source_audio=source_audio,
+        source_audio_lens=source_audio_lens,
+        batch_idx=0,
+    )
+
+    # Verify: Should still have valid BOS and EOS tokens
+    bos_positions = (target_tokens[0] == 1).nonzero(as_tuple=True)[0]
+    eos_positions = (target_tokens[0] == 2).nonzero(as_tuple=True)[0]
+
+    assert len(bos_positions) >= 1, "Should have at least one BOS token"
+    assert len(eos_positions) >= 1, "Should have at least one EOS token"
+
+    # Each BOS should be followed eventually by an EOS
+    for bos_pos in bos_positions:
+        matching_eos = eos_positions[eos_positions > bos_pos]
+        assert len(matching_eos) > 0, f"BOS at position {bos_pos} should have a matching EOS"
+
+
+def test_early_interruption_overlap_tokens(dataset_with_early_interruption):
+    """Test that overlap tokens parameter works correctly."""
+    # Test with custom overlap tokens
+    dataset_with_early_interruption.cfg["early_interruption_overlap_tokens"] = 3
+
+    batch_size = 1
+    seq_len = 25
+
+    target_tokens = torch.full((batch_size, seq_len), 0, dtype=torch.long)
+    target_tokens[0, 0] = 1  # BOS
+    target_tokens[0, 1:11] = torch.arange(10, 20)  # 10 content tokens (positions 1-10)
+    # Positions 11-14 are PAD
+    target_tokens[0, 15] = 2  # Original EOS at position 15 (after padding)
+    original_eos_pos = 15
+    overlap_tokens = 3
+
+    # Place a marker token in source_tokens AFTER original_eos_pos to track the shift
+    # The implementation shifts source_tokens from original_eos_pos+1 to cutoff_pos+1
+    marker_token = 999
+    marker_original_pos = 20  # Position after original EOS
+    source_tokens = torch.full((batch_size, seq_len), 0, dtype=torch.long)
+    source_tokens[0, marker_original_pos] = marker_token
+
+    target_audio = torch.zeros((batch_size, 20000), dtype=torch.float32)
+    source_audio = torch.zeros((batch_size, 16000), dtype=torch.float32)
+
+    target_token_lens = torch.tensor([16], dtype=torch.long)
+    source_token_lens = torch.tensor([1], dtype=torch.long)
+    target_audio_lens = torch.tensor([20000], dtype=torch.long)
+    source_audio_lens = torch.tensor([16000], dtype=torch.long)
+
+    # Apply early interruption
+    dataset_with_early_interruption._apply_early_interruption_augmentation(
+        target_tokens=target_tokens,
+        source_tokens=source_tokens,
+        source_audio=source_audio,
+        source_audio_lens=source_audio_lens,
+        batch_idx=0,
+    )
+
+    # Find new EOS position
+    eos_positions = (target_tokens[0] == 2).nonzero(as_tuple=True)[0]
+    assert len(eos_positions) > 0, "EOS should exist after early interruption"
+    new_eos_pos = eos_positions[0].item()
+
+    # Find where the marker moved to in source_tokens
+    # The shift is: source_tokens[cutoff_pos+1:...] = source_tokens[original_eos_pos+1:...]
+    # So marker moves from marker_original_pos to cutoff_pos + 1 + (marker_original_pos - original_eos_pos - 1)
+    # = cutoff_pos + (marker_original_pos - original_eos_pos)
+    marker_new_positions = (source_tokens[0] == marker_token).nonzero(as_tuple=True)[0]
+    assert len(marker_new_positions) > 0, "Marker token should still exist after transformation"
+    marker_new_pos = marker_new_positions[0].item()
+
+    # Calculate actual cutoff position from the marker shift
+    # marker_new_pos = cutoff_pos + (marker_original_pos - original_eos_pos)
+    # => cutoff_pos = marker_new_pos - (marker_original_pos - original_eos_pos)
+    actual_cutoff_pos = marker_new_pos - (marker_original_pos - original_eos_pos)
+
+    # Verify the overlap is exactly overlap_tokens
+    actual_overlap = new_eos_pos - actual_cutoff_pos
+    assert actual_overlap == overlap_tokens, (
+        f"Overlap should be {overlap_tokens} tokens, but got {actual_overlap} "
+        f"(cutoff at {actual_cutoff_pos}, new EOS at {new_eos_pos})"
+    )
+
+    # Cutoff should be in non-pad region (user interrupts during non-pad region)
+    assert 1 <= actual_cutoff_pos <= 10, (
+        f"Cutoff at {actual_cutoff_pos} should be in non-pad region (1-10), "
+        f"meaning user interrupts during active speech"
+    )
+
+    print(f"\n✓ Overlap tokens verification:")
+    print(f"  - Configured overlap: {overlap_tokens} tokens")
+    print(f"  - Actual cutoff position (from marker shift): {actual_cutoff_pos}")
+    print(f"  - New EOS position: {new_eos_pos}")
+    print(f"  - Actual overlap: {actual_overlap} tokens ✓")
+
+
+def test_early_interruption_no_valid_turns():
+    """Test that early interruption handles cases with no valid turns gracefully."""
+    mock_tokenizer = MagicMock()
+    type(mock_tokenizer).bos = PropertyMock(return_value=1)
+    type(mock_tokenizer).eos = PropertyMock(return_value=2)
+    type(mock_tokenizer).pad = PropertyMock(return_value=0)
+    type(mock_tokenizer).pad_id = PropertyMock(return_value=0)
+    type(mock_tokenizer).unk_id = PropertyMock(return_value=None)
+    mock_tokenizer.text_to_ids = MagicMock(return_value=[1])
+
+    cfg = {"early_interruption_prob": 1.0, "early_interruption_overlap_tokens": 5}
+    model_cfg = {"predict_user_text": False, "force_align_user_text": False}
+
+    dataset = DuplexSTTDataset(
+        tokenizer=mock_tokenizer,
+        frame_length=0.08,
+        source_sample_rate=16000,
+        input_roles=["user"],
+        output_roles=["assistant"],
+        cfg=cfg,
+        model_cfg=model_cfg,
+    )
+
+    # Create tensors with no valid turns (all padding)
+    batch_size = 1
+    seq_len = 20
+
+    target_tokens = torch.full((batch_size, seq_len), 0, dtype=torch.long)
+    source_tokens = torch.full((batch_size, seq_len), 0, dtype=torch.long)
+    target_audio = torch.zeros((batch_size, 16000), dtype=torch.float32)
+    source_audio = torch.zeros((batch_size, 16000), dtype=torch.float32)
+
+    target_token_lens = torch.tensor([1], dtype=torch.long)
+    source_token_lens = torch.tensor([1], dtype=torch.long)
+    target_audio_lens = torch.tensor([16000], dtype=torch.long)
+    source_audio_lens = torch.tensor([16000], dtype=torch.long)
+
+    # Apply early interruption - should not crash
+    dataset._apply_early_interruption_augmentation(
+        target_tokens=target_tokens,
+        source_tokens=source_tokens,
+        source_audio=source_audio,
+        source_audio_lens=source_audio_lens,
+        batch_idx=0,
+    )
+
+    # Verify: Tokens should remain unchanged (all zeros)
+    assert torch.all(target_tokens == 0), "Tokens should remain unchanged when no valid turns exist"
+
+
+def test_early_interruption_frames_to_remove_calculation(dataset_with_early_interruption):
+    """Test that frames_to_remove is calculated correctly."""
+    batch_size = 1
+    seq_len = 25
+
+    target_tokens = torch.full((batch_size, seq_len), 0, dtype=torch.long)
+    target_tokens[0, 0] = 1  # BOS
+    target_tokens[0, 1:11] = torch.arange(10, 20)  # 10 content tokens (positions 1-10)
+    # Positions 11-14 are PAD
+    target_tokens[0, 15] = 2  # Original EOS at position 15 (after padding)
+    original_eos_pos = 15
+
+    source_tokens = torch.full((batch_size, seq_len), 0, dtype=torch.long)
+    target_audio = torch.zeros((batch_size, 20000), dtype=torch.float32)
+    source_audio = torch.zeros((batch_size, 16000), dtype=torch.float32)
+
+    target_token_lens = torch.tensor([16], dtype=torch.long)
+    source_token_lens = torch.tensor([1], dtype=torch.long)
+    target_audio_lens = torch.tensor([20000], dtype=torch.long)
+    source_audio_lens = torch.tensor([16000], dtype=torch.long)
+
+    # Apply early interruption
+    dataset_with_early_interruption._apply_early_interruption_augmentation(
+        target_tokens=target_tokens,
+        source_tokens=source_tokens,
+        source_audio=source_audio,
+        source_audio_lens=source_audio_lens,
+        batch_idx=0,
+    )
+
+    # Verify: New EOS should be in the non-pad region
+    new_eos_positions = (target_tokens[0] == 2).nonzero(as_tuple=True)[0]
+    assert len(new_eos_positions) > 0, "EOS should exist after early interruption"
+    new_eos_pos = new_eos_positions[0].item()
+
+    # cutoff_pos must satisfy (eos_pos - pos) > overlap_tokens, i.e., pos < 10
+    # So valid cutoff positions are 1-9, and new_eos_pos = cutoff_pos + 5
+    # Range: (1+5) to (9+5) = 6 to 14
+    overlap_tokens = 5  # from fixture cfg
+    assert 6 <= new_eos_pos <= 14, (
+        f"New EOS at {new_eos_pos} should be within cutoff + overlap range (6-14), "
+        f"meaning agent continues for {overlap_tokens} tokens after user interruption"
+    )
+
+    # Verify: Check that padding is added at the end
+    # Since we always use frames_to_remove = original_eos_pos - cutoff_pos,
+    # we should have more padding at the end after truncation
+    num_pad_tokens = (target_tokens[0] == 0).sum().item()
+    # 9 is a conservative lower bound: new_eos_pos ranges 6-14, so minimum trailing
+    # PAD is 10 (positions 15-24), plus PAD positions 11-13 before EOS = 13 total minimum
+    assert num_pad_tokens >= 9, "Should have increased padding after truncation"
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

tests/collections/speechlm2/test_encoder_chunking.py

@@ -0,0 +1,229 @@
+# 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.
+import pytest
+import torch
+
+from nemo.collections.speechlm2.data.salm_dataset import MultiSpeakerConfig
+from nemo.collections.speechlm2.parts.encoder_chunking import (
+    _recombine_chunked_audio_embeddings,
+    _split_audio_into_chunks,
+    _split_spk_targets_into_chunks,
+    encode_audio_with_optional_chunking,
+)
+from tests.collections.speechlm2._chunking_helpers import ChunkingTestPerception
+
+
+@pytest.mark.parametrize(
+    (
+        "input_signal_lengths",
+        "chunk_size_samples",
+        "min_chunk_size_samples",
+        "expected_chunk_lens",
+        "expected_chunks_per_audio",
+        "expected_chunk_spans",
+    ),
+    [
+        (
+            [5, 0, 7],
+            3,
+            2,
+            [3, 2, 0, 3, 4],
+            [2, 1, 2],
+            [(0, 0, 3), (0, 3, 5), (1, 0, 0), (2, 0, 3), (2, 3, 7)],
+        ),
+        (
+            [6],
+            2,
+            1,
+            [2, 2, 2],
+            [3],
+            [(0, 0, 2), (0, 2, 4), (0, 4, 6)],
+        ),
+    ],
+)
+def test_split_audio_into_chunks_returns_spans_independent_of_spk_targets(
+    input_signal_lengths,
+    chunk_size_samples,
+    min_chunk_size_samples,
+    expected_chunk_lens,
+    expected_chunks_per_audio,
+    expected_chunk_spans,
+):
+    max_signal_len = max(input_signal_lengths, default=0)
+    input_signal = torch.arange(len(input_signal_lengths) * max_signal_len, dtype=torch.float32).reshape(
+        len(input_signal_lengths), max_signal_len
+    )
+    chunks, chunk_lens, chunks_per_audio, chunk_spans = _split_audio_into_chunks(
+        input_signal=input_signal,
+        input_signal_lengths=input_signal_lengths,
+        chunk_size_samples=chunk_size_samples,
+        min_chunk_size_samples=min_chunk_size_samples,
+    )
+
+    assert chunk_lens == expected_chunk_lens
+    assert chunks_per_audio == expected_chunks_per_audio
+    assert chunk_spans == expected_chunk_spans
+    for chunk, (audio_idx, begin, end) in zip(chunks, chunk_spans):
+        assert torch.equal(chunk, input_signal[audio_idx, begin:end])
+
+    assert _split_spk_targets_into_chunks(None, input_signal_lengths, chunk_spans) is None
+
+
+@pytest.mark.parametrize(
+    ("chunk_values", "chunk_lens", "chunks_per_audio", "expected_audio_values"),
+    [
+        (
+            [[1.0, 2.0, 0.0], [3.0, 4.0, 5.0]],
+            [2, 3],
+            [2],
+            [[1.0, 2.0, 3.0, 4.0, 5.0]],
+        ),
+        (
+            [[1.0, 2.0, 0.0], [3.0, 4.0, 5.0], [10.0, 11.0, 0.0], [12.0, 13.0, 14.0]],
+            [2, 3, 2, 3],
+            [2, 2],
+            [[1.0, 2.0, 3.0, 4.0, 5.0], [10.0, 11.0, 12.0, 13.0, 14.0]],
+        ),
+    ],
+)
+def test_recombine_chunked_audio_embeddings_reconstructs_original_rows(
+    chunk_values,
+    chunk_lens,
+    chunks_per_audio,
+    expected_audio_values,
+):
+    chunked_embs = torch.tensor(chunk_values, dtype=torch.float32).unsqueeze(-1)
+    chunked_emb_lens = torch.tensor(chunk_lens, dtype=torch.long)
+
+    audio_embs = _recombine_chunked_audio_embeddings(chunked_embs, chunked_emb_lens, chunks_per_audio)
+
+    assert len(audio_embs) == len(expected_audio_values)
+    for audio_emb, expected_values in zip(audio_embs, expected_audio_values):
+        assert torch.equal(audio_emb.squeeze(-1), torch.tensor(expected_values))
+
+
+@pytest.mark.parametrize(
+    (
+        "input_signal_lengths",
+        "chunk_spans",
+        "expected_chunks",
+    ),
+    [
+        (
+            [5],
+            [(0, 0, 2), (0, 2, 5)],
+            [
+                [[0.0, 1.0, 2.0, 3.0], [4.0, 5.0, 6.0, 7.0], [4.0, 5.0, 6.0, 7.0]],
+                [[8.0, 9.0, 10.0, 11.0], [12.0, 13.0, 14.0, 15.0], [16.0, 17.0, 18.0, 19.0]],
+            ],
+        ),
+        (
+            [4],
+            [(0, 0, 2), (0, 2, 4)],
+            [
+                [[0.0, 1.0, 2.0, 3.0], [4.0, 5.0, 6.0, 7.0], [4.0, 5.0, 6.0, 7.0]],
+                [[8.0, 9.0, 10.0, 11.0], [12.0, 13.0, 14.0, 15.0], [16.0, 17.0, 18.0, 19.0]],
+            ],
+        ),
+    ],
+)
+def test_split_spk_targets_into_chunks_uses_chunk_spans(
+    input_signal_lengths,
+    chunk_spans,
+    expected_chunks,
+):
+    cfg = MultiSpeakerConfig()
+    spk_targets = torch.arange(5 * cfg.num_speakers, dtype=torch.float32).reshape(1, 5, cfg.num_speakers)
+
+    chunked_spk_targets = _split_spk_targets_into_chunks(spk_targets, input_signal_lengths, chunk_spans)
+
+    assert torch.equal(chunked_spk_targets, torch.tensor(expected_chunks))
+
+
+@pytest.mark.parametrize(
+    ("audio_values", "audio_len", "expected_chunk_lens"),
+    [
+        ([1.0, 2.0, 3.0, 4.0, 5.0], 5, [2, 3]),
+        ([1.0, 2.0, 3.0, 4.0], 4, [2, 2]),
+    ],
+)
+def test_encode_audio_with_optional_chunking_does_not_forward_absent_spk_targets(
+    audio_values, audio_len, expected_chunk_lens
+):
+    perception = ChunkingTestPerception(sampling_rate=2, hop_length=1)
+    audios = torch.tensor([audio_values])
+    audio_lens = torch.tensor([audio_len], dtype=torch.long)
+
+    embs = encode_audio_with_optional_chunking(
+        perception,
+        audios,
+        audio_lens,
+        chunk_size_seconds=1.0,
+        sampling_rate=2,
+        spk_targets=None,
+    )
+
+    chunked_signal, chunked_lens = perception.calls[0]
+    assert chunked_signal.shape == (len(expected_chunk_lens), max(expected_chunk_lens))
+    assert torch.equal(chunked_lens, torch.tensor(expected_chunk_lens, dtype=torch.long))
+    assert perception.spk_targets_calls[0] is None
+    assert torch.equal(embs[0].squeeze(-1), audios[0])
+
+
+@pytest.mark.parametrize(
+    ("audio_values", "audio_len", "expected_chunk_lens", "expected_spk_targets"),
+    [
+        (
+            [1.0, 2.0, 3.0, 4.0, 5.0],
+            5,
+            [2, 3],
+            [
+                [[0.0, 1.0, 2.0, 3.0], [4.0, 5.0, 6.0, 7.0], [4.0, 5.0, 6.0, 7.0]],
+                [[8.0, 9.0, 10.0, 11.0], [12.0, 13.0, 14.0, 15.0], [16.0, 17.0, 18.0, 19.0]],
+            ],
+        ),
+        (
+            [1.0, 2.0, 3.0, 4.0],
+            4,
+            [2, 2],
+            [
+                [[0.0, 1.0, 2.0, 3.0], [4.0, 5.0, 6.0, 7.0], [4.0, 5.0, 6.0, 7.0]],
+                [[8.0, 9.0, 10.0, 11.0], [12.0, 13.0, 14.0, 15.0], [16.0, 17.0, 18.0, 19.0]],
+            ],
+        ),
+    ],
+)
+def test_encode_audio_with_optional_chunking_forwards_chunked_spk_targets(
+    audio_values, audio_len, expected_chunk_lens, expected_spk_targets
+):
+    cfg = MultiSpeakerConfig()
+    perception = ChunkingTestPerception(sampling_rate=2, hop_length=1)
+    audios = torch.tensor([audio_values])
+    audio_lens = torch.tensor([audio_len], dtype=torch.long)
+    spk_targets = torch.arange(5 * cfg.num_speakers, dtype=torch.float32).reshape(1, 5, cfg.num_speakers)
+
+    embs = encode_audio_with_optional_chunking(
+        perception,
+        audios,
+        audio_lens,
+        chunk_size_seconds=1.0,
+        sampling_rate=2,
+        spk_targets=spk_targets,
+    )
+
+    chunked_signal, chunked_lens = perception.calls[0]
+    assert chunked_signal.shape == (len(expected_chunk_lens), max(expected_chunk_lens))
+    assert torch.equal(chunked_lens, torch.tensor(expected_chunk_lens, dtype=torch.long))
+    assert torch.equal(perception.spk_targets_calls[0], torch.tensor(expected_spk_targets))
+    assert torch.equal(embs[0].squeeze(-1), audios[0])

tests/collections/speechlm2/test_force_align.py

@@ -0,0 +1,169 @@
+# Copyright (c) 2025, 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.
+
+import multiprocessing as mp
+import os
+
+import pytest
+import torch
+from lhotse import CutSet, Recording, SupervisionSegment
+from lhotse.testing.dummies import dummy_cut, dummy_recording
+
+from nemo.collections.speechlm2.data.force_align import ForceAligner
+
+# Set spawn method to avoid fork+CUDA conflicts that cause ForceAligner to fall back to CPU
+if mp.get_start_method(allow_none=True) != 'spawn':
+    mp.set_start_method('spawn', force=True)
+
+TEST_DATA_DIR = os.path.join(os.path.dirname(__file__), "test_data")
+
+
+@pytest.fixture(scope="module")
+def force_aligner():
+    """Create a ForceAligner instance for testing"""
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    aligner = ForceAligner(device=device, frame_length=0.08)
+    return aligner
+
+
+@pytest.fixture(scope="module")
+def test_cutset_from_audio_file():
+    """Create a test cutset from a pre-recorded audio file."""
+    audio_path = os.path.join(TEST_DATA_DIR, "force_align_test.mp3")
+    text = "ten companies that let you teach english"
+
+    rec = Recording.from_file(audio_path)
+    cut = rec.to_cut()
+    cut.supervisions = [
+        SupervisionSegment(
+            id=f"{cut.id}-0",
+            recording_id=cut.recording_id,
+            start=0.0,
+            duration=rec.duration,
+            text=text,
+            speaker="user",
+        ),
+    ]
+
+    return CutSet([cut])
+
+
+def test_force_align_audio_file(force_aligner, test_cutset_from_audio_file):
+    """Test force alignment with a pre-recorded audio file."""
+    import re
+
+    # Store original texts before alignment
+    original_texts = {}
+    for cut in test_cutset_from_audio_file:
+        for sup in cut.supervisions:
+            if sup.speaker == "user":
+                original_texts[sup.id] = sup.text
+
+    result_cuts = force_aligner.batch_force_align_user_audio(test_cutset_from_audio_file, source_sample_rate=24000)
+
+    assert len(result_cuts) == len(test_cutset_from_audio_file)
+    assert len(result_cuts) == 1
+
+    for cut in result_cuts:
+        user_supervisions = [s for s in cut.supervisions if s.speaker == "user"]
+        assert len(user_supervisions) > 0
+
+        for sup in user_supervisions:
+            original_text = original_texts.get(sup.id, "")
+            aligned_text = sup.text
+
+            print(f"\n{'='*80}")
+            print(f"Supervision ID: {sup.id}")
+            print(f"{'='*80}")
+            print(f"ORIGINAL TEXT:\n  {original_text}")
+            print(f"\nALIGNED TEXT:\n  {aligned_text}")
+            print(f"{'='*80}")
+
+            if "<|" not in aligned_text:
+                # TODO(kevinhu): Fix CUDA/numpy device mismatch in NeMo aligner utils
+                # (get_batch_variables returns CUDA tensors that fail on .numpy() calls)
+                pytest.skip("Force alignment did not produce timestamps (likely CUDA/numpy device mismatch in CI)")
+
+            # Extract timestamp-word-timestamp patterns: <|start|> word <|end|>
+            pattern = r'<\|(\d+)\|>\s+(\S+)\s+<\|(\d+)\|>'
+            matches = re.findall(pattern, aligned_text)
+
+            words_only = re.sub(r'<\|\d+\|>', '', aligned_text).split()
+            words_only = [w for w in words_only if w]
+
+            print(f"\nValidation: Found {len(matches)} timestamped words out of {len(words_only)} total words")
+
+            assert len(matches) > 0, "Should have at least one timestamped word"
+            assert len(matches) == len(
+                words_only
+            ), f"Every word should have timestamps. Found {len(matches)} timestamped words but {len(words_only)} total words"
+
+            for start_frame, word, end_frame in matches:
+                start_frame = int(start_frame)
+                end_frame = int(end_frame)
+
+                assert (
+                    start_frame <= end_frame
+                ), f"Start frame {start_frame} should be before or equal to end frame {end_frame} for word '{word}'"
+                assert start_frame >= 0, f"Start frame should be non-negative for word '{word}'"
+                assert end_frame >= 0, f"End frame should be non-negative for word '{word}'"
+
+
+def test_force_align_no_user_supervisions(force_aligner):
+    """Test with cutset containing no user supervisions"""
+    cut = dummy_cut(0, recording=dummy_recording(0, with_data=True))
+    cut.supervisions = [
+        SupervisionSegment(
+            id=cut.id,
+            recording_id=cut.recording_id,
+            start=0,
+            duration=0.5,
+            text='hello',
+            speaker="assistant",
+        ),
+    ]
+    cutset = CutSet([cut])
+
+    result_cuts = force_aligner.batch_force_align_user_audio(cutset)
+
+    assert len(result_cuts) == 1
+    result_supervisions = list(result_cuts)[0].supervisions
+    assert result_supervisions[0].text == 'hello'
+
+
+def test_force_align_empty_cutset(force_aligner):
+    """Test with empty cutset"""
+    empty_cutset = CutSet.from_cuts([])
+    result_cuts = force_aligner.batch_force_align_user_audio(empty_cutset)
+    assert len(result_cuts) == 0
+
+
+def test_strip_timestamps(force_aligner):
+    """Test timestamp stripping utility"""
+    text_with_timestamps = "<|10|> hello <|20|> world <|30|>"
+    result = force_aligner._strip_timestamps(text_with_timestamps)
+    assert result == "hello world"
+    assert "<|" not in result
+
+    text_without_timestamps = "hello world"
+    result = force_aligner._strip_timestamps(text_without_timestamps)
+    assert result == "hello world"
+
+
+def test_normalize_transcript(force_aligner):
+    """Test transcript normalization"""
+    assert force_aligner._normalize_transcript("Hello World!") == "hello world"
+    assert force_aligner._normalize_transcript("don't worry") == "don't worry"
+    assert force_aligner._normalize_transcript("test123") == "test"
+    assert force_aligner._normalize_transcript("A,B.C!D?E") == "a b c d e"

tests/collections/speechlm2/test_freezing_params.py

@@ -0,0 +1,96 @@
+# Copyright (c) 2025, 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.
+import torch.nn
+from omegaconf import DictConfig
+
+import nemo.core.optim.lr_scheduler
+from nemo.collections.speechlm2.parts.optim_setup import configure_optimizers, freeze_and_subset
+
+
+class DummyModel(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.linear = torch.nn.Linear(1, 1)
+        self.conv = torch.nn.Conv1d(1, 1, 1)
+
+    def forward(self, x):
+        x = self.linear(x)
+        x = self.conv(x)
+        return x
+
+
+def test_freezing_params():
+    model = DummyModel().train()
+    assert model.linear.weight.requires_grad
+    assert model.linear.bias.requires_grad
+    assert model.conv.weight.requires_grad
+    assert model.conv.bias.requires_grad
+    params = freeze_and_subset(model.named_parameters(), exclude_patterns=[r"linear\..+"])
+    list(params)  # execute generator
+    assert not model.linear.weight.requires_grad
+    assert not model.linear.bias.requires_grad
+    assert model.conv.weight.requires_grad
+    assert model.conv.bias.requires_grad
+
+
+def test_keeping_unfrozen_params():
+    model = DummyModel().train()
+    assert model.linear.weight.requires_grad
+    assert model.linear.bias.requires_grad
+    assert model.conv.weight.requires_grad
+    assert model.conv.bias.requires_grad
+    params = freeze_and_subset(
+        model.named_parameters(), exclude_patterns=[r"linear\..+"], keep_patterns=[r"linear.bias"]
+    )
+    list(params)  # execute generator
+    assert not model.linear.weight.requires_grad
+    assert model.linear.bias.requires_grad
+    assert model.conv.weight.requires_grad
+    assert model.conv.bias.requires_grad
+
+
+def test_configure_optimizers():
+    model = DummyModel()
+    model.cfg = DictConfig(
+        {
+            "optimizer": {"_target_": "torch.optim.adamw.AdamW"},
+            "freeze_params": [r"conv\..+"],
+        }
+    )
+    ans = configure_optimizers(model)
+    assert ans.keys() == {"optimizer"}
+    assert isinstance(ans["optimizer"], torch.optim.AdamW)
+    parameters = ans["optimizer"].param_groups[0]['params']
+    assert len(parameters) == 2
+    assert parameters[0] == model.linear.weight
+    assert parameters[1] == model.linear.bias
+
+
+def test_configure_optimizers_with_lr_scheduler():
+    model = DummyModel()
+    model.cfg = DictConfig(
+        {
+            "optimizer": {"_target_": "torch.optim.adamw.AdamW"},
+            "lr_scheduler": {
+                "_target_": "nemo.core.optim.lr_scheduler.CosineAnnealing",
+                "warmup_steps": 0,
+                "min_lr": 1e-6,
+                "max_steps": 100000,
+            },
+        }
+    )
+    ans = configure_optimizers(model)
+    assert ans.keys() == {"optimizer", "lr_scheduler"}
+    assert isinstance(ans["optimizer"], torch.optim.AdamW)
+    assert isinstance(ans["lr_scheduler"]["scheduler"], nemo.core.optim.lr_scheduler.CosineAnnealing)

tests/collections/speechlm2/test_hf_hub.py

@@ -0,0 +1,75 @@
+# 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.speechlm2.parts.hf_hub import _inject_local_artifact_paths
+
+
+def _cached_file_kwargs():
+    return {
+        "cache_dir": None,
+        "force_download": False,
+        "local_files_only": True,
+        "token": None,
+        "revision": None,
+        "_raise_exceptions_for_gated_repo": False,
+        "_raise_exceptions_for_missing_entries": False,
+        "_raise_exceptions_for_connection_errors": False,
+    }
+
+
+def _write_local_export_artifacts(tmp_path):
+    (tmp_path / "tokenizer_config.json").write_text("{}")
+    (tmp_path / "llm_backbone").mkdir()
+    (tmp_path / "llm_backbone" / "config.json").write_text("{}")
+
+
+def test_inject_local_artifact_paths_salm_config(tmp_path):
+    _write_local_export_artifacts(tmp_path)
+    cfg = {
+        "pretrained_llm": "remote-llm",
+        "pretrained_asr": "remote-asr",
+    }
+
+    _inject_local_artifact_paths(cfg, str(tmp_path), _cached_file_kwargs())
+
+    assert cfg["pretrained_llm"] == str(tmp_path / "llm_backbone")
+    assert cfg["pretrained_asr"] == "remote-asr"
+    assert cfg["tokenizer_path"] == str(tmp_path)
+
+
+def test_inject_local_artifact_paths_duplex_eartts_config(tmp_path):
+    _write_local_export_artifacts(tmp_path)
+    cfg = {
+        "pretrained_lm_name": "remote-llm",
+        "tts_config": {},
+    }
+
+    _inject_local_artifact_paths(cfg, str(tmp_path), _cached_file_kwargs())
+
+    assert cfg["pretrained_lm_name"] == str(tmp_path / "llm_backbone")
+    assert cfg["tokenizer_path"] == str(tmp_path)
+
+
+def test_inject_local_artifact_paths_no_artifacts_keeps_old_config(tmp_path):
+    cfg = {
+        "pretrained_llm": "remote-llm",
+        "pretrained_weights": True,
+    }
+
+    _inject_local_artifact_paths(cfg, str(tmp_path), _cached_file_kwargs())
+
+    assert cfg == {
+        "pretrained_llm": "remote-llm",
+        "pretrained_weights": True,
+    }

tests/collections/speechlm2/test_init_from_checkpoint.py

@@ -0,0 +1,368 @@
+# Copyright (c) 2025, 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.
+"""Tests for init_from_checkpoint functionality in speechlm2.
+
+Unit tests use simple nn.Module subclasses (no HF downloads, no CUDA).
+Integration tests use real SALM / SALMAutomodel (require HF config download;
+SALMAutomodel tests also require CUDA).
+"""
+import os
+from unittest.mock import patch
+
+import pytest
+import torch
+from omegaconf import DictConfig
+from safetensors.torch import save_file
+
+from nemo.collections.speechlm2.parts.pretrained import (
+    _is_dcp_checkpoint,
+    init_from_training_checkpoint,
+    maybe_load_pretrained_models,
+)
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+class SimpleModel(torch.nn.Module):
+    """Tiny model used for fast, self-contained unit tests."""
+
+    def __init__(self):
+        super().__init__()
+        self.linear = torch.nn.Linear(4, 4, bias=False)
+        self.norm = torch.nn.LayerNorm(4)
+
+
+class ConfigurableModel(torch.nn.Module):
+    """SimpleModel that carries a ``cfg`` attribute, like SALM/SALMAutomodel."""
+
+    def __init__(self, cfg: dict):
+        super().__init__()
+        self.cfg = DictConfig(cfg)
+        self.linear = torch.nn.Linear(4, 4, bias=False)
+        self.norm = torch.nn.LayerNorm(4)
+
+
+def _save_ckpt(model, path):
+    """Save model state_dict in Lightning checkpoint format."""
+    torch.save({"state_dict": model.state_dict()}, path)
+
+
+def _assert_state_dicts_equal(sd1, sd2):
+    assert set(sd1.keys()) == set(sd2.keys()), f"Key mismatch: {set(sd1.keys()) ^ set(sd2.keys())}"
+    for key in sd1:
+        assert torch.equal(sd1[key].cpu(), sd2[key].cpu()), f"Tensor mismatch at key: {key}"
+
+
+def _assert_state_dicts_not_equal(sd1, sd2):
+    """Assert at least one tensor differs (sanity-check before loading)."""
+    assert set(sd1.keys()) == set(sd2.keys())
+    any_diff = any(not torch.equal(sd1[k].cpu(), sd2[k].cpu()) for k in sd1)
+    assert any_diff, "Expected state dicts to differ before checkpoint loading"
+
+
+# ---------------------------------------------------------------------------
+# _is_dcp_checkpoint
+# ---------------------------------------------------------------------------
+
+
+class TestIsDcpCheckpoint:
+    def test_with_metadata(self, tmp_path):
+        ckpt_dir = tmp_path / "step=100.ckpt"
+        ckpt_dir.mkdir()
+        (ckpt_dir / ".metadata").touch()
+        assert _is_dcp_checkpoint(str(ckpt_dir))
+
+    def test_without_metadata(self, tmp_path):
+        ckpt_dir = tmp_path / "step=100.ckpt"
+        ckpt_dir.mkdir()
+        assert not _is_dcp_checkpoint(str(ckpt_dir))
+
+    def test_regular_file(self, tmp_path):
+        ckpt_file = tmp_path / "step=100.ckpt"
+        ckpt_file.touch()
+        assert not _is_dcp_checkpoint(str(ckpt_file))
+
+    def test_nonexistent_path(self, tmp_path):
+        assert not _is_dcp_checkpoint(str(tmp_path / "nonexistent"))
+
+    def test_hf_dir_without_metadata(self, tmp_path):
+        """HF directory (model.safetensors) should NOT be detected as DCP."""
+        hf_dir = tmp_path / "hf_model"
+        hf_dir.mkdir()
+        (hf_dir / "model.safetensors").touch()
+        assert not _is_dcp_checkpoint(str(hf_dir))
+
+
+# ---------------------------------------------------------------------------
+# init_from_training_checkpoint — non-DCP paths
+# ---------------------------------------------------------------------------
+
+
+class TestInitFromTrainingCheckpoint:
+    def test_none_is_noop(self):
+        model = SimpleModel()
+        original = model.linear.weight.clone()
+        init_from_training_checkpoint(model, None)
+        assert torch.equal(model.linear.weight, original)
+
+    def test_single_file_ckpt(self, tmp_path):
+        source = SimpleModel()
+        torch.nn.init.ones_(source.linear.weight)
+
+        ckpt_path = str(tmp_path / "model.ckpt")
+        _save_ckpt(source, ckpt_path)
+
+        target = SimpleModel()  # different random init
+        _assert_state_dicts_not_equal(source.state_dict(), target.state_dict())
+
+        init_from_training_checkpoint(target, ckpt_path)
+        _assert_state_dicts_equal(target.state_dict(), source.state_dict())
+
+    def test_hf_directory(self, tmp_path):
+        source = SimpleModel()
+        torch.nn.init.ones_(source.linear.weight)
+
+        hf_dir = tmp_path / "hf_model"
+        hf_dir.mkdir()
+        save_file(source.state_dict(), str(hf_dir / "model.safetensors"))
+
+        target = SimpleModel()
+        _assert_state_dicts_not_equal(source.state_dict(), target.state_dict())
+
+        init_from_training_checkpoint(target, str(hf_dir))
+        _assert_state_dicts_equal(target.state_dict(), source.state_dict())
+
+
+# ---------------------------------------------------------------------------
+# init_from_training_checkpoint — DCP path (mocked)
+# ---------------------------------------------------------------------------
+
+
+class TestInitFromTrainingCheckpointDCP:
+    def test_dcp_calls_distributed_load(self, tmp_path):
+        """Verify DCP checkpoint triggers torch.distributed.checkpoint.load."""
+        ckpt_dir = tmp_path / "step=100.ckpt"
+        ckpt_dir.mkdir()
+        (ckpt_dir / ".metadata").touch()
+
+        model = SimpleModel()
+
+        with patch("nemo.collections.speechlm2.parts.pretrained.torch.distributed.checkpoint.load") as mock_load:
+            init_from_training_checkpoint(model, str(ckpt_dir))
+
+            mock_load.assert_called_once()
+            args, kwargs = mock_load.call_args
+            state_dict_wrapper = args[0]
+            assert "state_dict" in state_dict_wrapper
+            assert kwargs["checkpoint_id"] == str(ckpt_dir)
+
+    def test_dcp_state_dict_has_model_keys(self, tmp_path):
+        """The state dict passed to dcp.load should contain model parameter keys."""
+        ckpt_dir = tmp_path / "step=100.ckpt"
+        ckpt_dir.mkdir()
+        (ckpt_dir / ".metadata").touch()
+
+        model = SimpleModel()
+
+        with patch("nemo.collections.speechlm2.parts.pretrained.torch.distributed.checkpoint.load") as mock_load:
+            init_from_training_checkpoint(model, str(ckpt_dir))
+
+            state_dict_wrapper = mock_load.call_args[0][0]
+            model_sd = state_dict_wrapper["state_dict"]
+            assert "linear.weight" in model_sd
+            assert "norm.weight" in model_sd
+
+
+# ---------------------------------------------------------------------------
+# maybe_load_pretrained_models — init_from_checkpoint config key
+# ---------------------------------------------------------------------------
+
+
+class TestMaybeLoadPretrainedModels:
+    def test_init_from_checkpoint_loads_weights(self, tmp_path):
+        source = ConfigurableModel(cfg={})
+        torch.nn.init.ones_(source.linear.weight)
+        ckpt_path = str(tmp_path / "model.ckpt")
+        _save_ckpt(source, ckpt_path)
+
+        target = ConfigurableModel(cfg={"init_from_checkpoint": ckpt_path})
+        _assert_state_dicts_not_equal(source.state_dict(), target.state_dict())
+
+        maybe_load_pretrained_models(target)
+        _assert_state_dicts_equal(target.state_dict(), source.state_dict())
+
+    def test_init_from_checkpoint_null_is_noop(self):
+        model = ConfigurableModel(cfg={"init_from_checkpoint": None})
+        original = model.linear.weight.clone()
+        maybe_load_pretrained_models(model)
+        assert torch.equal(model.linear.weight, original)
+
+    def test_init_from_checkpoint_key_missing_is_noop(self):
+        model = ConfigurableModel(cfg={})
+        original = model.linear.weight.clone()
+        maybe_load_pretrained_models(model)
+        assert torch.equal(model.linear.weight, original)
+
+    def test_pretrained_s2s_model_still_works(self, tmp_path):
+        """Backward compat: pretrained_s2s_model should still be recognized."""
+        source = ConfigurableModel(cfg={})
+        torch.nn.init.ones_(source.linear.weight)
+        ckpt_path = str(tmp_path / "model.ckpt")
+        _save_ckpt(source, ckpt_path)
+
+        target = ConfigurableModel(cfg={"pretrained_s2s_model": ckpt_path})
+        maybe_load_pretrained_models(target)
+        _assert_state_dicts_equal(target.state_dict(), source.state_dict())
+
+
+# ---------------------------------------------------------------------------
+# SALM integration test
+# ---------------------------------------------------------------------------
+
+AUDIO_LOCATOR_TAG = "<|audioplaceholder|>"
+SALM_PERCEPTION_CFG = {
+    "target": "nemo.collections.speechlm2.modules.perception.AudioPerceptionModule",
+    "output_dim": 2048,
+    "encoder": {
+        "_target_": "nemo.collections.asr.modules.ConformerEncoder",
+        "att_context_size": [-1, -1],
+        "causal_downsampling": False,
+        "conv_context_size": None,
+        "conv_kernel_size": 9,
+        "conv_norm_type": "batch_norm",
+        "d_model": 1024,
+        "dropout": 0.1,
+        "dropout_att": 0.1,
+        "dropout_emb": 0.0,
+        "dropout_pre_encoder": 0.1,
+        "feat_in": 128,
+        "feat_out": -1,
+        "ff_expansion_factor": 4,
+        "n_heads": 8,
+        "n_layers": 2,
+        "pos_emb_max_len": 5000,
+        "self_attention_model": "rel_pos",
+        "subsampling": "dw_striding",
+        "subsampling_conv_channels": 256,
+        "subsampling_factor": 8,
+    },
+    "modality_adapter": {
+        "_target_": "nemo.collections.speechlm2.modules.perception.IdentityConnector",
+        "d_model": 1024,
+    },
+    "preprocessor": {
+        "_target_": "nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor",
+        "dither": 1e-05,
+        "features": 128,
+        "frame_splicing": 1,
+        "log": True,
+        "n_fft": 512,
+        "normalize": "per_feature",
+        "pad_to": 0,
+        "pad_value": 0.0,
+        "sample_rate": 16000,
+        "window": "hann",
+        "window_size": 0.025,
+        "window_stride": 0.01,
+    },
+}
+
+
+def _resolve_pretrained_salm():
+    if os.path.exists("/home/TestData/speechlm/pretrained_models"):
+        return {
+            "pretrained_llm": "/home/TestData/speechlm/pretrained_models/TinyLlama--TinyLlama_v1.1",
+            "pretrained_asr": "/home/TestData/speechlm/pretrained_models/canary-1b-flash.nemo",
+        }
+    return {
+        "pretrained_asr": "nvidia/canary-1b-flash",
+        "pretrained_llm": "TinyLlama/TinyLlama_v1.1",
+    }
+
+
+def _resolve_pretrained_automodel():
+    if os.path.exists("/home/TestData/speechlm/pretrained_models"):
+        return {
+            "pretrained_llm": "/home/TestData/speechlm/pretrained_models/Qwen--Qwen3-1.7B",
+            "pretrained_asr": "/home/TestData/speechlm/pretrained_models/canary-1b-flash.nemo",
+        }
+    return {
+        "pretrained_asr": "nvidia/canary-1b-flash",
+        "pretrained_llm": "Qwen/Qwen3-1.7B",
+    }
+
+
+def _make_salm_cfg(**overrides):
+    cfg = {
+        **_resolve_pretrained_salm(),
+        "pretrained_weights": False,
+        "prompt_format": "llama2",
+        "audio_locator_tag": AUDIO_LOCATOR_TAG,
+        "perception": SALM_PERCEPTION_CFG,
+        "optimizer": {"_target_": "torch.optim.AdamW"},
+    }
+    cfg.update(overrides)
+    return cfg
+
+
+def _make_automodel_cfg(**overrides):
+    cfg = {
+        **_resolve_pretrained_automodel(),
+        "pretrained_weights": False,
+        "prompt_format": "qwen",
+        "audio_locator_tag": AUDIO_LOCATOR_TAG,
+        "perception": SALM_PERCEPTION_CFG,
+        "optimizer": {"_target_": "torch.optim.AdamW"},
+        "torch_dtype": "bfloat16",
+    }
+    cfg.update(overrides)
+    return cfg
+
+
+def test_salm_init_from_checkpoint(tmp_path):
+    from nemo.collections.speechlm2.models import SALM
+
+    # Create source model and save checkpoint
+    model1 = SALM(_make_salm_cfg())
+    expected_sd = {k: v.clone().cpu() for k, v in model1.state_dict().items()}
+    ckpt_path = str(tmp_path / "source.ckpt")
+    _save_ckpt(model1, ckpt_path)
+    del model1
+
+    # Create target model — init_from_checkpoint overrides the random init
+    model2 = SALM(_make_salm_cfg(init_from_checkpoint=ckpt_path))
+    _assert_state_dicts_equal(model2.state_dict(), expected_sd)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="SALMAutomodel requires CUDA")
+def test_salm_automodel_init_from_checkpoint(tmp_path):
+    from nemo.collections.speechlm2.models import SALMAutomodel
+
+    # Create source model and save checkpoint
+    model1 = SALMAutomodel(_make_automodel_cfg())
+    model1.configure_model()
+    expected_sd = {k: v.clone().cpu() for k, v in model1.state_dict().items()}
+    ckpt_path = str(tmp_path / "source.ckpt")
+    _save_ckpt(model1, ckpt_path)
+    del model1
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+    # Create target model — configure_model loads checkpoint via maybe_load_pretrained_models
+    model2 = SALMAutomodel(_make_automodel_cfg(init_from_checkpoint=ckpt_path))
+    model2.configure_model()
+    _assert_state_dicts_equal(model2.state_dict(), expected_sd)

tests/collections/speechlm2/test_label_prep.py

@@ -0,0 +1,138 @@
+# Copyright (c) 2025, 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.
+
+import torch
+
+from nemo.collections.speechlm2.parts.label_prep import maybe_prepend_prompt_tokens
+
+
+def test_maybe_prepend_prompt_tokens_with_source_tokens():
+    """Test that prompt tokens are correctly prepended to all sequences and lengths are updated."""
+    B, T_src, T_tgt, H = 2, 6, 6, 4
+    PAD = 0
+    prompt_len_0, prompt_len_1 = 3, 2
+
+    # Prompt token IDs (will be passed through embed_fn)
+    max_prompt_len = max(prompt_len_0, prompt_len_1)
+    prompt_tokens = torch.full((B, max_prompt_len), PAD, dtype=torch.long)
+    prompt_tokens[0, :prompt_len_0] = torch.tensor([10, 11, 12])
+    prompt_tokens[1, :prompt_len_1] = torch.tensor([20, 21])
+
+    # Use a simple embedding: token_id -> [token_id, token_id, token_id, token_id]
+    def embed_fn(token_ids):
+        return token_ids.unsqueeze(-1).expand(-1, -1, H).float()
+
+    # Source encoded (audio features)
+    source_encoded = torch.arange(B * T_src * H).reshape(B, T_src, H).float()
+    source_encoded_lens = torch.tensor([5, 4])
+
+    # Target tokens
+    target_tokens = torch.tensor(
+        [
+            [1, 100, 101, 102, 2, 0],  # BOS, tokens, EOS, PAD
+            [1, 200, 201, 2, 0, 0],
+        ]
+    )
+    target_token_lens = torch.tensor([5, 4])
+
+    # Source tokens (for ASR head)
+    source_tokens = torch.tensor(
+        [
+            [1, 50, 51, 52, 2, 0],
+            [1, 60, 61, 2, 0, 0],
+        ]
+    )
+    source_token_lens = torch.tensor([5, 4])
+
+    batch = {
+        "prompt_tokens": prompt_tokens,
+        "prompt_token_lens": torch.tensor([prompt_len_0, prompt_len_1]),
+        "target_tokens": target_tokens,
+        "target_token_lens": target_token_lens,
+        "source_tokens": source_tokens,
+        "source_token_lens": source_token_lens,
+    }
+
+    new_source_encoded, new_source_encoded_lens, new_target_tokens = maybe_prepend_prompt_tokens(
+        batch=batch,
+        embed_fn=embed_fn,
+        source_encoded=source_encoded,
+        source_encoded_lens=source_encoded_lens,
+        text_pad_id=PAD,
+    )
+
+    # Check output shapes are extended by max_prompt_len
+    assert new_source_encoded.shape == (B, max_prompt_len + T_src, H)
+    assert new_target_tokens.shape == (B, max_prompt_len + T_tgt)
+    assert batch["source_tokens"].shape == (B, max_prompt_len + T_tgt)
+
+    # Check lengths are updated: original_len + prompt_len
+    assert new_source_encoded_lens[0].item() == 5 + prompt_len_0
+    assert new_source_encoded_lens[1].item() == 4 + prompt_len_1
+    assert batch["target_token_lens"][0].item() == 5 + prompt_len_0
+    assert batch["target_token_lens"][1].item() == 4 + prompt_len_1
+    assert batch["source_token_lens"][0].item() == 5 + prompt_len_0
+    assert batch["source_token_lens"][1].item() == 4 + prompt_len_1
+
+    # Check prompt embeddings are at the beginning of source_encoded
+    # embed_fn maps token_id -> [token_id]*H, so prompt region should match
+    for h in range(H):
+        assert new_source_encoded[0, 0, h].item() == 10.0
+        assert new_source_encoded[0, 1, h].item() == 11.0
+        assert new_source_encoded[0, 2, h].item() == 12.0
+        assert new_source_encoded[1, 0, h].item() == 20.0
+        assert new_source_encoded[1, 1, h].item() == 21.0
+
+    # Check original audio features follow the prompt
+    for t in range(5):  # source_encoded_lens[0] was 5
+        assert torch.equal(new_source_encoded[0, prompt_len_0 + t], source_encoded[0, t])
+    for t in range(4):  # source_encoded_lens[1] was 4
+        assert torch.equal(new_source_encoded[1, prompt_len_1 + t], source_encoded[1, t])
+
+    # Check target tokens are shifted by prompt_len
+    assert new_target_tokens[0, :prompt_len_0].tolist() == [PAD] * prompt_len_0
+    assert new_target_tokens[0, prompt_len_0 : prompt_len_0 + 5].tolist() == [1, 100, 101, 102, 2]
+    assert new_target_tokens[1, :prompt_len_1].tolist() == [PAD] * prompt_len_1
+    assert new_target_tokens[1, prompt_len_1 : prompt_len_1 + 4].tolist() == [1, 200, 201, 2]
+
+    # Check source tokens are shifted by prompt_len
+    assert batch["source_tokens"][0, :prompt_len_0].tolist() == [PAD] * prompt_len_0
+    assert batch["source_tokens"][0, prompt_len_0 : prompt_len_0 + 5].tolist() == [1, 50, 51, 52, 2]
+    assert batch["source_tokens"][1, :prompt_len_1].tolist() == [PAD] * prompt_len_1
+    assert batch["source_tokens"][1, prompt_len_1 : prompt_len_1 + 4].tolist() == [1, 60, 61, 2]
+
+
+def test_maybe_prepend_prompt_tokens_no_prompt():
+    """Test that without prompt_tokens in batch, inputs are returned unchanged."""
+    B, T_src, H = 1, 4, 4
+    source_encoded = torch.randn(B, T_src, H)
+    source_encoded_lens = torch.tensor([3])
+    target_tokens = torch.tensor([[1, 100, 2, 0]])
+
+    batch = {
+        "target_tokens": target_tokens,
+        "target_token_lens": torch.tensor([3]),
+    }
+
+    out_encoded, out_lens, out_tokens = maybe_prepend_prompt_tokens(
+        batch=batch,
+        embed_fn=lambda x: x.unsqueeze(-1).expand(-1, -1, H).float(),
+        source_encoded=source_encoded,
+        source_encoded_lens=source_encoded_lens,
+        text_pad_id=0,
+    )
+
+    assert torch.equal(out_encoded, source_encoded)
+    assert torch.equal(out_lens, source_encoded_lens)
+    assert torch.equal(out_tokens, target_tokens)

tests/collections/speechlm2/test_metrics.py

@@ -0,0 +1,392 @@
+# Copyright (c) 2025, 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.
+import pytest
+import torch
+
+from nemo.collections.speechlm2.parts.metrics import BLEU, WER, Intelligibility
+from nemo.collections.speechlm2.parts.metrics.empty_text import EmptyTextMetric
+from nemo.collections.speechlm2.parts.metrics.perplexity import Perplexity
+
+
+def test_bleu():
+    metric = BLEU(verbose=False)
+    metric.update(
+        name="dataset_1",
+        refs=["a b c d e f g h i j k l", "m n o p r s t u v"],
+        hyps=["a b c d e f g h i j k l", "m n o p r s t u v"],
+    )
+    metric.update(
+        name="dataset_2",
+        refs=["a b c"],
+        hyps=["a b d"],
+    )
+    ans = metric.compute()
+    assert ans["txt_bleu_dataset_1"] == 100.0
+    assert ans["txt_bleu_dataset_2"] == 0.0
+    assert ans["txt_bleu"] == 50.0  # average across datasets
+
+
+def test_wer():
+    metric = WER(verbose=False)
+    metric.update(
+        name="dataset_1",
+        refs=["a b c d e f g h i j k l", "m n o p r s t u v"],
+        hyps=["a b c d e f g h i j k l", "m n o p r s t u v"],
+    )
+    metric.update(
+        name="dataset_2",
+        refs=["a b c"],
+        hyps=["a b d"],
+    )
+    ans = metric.compute()
+    assert ans["wer_dataset_1"] == 0.0
+    assert ans["wer_dataset_2"] == 1 / 3
+    assert ans["wer"] == 1 / 6  # average across datasets
+
+
+def test_empty_text_metric():
+    """Test EmptyTextMetric for detecting empty hypotheses"""
+    metric = EmptyTextMetric(verbose=False)
+
+    # Test with some empty and non-empty texts
+    hyps = ["hello world", "", "  ", "test", "   \n  "]
+    metric.update("test_batch", hyps)
+
+    results = metric.compute()
+
+    # Should detect 3 empty texts out of 5
+    assert "empty_text_rate_test_batch" in results
+    assert results["empty_text_rate_test_batch"].item() == pytest.approx(0.6, abs=0.01)  # 3/5 = 0.6
+
+
+def test_empty_text_metric_reset():
+    """Test EmptyTextMetric reset functionality"""
+    metric = EmptyTextMetric(verbose=False)
+
+    # Add some data
+    metric.update("test", ["hello", ""])
+    metric.reset()
+
+    # After reset, should have no data
+    results = metric.compute()
+    assert len(results) == 0
+
+
+def test_empty_text_all_valid():
+    """Test EmptyTextMetric with no empty texts"""
+    metric = EmptyTextMetric(verbose=False)
+
+    metric.update("test", ["hello", "world", "test"])
+
+    results = metric.compute()
+    assert results["empty_text_rate_test"].item() == 0.0
+
+
+def test_empty_text_all_empty():
+    """Test EmptyTextMetric with all empty texts"""
+    metric = EmptyTextMetric(verbose=False)
+
+    metric.update("test", ["", "  ", "\n"])
+
+    results = metric.compute()
+    assert results["empty_text_rate_test"].item() == 1.0
+
+
+def test_perplexity_basic():
+    """Test basic perplexity calculation"""
+    metric = Perplexity(ignore_index=-100, verbose=False)
+
+    # Create simple logits and targets
+    vocab_size = 10
+    batch_size = 2
+    seq_len = 3
+
+    # Create perfect predictions (all correct)
+    logits = torch.zeros(batch_size, seq_len, vocab_size)
+    targets = torch.tensor([[1, 2, 3], [4, 5, 6]])
+
+    # Set logits to have high probability for correct tokens
+    for b in range(batch_size):
+        for s in range(seq_len):
+            logits[b, s, targets[b, s]] = 10.0  # High logit for correct token
+
+    ppl = metric.update("test", logits, targets)
+
+    # With perfect predictions, perplexity should be close to 1
+    assert ppl < 2.0  # Should be very low
+
+
+def test_perplexity_with_padding():
+    """Test perplexity with padding tokens"""
+    metric = Perplexity(ignore_index=-100, verbose=False)
+
+    vocab_size = 10
+    batch_size = 2
+    seq_len = 4
+
+    logits = torch.randn(batch_size, seq_len, vocab_size)
+    # Include padding tokens (ignore_index = -100)
+    targets = torch.tensor([[1, 2, 3, -100], [4, 5, -100, -100]])
+
+    ppl = metric.update("test", logits, targets)
+
+    # Should not fail with padding
+    assert ppl > 0
+    assert not torch.isnan(torch.tensor(ppl))
+
+
+def test_perplexity_compute():
+    """Test perplexity compute aggregation"""
+    metric = Perplexity(ignore_index=-100, verbose=False)
+
+    vocab_size = 10
+    logits = torch.randn(2, 3, vocab_size)
+    targets = torch.tensor([[1, 2, 3], [4, 5, 6]])
+
+    metric.update("test", logits, targets)
+
+    results = metric.compute()
+    assert "perplexity_test" in results
+    assert results["perplexity_test"] > 0
+
+
+def test_turn_taking_import():
+    """Test that turn taking metric function can be imported"""
+    from nemo.collections.speechlm2.parts.metrics.turn_taking import compute_turn_taking_metrics
+
+    # Test with dummy data
+    source_tokens = torch.tensor([[1, 2, 3, 4]])  # dummy tokens
+    pred_tokens = torch.tensor([[5, 6, 7, 8]])  # dummy tokens
+    eos_token_id = 4
+    bos_token_id = 5
+
+    accuracy, latency = compute_turn_taking_metrics(source_tokens, pred_tokens, eos_token_id, bos_token_id)
+    assert isinstance(accuracy, float)
+    assert isinstance(latency, float)
+
+
+def test_mcq_evaluator_import():
+    """Test that MCQ evaluator can be imported and initialized"""
+    import tempfile
+
+    from nemo.collections.speechlm2.parts.metrics.mcq_evaluator import MCQEvaluator
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        evaluator = MCQEvaluator(manifest_dir=tmpdir)
+        assert evaluator is not None
+        assert evaluator.manifest_dir == tmpdir
+
+
+def test_results_logger_import():
+    """Test that results logger can be imported"""
+    import tempfile
+
+    from nemo.collections.speechlm2.parts.metrics.results_logger import ResultsLogger
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        logger = ResultsLogger(save_path=tmpdir)
+        assert logger is not None
+
+
+def test_results_logger_single_rank(tmp_path):
+    """Test ResultsLogger with single rank (non-distributed)"""
+    from unittest.mock import patch
+
+    from nemo.collections.speechlm2.parts.metrics.results_logger import ResultsLogger
+
+    save_path = str(tmp_path / "single_rank")
+
+    # Mock distributed functions to simulate single rank
+    with (
+        patch('nemo.collections.speechlm2.parts.metrics.results_logger.get_rank', return_value=0),
+        patch('nemo.collections.speechlm2.parts.metrics.results_logger.get_world_size', return_value=1),
+        patch('torch.distributed.is_available', return_value=False),
+    ):
+
+        logger = ResultsLogger(save_path=save_path)
+
+        # Add some test data
+        logger.update(
+            name="test_dataset",
+            refs=["hello world", "goodbye"],
+            hyps=["hello there", "goodbye"],
+            asr_hyps=[None, None],
+            samples_id=["sample1", "sample2"],
+            pred_audio=None,
+            pred_audio_sr=16000,
+            user_audio=None,
+            user_audio_sr=16000,
+            src_refs=["user input 1", "user input 2"],
+            src_hyps=["", ""],
+        )
+
+        # Compute and save
+        metrics = logger.compute_and_save(special_subset_names=[], mcq_subset_names=[])
+
+        # Check that files were created
+        import os
+
+        metadata_path = os.path.join(save_path, "metadatas", "test_dataset_rank0.json")
+        assert os.path.exists(metadata_path)
+
+        # Check final merged file (should be same as rank0 in single-rank case)
+        final_path = os.path.join(save_path, "metadatas", "test_dataset.json")
+        assert os.path.exists(final_path)
+
+
+def test_results_logger_multi_rank(tmp_path):
+    """Test ResultsLogger with multiple ranks (simulated distributed training)"""
+    import json
+    import os
+    from unittest.mock import MagicMock, patch
+
+    from nemo.collections.speechlm2.parts.metrics.results_logger import ResultsLogger
+
+    save_path = str(tmp_path / "multi_rank")
+    world_size = 4
+
+    # Simulate each rank saving its own results
+    for rank in range(world_size):
+        with (
+            patch('nemo.collections.speechlm2.parts.metrics.results_logger.get_rank', return_value=rank),
+            patch('nemo.collections.speechlm2.parts.metrics.results_logger.get_world_size', return_value=world_size),
+            patch('torch.distributed.is_available', return_value=True),
+            patch('torch.distributed.is_initialized', return_value=True),
+            patch('torch.distributed.barrier'),
+        ):  # Mock barrier to avoid actual distributed ops
+
+            logger = ResultsLogger(save_path=save_path)
+
+            # Each rank processes different samples
+            start_idx = rank * 2
+            logger.update(
+                name="test_dataset",
+                refs=[f"ref_{start_idx}", f"ref_{start_idx+1}"],
+                hyps=[f"hyp_{start_idx}", f"hyp_{start_idx+1}"],
+                asr_hyps=[None, None],
+                samples_id=[f"sample_{start_idx}", f"sample_{start_idx+1}"],
+                pred_audio=None,
+                pred_audio_sr=16000,
+                user_audio=None,
+                user_audio_sr=16000,
+                src_refs=[f"src_{start_idx}", f"src_{start_idx+1}"],
+                src_hyps=["", ""],
+            )
+
+            # Save rank-specific files (only this part, not the merge)
+            rank_json_path = os.path.join(save_path, "metadatas", f"test_dataset_rank{rank}.json")
+            os.makedirs(os.path.dirname(rank_json_path), exist_ok=True)
+            with open(rank_json_path, 'w', encoding='utf-8') as fout:
+                for item in logger.cached_results["test_dataset"]:
+                    fout.write(json.dumps(item, ensure_ascii=False) + '\n')
+
+    # Now simulate rank 0 merging all results
+    with (
+        patch('nemo.collections.speechlm2.parts.metrics.results_logger.get_rank', return_value=0),
+        patch('nemo.collections.speechlm2.parts.metrics.results_logger.get_world_size', return_value=world_size),
+        patch('torch.distributed.is_available', return_value=True),
+        patch('torch.distributed.is_initialized', return_value=True),
+        patch('torch.distributed.barrier'),
+        patch('torch.distributed.broadcast_object_list'),
+    ):  # Mock broadcast
+
+        logger = ResultsLogger(save_path=save_path)
+        # Manually set cached_results to simulate what rank 0 would have
+        logger.cached_results["test_dataset"] = []
+
+        # Call merge function directly
+        merged_results = logger._merge_rank_files("test_dataset")
+
+        # Verify all samples from all ranks are merged
+        assert len(merged_results) == world_size * 2  # 4 ranks * 2 samples each
+
+        # Verify sample IDs are from all ranks
+        sample_ids = [item["id"] for item in merged_results]
+        assert "sample_0" in sample_ids
+        assert "sample_1" in sample_ids
+        assert "sample_6" in sample_ids  # Last rank's samples
+        assert "sample_7" in sample_ids
+
+
+def test_results_logger_rank_file_wait(tmp_path):
+    """Test that rank 0 waits for other ranks' files"""
+    import json
+    import os
+    import threading
+    import time
+    from unittest.mock import patch
+
+    from nemo.collections.speechlm2.parts.metrics.results_logger import ResultsLogger
+
+    save_path = str(tmp_path / "wait_test")
+    os.makedirs(os.path.join(save_path, "metadatas"), exist_ok=True)
+    world_size = 2
+
+    # Create rank 0's file immediately
+    rank0_file = os.path.join(save_path, "metadatas", "test_dataset_rank0.json")
+    with open(rank0_file, 'w') as f:
+        json.dump({"id": "sample_0", "pred_text": "test0"}, f)
+        f.write('\n')
+
+    # Simulate rank 1's file appearing after a delay
+    def create_rank1_file_delayed():
+        time.sleep(1)  # 1 second delay
+        rank1_file = os.path.join(save_path, "metadatas", "test_dataset_rank1.json")
+        with open(rank1_file, 'w') as f:
+            json.dump({"id": "sample_1", "pred_text": "test1"}, f)
+            f.write('\n')
+
+    thread = threading.Thread(target=create_rank1_file_delayed)
+    thread.start()
+
+    # Rank 0 tries to merge - should wait for rank 1's file
+    with (
+        patch('nemo.collections.speechlm2.parts.metrics.results_logger.get_rank', return_value=0),
+        patch('nemo.collections.speechlm2.parts.metrics.results_logger.get_world_size', return_value=world_size),
+    ):
+
+        logger = ResultsLogger(save_path=save_path)
+        merged_results = logger._merge_rank_files("test_dataset")
+
+        # Should have results from both ranks
+        assert len(merged_results) == 2
+        sample_ids = [item["id"] for item in merged_results]
+        assert "sample_0" in sample_ids
+        assert "sample_1" in sample_ids
+
+    thread.join()
+
+
+def test_intelligibility():
+    metric = Intelligibility(pretrained_asr=None, verbose=False, reuse_asr_hyps=True)
+    metric.update(
+        name="dataset_1",
+        refs=["a b c d e f g h i j k l", "m n o p r s t u v"],
+        asr_hyps=["a b c d e f g h i j k l", "m n o p r s t u v"],
+        pred_audio=None,
+    )
+    metric.update(
+        name="dataset_2",
+        refs=["a b c"],
+        asr_hyps=["a b d"],
+        pred_audio=None,
+    )
+    ans = metric.compute()
+    # wer
+    assert ans["wer_dataset_1"] == 0.0
+    assert ans["wer_dataset_2"] == 1 / 3
+    assert ans["wer"] == 1 / 6  # average across datasets
+    # cer
+    assert ans["cer_dataset_1"] == 0.0
+    assert ans["cer_dataset_2"] == 1 / 5

tests/collections/speechlm2/test_nemotron_voicechat.py

@@ -0,0 +1,336 @@
+# Copyright (c) 2025, 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.
+
+import os
+
+import pytest
+import torch
+from lhotse import CutSet, SupervisionSegment
+from lhotse.testing.dummies import dummy_cut, dummy_recording
+
+from nemo.collections.common.data.utils import move_data_to_device
+from nemo.collections.speechlm2 import DuplexSTTDataset
+from nemo.collections.speechlm2.models import NemotronVoiceChat
+
+if torch.cuda.is_available():
+    torch.set_default_device('cuda')
+
+
+pretrained_llm = "TinyLlama/TinyLlama_v1.1"
+if os.path.exists("/home/TestData/speechlm/pretrained_models"):
+    pretrained_llm = "/home/TestData/speechlm/pretrained_models/TinyLlama--TinyLlama_v1.1"
+
+# STT sampling rate
+source_sample_rate = 16000
+# TTS sampling rate
+target_sample_rate = 22050
+
+
+def create_model(
+    predict_user_text=False,
+    force_use_noise_augmentation=False,
+    old_noise_prob=0.0,
+    old_noise_min_snr=0.0,
+    old_noise_max_snr=0.0,
+):
+    """Helper function to create a model with configurable settings."""
+    test_stt_cfg = {
+        "model": {
+            "pretrained_llm": pretrained_llm,
+            "pretrained_weights": False,
+            "audio_loss_weight": 1,
+            "text_loss_weight": 3,
+            "source_sample_rate": source_sample_rate,
+            "validation_save_path": "/tmp/test_duplex_stt_logs",
+            "perception": {
+                "_target_": "nemo.collections.speechlm2.modules.perception.AudioPerceptionModule",
+                "preprocessor": {
+                    "_target_": "nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor",
+                    "features": 80,
+                },
+                "encoder": {
+                    "_target_": "nemo.collections.asr.modules.ConformerEncoder",
+                    "feat_in": 80,
+                    "d_model": 512,
+                    "n_heads": 8,
+                    "n_layers": 1,
+                    "subsampling_factor": 8,
+                },
+                "modality_adapter": {
+                    "_target_": "nemo.collections.speechlm2.modules.perception.IdentityConnector",
+                    "d_model": 512,
+                },
+                "output_dim": 2048,
+            },
+            "predict_user_text": predict_user_text,
+            "force_use_noise_augmentation": force_use_noise_augmentation,
+            "old_noise_prob": old_noise_prob,
+            "old_noise_min_snr": old_noise_min_snr,
+            "old_noise_max_snr": old_noise_max_snr,
+            "optimizer": {"_target_": "torch.optim.AdamW"},
+        },
+        "data": {
+            "source_sample_rate": 16000,
+        },
+        "exp_manager": {
+            "explicit_log_dir": "/tmp/test_duplex_stt_logs",
+        },
+    }
+
+    test_tts_config = {
+        "model": {
+            "pretrained_lm_name": pretrained_llm,
+            "pretrained_ae_dir": None,
+            "pretrained_tts_model": None,
+            "scoring_asr": "stt_en_fastconformer_transducer_large",
+            "freeze_params": [
+                r"^audio_codec\..+$",  # Keep audio codec frozen as it only provides supervision for training.
+                r"^embed_tokens\..+$",  # Keep embed_tokens frozen as done in eartts
+            ],
+            "bos_token": "<s>",
+            "eos_token": "</s>",
+            "pad_token": "<SPECIAL_12>",
+            "audio_codec_run_dtype": "float32",
+            "prevent_freeze_params": [],
+            "audio_save_path": "",
+            "inference_guidance_scale": 0.5,
+            "inference_noise_scale": 0.8,
+            "inference_top_p_or_k": 0.8,
+            "inference_guidance_enabled": False,
+            "subword_mask_exactly_as_eartts": False,
+            "context_hidden_mask_exactly_as_eartts": False,
+            "optimizer": {
+                "_target_": "torch.optim.AdamW",
+                "lr": 4e-5,
+                "betas": [0.9, 0.98],
+                "weight_decay": 0,
+                "foreach": True,
+            },
+            "lr_scheduler": {
+                "_target_": "nemo.core.optim.lr_scheduler.InverseSquareRootAnnealing",
+                "warmup_steps": 2500,
+                "min_lr": 1e-6,
+                "max_steps": 100_000_000,
+            },
+            "codec_config": {
+                "latent_size": 512,
+                "n_fft": 16,
+                "hop_length": 4,
+                "base_hidden_size": 384,
+                "channel_mult": [1, 2, 4],
+                "rates": [7, 7, 9],
+                "num_blocks": 3,
+                "kernel_size": 7,
+                "groups": 1,
+                "codebook_size": 1024,
+                "num_quantizers": 31,
+                "wav_to_token_ratio": 1764,
+            },
+            "tts_config": {
+                "use_gated_fusion_for_text_audio": True,
+                "disable_eos_prediction": True,
+                "use_bos_eos_emb": True,
+                "use_subword_flag_emb": True,
+                "num_delay_speech_tokens": 2,
+                "backbone_type": "gemma3_text",
+                "backbone_model_class": None,
+                "backbone_config_class": None,
+                "backbone_config": {
+                    "hidden_size": 1152,
+                    "intermediate_size": 4608,
+                    "num_hidden_layers": 1,
+                    "num_attention_heads": 16,
+                    "num_key_value_heads": 16,
+                    "head_dim": 72,
+                    "attention_dropout": 0.1,
+                    "use_cache": False,
+                },
+                "latent_size": 512,
+                "codebook_size": 1024,
+                "num_quantizers": 31,
+                "context_hidden_size": None,
+                "cas_config": {
+                    "backbone_type": "t5gemma",
+                    "backbone_model_class": None,
+                    "backbone_config_class": None,
+                    "backbone_config": {
+                        "is_encoder_decoder": False,
+                        "encoder": {
+                            "hidden_size": 1152,
+                            "intermediate_size": 4608,
+                            "num_hidden_layers": 1,
+                            "num_attention_heads": 16,
+                            "num_key_value_heads": 16,
+                            "head_dim": 72,
+                            "use_cache": False,
+                            "attention_dropout": 0.1,
+                        },
+                    },
+                },
+                "mog_head_config": {
+                    "intermediate_size": 4608,
+                    "num_layers": 3,
+                    "low_rank": 64,
+                    "num_predictions": 1024,
+                    "min_log_std": -4.0,
+                    "eps": 1e-6,
+                },
+                "p_uncond": 0.1,
+                "label_smoothing": 0.01,
+                "max_training_rate": 0.8,
+                "quantizer_dropout": 0.5,
+                "random_target_masking": False,
+                "exponent": 3.0,
+            },
+        },
+        "data": {
+            "add_text_bos_and_eos_in_each_turn": True,
+            "add_audio_prompt": True,
+            "audio_prompt_duration": 3.0,
+            "frame_length": 0.08,
+            "source_sample_rate": source_sample_rate,
+            "target_sample_rate": target_sample_rate,
+        },
+        "exp_manager": {
+            "explicit_log_dir": "/tmp/test_duplex_stt_logs",
+        },
+    }
+
+    test_config = {
+        "model": {
+            "scoring_asr": "stt_en_fastconformer_transducer_large",
+            "stt": test_stt_cfg,
+            "speech_generation": test_tts_config,
+        },
+        "data": {
+            "frame_length": 0.08,
+            "source_sample_rate": source_sample_rate,
+            "target_sample_rate": target_sample_rate,
+            "input_roles": ["user", "User"],
+            "output_roles": ["agent", "Assistant", "assistant", "Agent"],
+        },
+        "exp_manager": {
+            "explicit_log_dir": "/tmp/test_nemotron_voicechat_logs",
+        },
+    }
+    model = NemotronVoiceChat(test_config)
+    if torch.cuda.is_available():
+        model.to("cuda")
+    return model
+
+
+@pytest.fixture(scope="session")
+def model():
+    return create_model(predict_user_text=False)
+
+
+@pytest.fixture(scope="session")
+def dataset(model):
+    return DuplexSTTDataset(
+        model.stt_model.tokenizer,
+        frame_length=0.08,
+        source_sample_rate=source_sample_rate,
+        input_roles=["user"],
+        output_roles=["assistant"],
+    )
+
+
+@pytest.fixture(scope="session")
+def training_cutset_batch():
+    cut = dummy_cut(0, recording=dummy_recording(0, with_data=True, duration=1.0, sampling_rate=22050))
+    cut.target_audio = dummy_recording(1, with_data=True, duration=1.0, sampling_rate=22050)
+    cut.supervisions = [
+        SupervisionSegment(
+            id=cut.id,
+            recording_id=cut.recording_id,
+            start=0,
+            duration=0.1,
+            text='hi',
+            speaker="user",
+        ),
+        SupervisionSegment(
+            id=cut.id,
+            recording_id=cut.recording_id,
+            start=0.3,
+            duration=0.1,
+            text='hello',
+            speaker="assistant",
+        ),
+        SupervisionSegment(
+            id=cut.id,
+            recording_id=cut.recording_id,
+            start=0.5,
+            duration=0.1,
+            text='ok',
+            speaker="user",
+        ),
+        SupervisionSegment(
+            id=cut.id,
+            recording_id=cut.recording_id,
+            start=0.6,
+            duration=0.1,
+            text='okay',
+            speaker="assistant",
+        ),
+    ]
+    return CutSet([cut])
+
+
+def test_e2e_validation_step(model, dataset, training_cutset_batch):
+    model.eval()
+    model.on_validation_epoch_start()
+    batch = dataset[training_cutset_batch]
+    batch = move_data_to_device(batch, device=model.device)
+    results = model.validation_step(
+        {"dummy_val_set": batch},
+        batch_idx=0,
+        speaker_audio=torch.randn(1, 22050, device=model.device),
+        speaker_audio_lens=torch.tensor([22050], device=model.device),
+    )
+    assert results is None  # no return value
+
+
+def test_e2s_offline_generation(model):
+    model.eval()
+    # 16000 samples == 1 second == 12.5 frames ~= 14 frames after encoder padding
+    ans = model.offline_inference(
+        input_signal=torch.randn(1, 16000, device=model.device),
+        input_signal_lens=torch.tensor([16000], device=model.device),
+        speaker_audio=torch.randn(1, 22050, device=model.device),
+        speaker_audio_lens=torch.tensor([22050], device=model.device),
+    )
+
+    assert ans.keys() == {
+        'text',
+        'src_text',
+        'tokens_text_src',
+        'tokens_text',
+        'tokens_len',
+        'source_audio',
+        'source_audio_len',
+        "audio",
+        "audio_len",
+    }
+
+    assert isinstance(ans["text"], list)
+    assert isinstance(ans["text"][0], str)
+
+    gen_text = ans["tokens_text"]
+    assert gen_text.shape == (1, 14)
+    assert gen_text.dtype == torch.long
+    assert (gen_text >= 0).all()
+    assert (gen_text < model.stt_model.text_vocab_size).all()
+    # 14 tokens = 24696 audio frames
+    gen_audio = ans["audio"]
+    assert gen_audio.shape == (1, 24696)

tests/collections/speechlm2/test_parallel.py

@@ -0,0 +1,305 @@
+# 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.
+
+import pytest
+from lightning.pytorch.strategies.model_parallel import ModelParallelStrategy
+from nemo_automodel.components.distributed.config import FSDP2Config
+from nemo_automodel.components.moe.config import MoEParallelizerConfig
+from omegaconf import DictConfig
+
+from nemo.collections.speechlm2.parts.parallel import AutomodelParallelStrategy
+from nemo.utils.trainer_utils import _resolve_automodel_configs, resolve_trainer_cfg
+
+
+# ---------------------------------------------------------------------------
+# AutomodelParallelStrategy
+# ---------------------------------------------------------------------------
+
+
+class TestAutomodelParallelStrategy:
+    def test_is_subclass_of_model_parallel_strategy(self):
+        assert issubclass(AutomodelParallelStrategy, ModelParallelStrategy)
+
+    def test_default_init(self):
+        strategy = AutomodelParallelStrategy()
+        assert strategy._dp_size is None
+        assert strategy._dp_replicate_size is None
+        assert strategy._tp_size == 1
+        assert strategy._pp_size == 1
+        assert strategy._cp_size == 1
+        assert strategy._ep_size == 1
+        assert strategy._distributed_config is None
+        assert strategy._moe_config is None
+        assert strategy._moe_mesh is None
+        assert strategy.activation_checkpointing_llm is False
+        assert strategy.activation_checkpointing_perception is False
+
+    def test_accepts_activation_checkpointing_flags(self):
+        strategy = AutomodelParallelStrategy(
+            activation_checkpointing_llm=True,
+            activation_checkpointing_perception=True,
+        )
+        assert strategy.activation_checkpointing_llm is True
+        assert strategy.activation_checkpointing_perception is True
+
+    def test_activation_checkpointing_flags_are_independent(self):
+        """Each AC flag can be set without the other."""
+        llm_only = AutomodelParallelStrategy(activation_checkpointing_llm=True)
+        assert llm_only.activation_checkpointing_llm is True
+        assert llm_only.activation_checkpointing_perception is False
+
+        perception_only = AutomodelParallelStrategy(activation_checkpointing_perception=True)
+        assert perception_only.activation_checkpointing_llm is False
+        assert perception_only.activation_checkpointing_perception is True
+
+    def test_custom_parallelism_sizes(self):
+        strategy = AutomodelParallelStrategy(
+            dp_size=4,
+            dp_replicate_size=2,
+            tp_size=2,
+            pp_size=2,
+            cp_size=2,
+            ep_size=4,
+        )
+        assert strategy._dp_size == 4
+        assert strategy._dp_replicate_size == 2
+        assert strategy._tp_size == 2
+        assert strategy._pp_size == 2
+        assert strategy._cp_size == 2
+        assert strategy._ep_size == 4
+
+    def test_accepts_distributed_config(self):
+        cfg = FSDP2Config(sequence_parallel=True, defer_fsdp_grad_sync=False)
+        strategy = AutomodelParallelStrategy(distributed_config=cfg)
+        assert strategy.distributed_config is cfg
+        assert strategy.distributed_config.sequence_parallel is True
+        assert strategy.distributed_config.defer_fsdp_grad_sync is False
+
+    def test_accepts_moe_config(self):
+        cfg = MoEParallelizerConfig()
+        strategy = AutomodelParallelStrategy(moe_config=cfg)
+        assert strategy.moe_config is cfg
+
+    def test_save_distributed_checkpoint_forwarded(self):
+        strategy = AutomodelParallelStrategy(save_distributed_checkpoint=False)
+        assert strategy._save_distributed_checkpoint is False
+
+    def test_moe_mesh_initially_none(self):
+        strategy = AutomodelParallelStrategy()
+        assert strategy.moe_mesh is None
+
+    def test_device_mesh_raises_before_setup(self):
+        strategy = AutomodelParallelStrategy()
+        with pytest.raises(RuntimeError):
+            _ = strategy.device_mesh
+
+    def test_distributed_sampler_kwargs_raises_before_setup(self):
+        strategy = AutomodelParallelStrategy()
+        with pytest.raises(RuntimeError):
+            _ = strategy.distributed_sampler_kwargs
+
+
+# ---------------------------------------------------------------------------
+# _resolve_automodel_configs
+# ---------------------------------------------------------------------------
+
+
+class TestResolveAutomodelConfigs:
+    """Tests for _resolve_automodel_configs which operates on an instantiated strategy object."""
+
+    def test_plain_dict_to_fsdp2_config(self):
+        strategy = AutomodelParallelStrategy(
+            distributed_config={"defer_fsdp_grad_sync": False, "sequence_parallel": True},
+        )
+        _resolve_automodel_configs(strategy)
+        assert isinstance(strategy.distributed_config, FSDP2Config)
+        assert strategy.distributed_config.defer_fsdp_grad_sync is False
+        assert strategy.distributed_config.sequence_parallel is True
+        # Check that __post_init__ created the default mp_policy
+        assert strategy.distributed_config.mp_policy is not None
+
+    def test_plain_dict_to_moe_config(self):
+        strategy = AutomodelParallelStrategy(
+            moe_config={"reshard_after_forward": True},
+        )
+        _resolve_automodel_configs(strategy)
+        assert isinstance(strategy.moe_config, MoEParallelizerConfig)
+        assert strategy.moe_config.reshard_after_forward is True
+
+    def test_noop_when_no_configs(self):
+        """Strategy without configs (e.g. DDPStrategy) is untouched."""
+        from lightning.pytorch.strategies import DDPStrategy
+
+        strategy = DDPStrategy()
+        _resolve_automodel_configs(strategy)  # should not raise
+
+    def test_noop_when_already_objects(self):
+        original_cfg = FSDP2Config()
+        original_moe = MoEParallelizerConfig()
+        strategy = AutomodelParallelStrategy(
+            distributed_config=original_cfg,
+            moe_config=original_moe,
+        )
+        _resolve_automodel_configs(strategy)
+        assert strategy.distributed_config is original_cfg
+        assert strategy.moe_config is original_moe
+
+    def test_noop_when_configs_are_none(self):
+        strategy = AutomodelParallelStrategy()
+        _resolve_automodel_configs(strategy)
+        assert strategy.distributed_config is None
+        assert strategy.moe_config is None
+
+    def test_empty_dict_creates_default_configs(self):
+        strategy = AutomodelParallelStrategy(distributed_config={}, moe_config={})
+        _resolve_automodel_configs(strategy)
+        assert isinstance(strategy.distributed_config, FSDP2Config)
+        assert isinstance(strategy.moe_config, MoEParallelizerConfig)
+        # All defaults should apply
+        assert strategy.distributed_config.defer_fsdp_grad_sync is True
+        assert strategy.distributed_config.sequence_parallel is False
+
+    def test_nested_target_in_distributed_config(self):
+        """A sub-field like mp_policy can use _target_ for Hydra instantiation."""
+        strategy = AutomodelParallelStrategy(
+            distributed_config={
+                "sequence_parallel": True,
+                "mp_policy": {
+                    "_target_": "torch.distributed.fsdp.MixedPrecisionPolicy",
+                    "param_dtype": "torch.bfloat16",
+                },
+            },
+        )
+        _resolve_automodel_configs(strategy)
+        assert isinstance(strategy.distributed_config, FSDP2Config)
+        assert strategy.distributed_config.sequence_parallel is True
+        from torch.distributed.fsdp import MixedPrecisionPolicy
+
+        assert isinstance(strategy.distributed_config.mp_policy, MixedPrecisionPolicy)
+
+    def test_both_configs_resolved_together(self):
+        strategy = AutomodelParallelStrategy(
+            distributed_config={"sequence_parallel": False},
+            moe_config={},
+        )
+        _resolve_automodel_configs(strategy)
+        assert isinstance(strategy.distributed_config, FSDP2Config)
+        assert isinstance(strategy.moe_config, MoEParallelizerConfig)
+
+
+# ---------------------------------------------------------------------------
+# resolve_trainer_cfg (end-to-end with AutomodelParallelStrategy)
+# ---------------------------------------------------------------------------
+
+
+class TestResolveTrainerCfg:
+    def test_automodel_strategy_from_yaml(self):
+        """End-to-end: YAML dict config -> instantiated AutomodelParallelStrategy."""
+        trainer_cfg = DictConfig(
+            {
+                "devices": 1,
+                "accelerator": "cpu",
+                "strategy": {
+                    "_target_": "nemo.collections.speechlm2.parts.parallel.AutomodelParallelStrategy",
+                    "tp_size": 2,
+                    "pp_size": 1,
+                    "distributed_config": {
+                        "sequence_parallel": True,
+                        "defer_fsdp_grad_sync": False,
+                    },
+                    "moe_config": {},
+                },
+            }
+        )
+        resolved = resolve_trainer_cfg(trainer_cfg)
+        strategy = resolved["strategy"]
+        assert isinstance(strategy, AutomodelParallelStrategy)
+        assert strategy._tp_size == 2
+        assert strategy._pp_size == 1
+        assert isinstance(strategy.distributed_config, FSDP2Config)
+        assert strategy.distributed_config.sequence_parallel is True
+        assert strategy.distributed_config.defer_fsdp_grad_sync is False
+        assert isinstance(strategy.moe_config, MoEParallelizerConfig)
+
+    def test_non_automodel_strategy_unaffected(self):
+        """Other strategies (e.g. DDPStrategy) should pass through unchanged."""
+        trainer_cfg = DictConfig(
+            {
+                "devices": 1,
+                "accelerator": "cpu",
+                "strategy": {
+                    "_target_": "lightning.pytorch.strategies.DDPStrategy",
+                    "gradient_as_bucket_view": True,
+                    "find_unused_parameters": True,
+                },
+            }
+        )
+        resolved = resolve_trainer_cfg(trainer_cfg)
+        from lightning.pytorch.strategies import DDPStrategy
+
+        assert isinstance(resolved["strategy"], DDPStrategy)
+
+    def test_string_strategy_unaffected(self):
+        """A plain string strategy (e.g. 'ddp') should pass through."""
+        trainer_cfg = DictConfig(
+            {
+                "devices": 1,
+                "accelerator": "cpu",
+                "strategy": "ddp",
+            }
+        )
+        resolved = resolve_trainer_cfg(trainer_cfg)
+        assert resolved["strategy"] == "ddp"
+
+    def test_automodel_strategy_with_ac_flags_from_yaml(self):
+        """Both AC flags can be set via YAML and survive Hydra instantiation."""
+        trainer_cfg = DictConfig(
+            {
+                "devices": 1,
+                "accelerator": "cpu",
+                "strategy": {
+                    "_target_": "nemo.collections.speechlm2.parts.parallel.AutomodelParallelStrategy",
+                    "ep_size": 4,
+                    "activation_checkpointing_llm": True,
+                    "activation_checkpointing_perception": True,
+                },
+            }
+        )
+        resolved = resolve_trainer_cfg(trainer_cfg)
+        strategy = resolved["strategy"]
+        assert isinstance(strategy, AutomodelParallelStrategy)
+        assert strategy.activation_checkpointing_llm is True
+        assert strategy.activation_checkpointing_perception is True
+
+    def test_automodel_strategy_without_configs(self):
+        """AutomodelParallelStrategy can be specified without distributed_config/moe_config."""
+        trainer_cfg = DictConfig(
+            {
+                "devices": 1,
+                "accelerator": "cpu",
+                "strategy": {
+                    "_target_": "nemo.collections.speechlm2.parts.parallel.AutomodelParallelStrategy",
+                    "tp_size": 4,
+                    "ep_size": 8,
+                },
+            }
+        )
+        resolved = resolve_trainer_cfg(trainer_cfg)
+        strategy = resolved["strategy"]
+        assert isinstance(strategy, AutomodelParallelStrategy)
+        assert strategy._tp_size == 4
+        assert strategy._ep_size == 8
+        # No configs passed → still None (will be defaulted in setup_environment)
+        assert strategy._distributed_config is None
+        assert strategy._moe_config is None

tests/collections/speechlm2/test_perception_activation_checkpointing.py

@@ -0,0 +1,241 @@
+# 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.
+
+"""Unit tests for perception-module helpers.
+
+These tests exercise:
+  * ``_set_encoder_activation_checkpointing`` (the module-level helper that
+    wraps ``encoder.pre_encode`` and each ``encoder.layers[i]``).
+  * ``AudioPerceptionModule.set_activation_checkpointing`` and
+    ``AudioTranscriptionPerceptionModule.set_activation_checkpointing`` —
+    both are thin delegators, so we assert that each class exposes the method
+    and that it routes through the helper.
+"""
+
+import pytest
+import torch
+from torch import nn
+from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import CheckpointWrapper
+
+from nemo.collections.speechlm2.modules.perception import (
+    AudioPerceptionModule,
+    AudioTranscriptionPerceptionModule,
+    _set_encoder_activation_checkpointing,
+)
+
+
+class _FakeConvSubsampling(nn.Module):
+    """Stand-in for ``ConvSubsampling``/``StackingSubsampling``: non-Linear, so
+    the helper should wrap it."""
+
+    def __init__(self):
+        super().__init__()
+        self.conv = nn.Conv1d(1, 1, kernel_size=1)
+
+    def forward(self, x, lengths):  # pragma: no cover — forward not exercised
+        return x, lengths
+
+
+class _FakeConformerLayer(nn.Module):
+    def __init__(self, dim: int = 4):
+        super().__init__()
+        self.linear = nn.Linear(dim, dim)
+
+    def forward(self, x):  # pragma: no cover — forward not exercised
+        return self.linear(x)
+
+
+class _FakeConformerEncoder(nn.Module):
+    """Mimics just the shape the helper inspects: ``.pre_encode`` and
+    ``.layers``. Good enough to test wrapping without pulling the real
+    ``ConformerEncoder``, which needs heavy deps."""
+
+    def __init__(self, num_layers: int = 3, pre_encode: nn.Module | None = None):
+        super().__init__()
+        if pre_encode is None:
+            pre_encode = _FakeConvSubsampling()
+        self.pre_encode = pre_encode
+        self.layers = nn.ModuleList([_FakeConformerLayer() for _ in range(num_layers)])
+
+
+class _FakePreprocessor(nn.Module):
+    def forward(self, input_signal, length):
+        return input_signal, length
+
+
+class _FakeUnsupportedEncoder(nn.Module):
+    def forward(self, audio_signal, length):
+        return audio_signal, length
+
+
+class _FakeIdentityAdapter(nn.Module):
+    def forward(self, audio_signal, length):
+        return audio_signal, length
+
+
+# ---------------------------------------------------------------------------
+# _set_encoder_activation_checkpointing
+# ---------------------------------------------------------------------------
+
+
+class TestSetEncoderActivationCheckpointing:
+    def test_noop_when_disabled(self):
+        encoder = _FakeConformerEncoder(num_layers=3)
+        original_pre_encode = encoder.pre_encode
+        original_layers = [encoder.layers[i] for i in range(len(encoder.layers))]
+
+        _set_encoder_activation_checkpointing(encoder, enabled=False)
+
+        # Nothing should be wrapped.
+        assert encoder.pre_encode is original_pre_encode
+        assert not isinstance(encoder.pre_encode, CheckpointWrapper)
+        for i, layer in enumerate(original_layers):
+            assert encoder.layers[i] is layer
+            assert not isinstance(encoder.layers[i], CheckpointWrapper)
+
+    def test_wraps_pre_encode_and_all_layers_when_enabled(self):
+        encoder = _FakeConformerEncoder(num_layers=4)
+
+        _set_encoder_activation_checkpointing(encoder, enabled=True)
+
+        assert isinstance(encoder.pre_encode, CheckpointWrapper)
+        assert len(encoder.layers) == 4
+        for layer in encoder.layers:
+            assert isinstance(layer, CheckpointWrapper)
+
+    def test_skips_linear_pre_encode(self):
+        """``ConformerEncoder.forward`` dispatches on ``isinstance(pre_encode,
+        nn.Linear)`` — wrapping would hide the type and route the call to the
+        wrong branch. The helper must leave Linear pre_encode untouched even
+        when ``enabled=True``."""
+        encoder = _FakeConformerEncoder(num_layers=2, pre_encode=nn.Linear(8, 8))
+
+        _set_encoder_activation_checkpointing(encoder, enabled=True)
+
+        # Linear pre_encode unchanged …
+        assert isinstance(encoder.pre_encode, nn.Linear)
+        assert not isinstance(encoder.pre_encode, CheckpointWrapper)
+        # … but layers are still wrapped.
+        for layer in encoder.layers:
+            assert isinstance(layer, CheckpointWrapper)
+
+    def test_noop_when_encoder_has_no_pre_encode(self):
+        """Graceful degradation: a non-Conformer encoder with ``.layers`` but
+        no ``.pre_encode`` should still have its layers wrapped."""
+
+        class EncoderNoPreEncode(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.layers = nn.ModuleList([_FakeConformerLayer() for _ in range(2)])
+
+        encoder = EncoderNoPreEncode()
+        _set_encoder_activation_checkpointing(encoder, enabled=True)
+        for layer in encoder.layers:
+            assert isinstance(layer, CheckpointWrapper)
+
+    def test_noop_when_encoder_has_no_layers(self):
+        """Graceful degradation: an encoder with ``.pre_encode`` but no
+        ``.layers`` should still have its pre_encode wrapped."""
+
+        class EncoderNoLayers(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.pre_encode = _FakeConvSubsampling()
+
+        encoder = EncoderNoLayers()
+        _set_encoder_activation_checkpointing(encoder, enabled=True)
+        assert isinstance(encoder.pre_encode, CheckpointWrapper)
+
+    def test_noop_when_encoder_has_neither(self):
+        """Foreign architecture with no recognisable structure: helper is a
+        no-op and does not raise."""
+
+        class UnrelatedEncoder(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.something = nn.Linear(4, 4)
+
+        encoder = UnrelatedEncoder()
+        _set_encoder_activation_checkpointing(encoder, enabled=True)  # must not raise
+        assert not isinstance(encoder.something, CheckpointWrapper)
+
+
+# ---------------------------------------------------------------------------
+# set_activation_checkpointing method bindings
+# ---------------------------------------------------------------------------
+
+
+class TestPerceptionSetActivationCheckpointingMethod:
+    """Verify the method exists and routes through the helper on both
+    perception classes. We don't instantiate the full module (that requires
+    loading an ASR checkpoint) — we inspect the class attribute and call it
+    on a lightweight fake ``self`` that exposes only the bits the method
+    touches (``self.encoder``)."""
+
+    def test_audio_perception_module_has_method(self):
+        assert callable(getattr(AudioPerceptionModule, "set_activation_checkpointing", None))
+
+    def test_audio_transcription_perception_module_has_method(self):
+        assert callable(getattr(AudioTranscriptionPerceptionModule, "set_activation_checkpointing", None))
+
+    def test_method_wraps_encoder_layers(self):
+        """Invoke the unbound method against a fake ``self`` that only has
+        ``self.encoder``. This is enough to verify that the method delegates
+        to ``_set_encoder_activation_checkpointing`` against the encoder
+        exposed via ``self.encoder``."""
+
+        class FakeSelf:
+            pass
+
+        fake = FakeSelf()
+        fake.encoder = _FakeConformerEncoder(num_layers=2)
+
+        AudioPerceptionModule.set_activation_checkpointing(fake, enabled=True)
+
+        assert isinstance(fake.encoder.pre_encode, CheckpointWrapper)
+        for layer in fake.encoder.layers:
+            assert isinstance(layer, CheckpointWrapper)
+
+    def test_method_disabled_is_noop(self):
+        class FakeSelf:
+            pass
+
+        fake = FakeSelf()
+        fake.encoder = _FakeConformerEncoder(num_layers=2)
+        original_pre_encode = fake.encoder.pre_encode
+        original_layers = list(fake.encoder.layers)
+
+        AudioPerceptionModule.set_activation_checkpointing(fake, enabled=False)
+
+        assert fake.encoder.pre_encode is original_pre_encode
+        for i, layer in enumerate(original_layers):
+            assert fake.encoder.layers[i] is layer
+
+
+class TestAudioPerceptionSpeakerTargets:
+    def test_raises_when_spk_targets_are_passed_to_unsupported_encoder(self):
+        module = AudioPerceptionModule.__new__(AudioPerceptionModule)
+        nn.Module.__init__(module)
+        module.preprocessor = _FakePreprocessor()
+        module.encoder = _FakeUnsupportedEncoder()
+        module.modality_adapter = _FakeIdentityAdapter()
+        module.proj = nn.Identity()
+        module.spec_augmentation = None
+
+        with pytest.raises(ValueError, match="spk_targets.*does not support"):
+            module(
+                input_signal=torch.zeros(1, 2, 4),
+                input_signal_length=torch.tensor([4]),
+                spk_targets=torch.zeros(1, 4, 2),
+            )

tests/collections/speechlm2/test_role_swap.py

@@ -0,0 +1,182 @@
+# Copyright (c) 2025, 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 io import BytesIO
+from unittest.mock import MagicMock, PropertyMock
+
+import numpy as np
+import pytest
+import soundfile as sf
+from lhotse import AudioSource, MonoCut, Recording, SupervisionSegment
+
+from nemo.collections.speechlm2.data import DuplexSTTDataset
+
+
+def _make_in_memory_recording(recording_id, duration, sampling_rate, signal_freq=440.0):
+    """Create a lhotse Recording backed by in-memory WAV bytes with a sine wave."""
+    num_samples = int(duration * sampling_rate)
+    t = np.linspace(0, duration, num_samples, endpoint=False, dtype=np.float32)
+    audio = 0.5 * np.sin(2 * np.pi * signal_freq * t)
+    buf = BytesIO()
+    sf.write(buf, audio, sampling_rate, format='wav')
+    buf.seek(0)
+    return Recording(
+        id=recording_id,
+        sampling_rate=sampling_rate,
+        num_samples=num_samples,
+        duration=duration,
+        sources=[AudioSource(type="memory", channels=[0], source=buf.getvalue())],
+    )
+
+
+def test_create_role_swapped_cut():
+    """
+    Test _create_role_swapped_cut with a 6-turn conversation.
+
+    Original:
+        Turn 1: User       0.0-1.0  "hello"
+        Turn 2: Assistant  1.5-2.5  "hi there"
+        Turn 3: User       3.0-4.0  "how are you"
+        Turn 4: Assistant  4.5-5.5  "doing well"
+        Turn 5: User       6.0-7.0  "great"
+        Turn 6: Assistant  7.5-8.5  "bye"
+
+    After swap + filter (remove first Assistant, last User):
+        User       0.0-1.0  "hi there"     (originally Assistant turn 2)
+        Assistant  1.5-2.5  "how are you"  (originally User turn 3)
+        User       3.0-4.0  "doing well"   (originally Assistant turn 4)
+        Assistant  4.5-5.5  "great"        (originally User turn 5)
+    """
+    tokenizer = MagicMock()
+    type(tokenizer).bos = PropertyMock(return_value=1)
+    type(tokenizer).eos = PropertyMock(return_value=2)
+    type(tokenizer).pad = PropertyMock(return_value=0)
+    type(tokenizer).pad_id = PropertyMock(return_value=0)
+    type(tokenizer).unk_id = PropertyMock(return_value=None)
+    tokenizer.text_to_ids = MagicMock(return_value=[1])
+
+    dataset = DuplexSTTDataset(
+        tokenizer=tokenizer,
+        frame_length=0.08,
+        source_sample_rate=16000,
+        input_roles=["User"],
+        output_roles=["Assistant"],
+        aug_by_swap_role=True,
+        cfg={"early_interruption_prob": 0.0},
+        model_cfg={"predict_user_text": False, "force_align_user_text": False},
+    )
+
+    # Build a 6-turn cut with source (440Hz) and target (880Hz) audio
+    sampling_rate = 16000
+    total_duration = 9.0
+    source_rec = _make_in_memory_recording("src", total_duration, sampling_rate, signal_freq=440.0)
+    target_rec = _make_in_memory_recording("tgt", total_duration, sampling_rate, signal_freq=880.0)
+
+    supervisions = [
+        SupervisionSegment(id="s1", recording_id="src", start=0.0, duration=1.0, speaker="User", text="hello"),
+        SupervisionSegment(id="s2", recording_id="src", start=1.5, duration=1.0, speaker="Assistant", text="hi there"),
+        SupervisionSegment(id="s3", recording_id="src", start=3.0, duration=1.0, speaker="User", text="how are you"),
+        SupervisionSegment(
+            id="s4", recording_id="src", start=4.5, duration=1.0, speaker="Assistant", text="doing well"
+        ),
+        SupervisionSegment(id="s5", recording_id="src", start=6.0, duration=1.0, speaker="User", text="great"),
+        SupervisionSegment(id="s6", recording_id="src", start=7.5, duration=1.0, speaker="Assistant", text="bye"),
+    ]
+
+    cut = MonoCut(
+        id="test_cut",
+        start=0,
+        duration=total_duration,
+        channel=0,
+        supervisions=supervisions,
+        recording=source_rec,
+        custom={'target_audio': target_rec, 'total_turns': 6},
+    )
+
+    swapped_cut = dataset._create_role_swapped_cut(cut)
+    assert swapped_cut is not None
+
+    sups = swapped_cut.supervisions
+
+    # 4 turns remain after removing first Assistant and last User
+    assert len(sups) == 4
+
+    # Roles alternate User, Assistant, User, Assistant
+    assert [s.speaker for s in sups] == ["User", "Assistant", "User", "Assistant"]
+
+    # Texts match the kept turns
+    assert [s.text for s in sups] == ["hi there", "how are you", "doing well", "great"]
+
+    # Timestamps adjusted (first_remaining_start=1.5 subtracted), all start at 0
+    assert abs(sups[0].start - 0.0) < 1e-6
+    assert abs(sups[1].start - 1.5) < 1e-6
+    assert abs(sups[2].start - 3.0) < 1e-6
+    assert abs(sups[3].start - 4.5) < 1e-6
+
+    # Source audio: User regions have signal (from target_audio), Assistant regions are silent
+    audio = swapped_cut.recording.to_cut().load_audio().squeeze()
+    for s in sups:
+        start_sample = int(s.start * sampling_rate)
+        end_sample = int((s.start + s.duration) * sampling_rate)
+        rms = np.sqrt(np.mean(audio[start_sample:end_sample] ** 2))
+        if s.speaker == 'User':
+            assert rms > 0.1, f"User turn '{s.text}' should have audio from target_audio, RMS={rms:.4f}"
+        else:
+            assert rms < 0.01, f"Assistant turn '{s.text}' should be silent in STT source, RMS={rms:.4f}"
+
+    # Verify User audio is similar to target audio (880Hz)
+    # Load the original target audio for comparison
+    target_audio = target_rec.to_cut().load_audio().squeeze()
+
+    for s in sups:
+        if s.speaker == 'User':
+            # User turns should be from target_audio (originally Assistant turns)
+            # The original supervisions had these start times: 1.5 (hi there), 4.5 (doing well)
+            # which map to indices 0 and 2 in sups (User turns after swap)
+            start_sample = int(s.start * sampling_rate)
+            end_sample = int((s.start + s.duration) * sampling_rate)
+            segment = audio[start_sample:end_sample]
+            idx = list(sups).index(s)
+
+            original_start_times = {0: 1.5, 2: 4.5}  # map swapped index to original time
+            if idx in original_start_times:
+                orig_start = original_start_times[idx]
+                orig_start_sample = int(orig_start * sampling_rate)
+                orig_end_sample = int((orig_start + s.duration) * sampling_rate)
+                target_segment = target_audio[orig_start_sample:orig_end_sample]
+
+                # Verify lengths match exactly
+                len_diff = abs(len(segment) - len(target_segment))
+                assert len_diff == 0, (
+                    f"User turn '{s.text}' length must match exactly: "
+                    f"segment={len(segment)} samples, "
+                    f"target_segment={len(target_segment)} samples, "
+                    f"diff={len_diff}"
+                )
+
+                # Compute correlation between user audio and target audio
+                # Waveforms should be nearly identical
+                correlation = np.corrcoef(segment, target_segment)[0, 1]
+                assert correlation > 0.999, (
+                    f"User turn '{s.text}' should have nearly identical waveform to target audio. "
+                    f"Correlation={correlation:.6f}, expected > 0.999"
+                )
+
+    # Metadata
+    assert swapped_cut.custom['role_swapped'] is True
+    assert swapped_cut.custom['total_turns'] == 4
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

tests/collections/speechlm2/test_rote.py

@@ -0,0 +1,83 @@
+# Copyright (c) 2025, 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.
+
+import math
+
+import pytest
+import torch
+
+from nemo.collections.speechlm2.modules.rote import RotaryTimeEmbedding, rotate_half
+
+
+def test_defaults():
+    """Defaults follow the Audio Flamingo Next audio-encoder config (theta=1200, partial_rotary_factor=0.2)."""
+    dim = 100
+    rote = RotaryTimeEmbedding(dim)
+    assert rote.theta == 1200.0
+    assert rote.rotary_fraction == 0.2
+    assert rote.rotary_dim == 20
+    assert rote.inv_freq.shape == (10,)
+
+    # Rotated width is rounded down to an even number: int(10 * 0.5) = 5 -> 4.
+    assert RotaryTimeEmbedding(dim=10, rotary_fraction=0.5).rotary_dim == 4
+
+    # rotary_fraction=1.0 recovers full-width rotation.
+    full = RotaryTimeEmbedding(dim, rotary_fraction=1.0)
+    assert full.rotary_dim == dim
+    assert full.inv_freq.shape == (dim // 2,)
+
+
+def test_rotate_half():
+    """GPT-J convention: ``[x0, x1, x2, x3] -> [-x1, x0, -x3, x2]``."""
+    x = torch.tensor([[1.0, 2.0, 3.0, 4.0]])
+    expected = torch.tensor([[-2.0, 1.0, -4.0, 3.0]])
+    assert torch.equal(rotate_half(x), expected)
+
+
+def test_known_value():
+    """Check the rotation against an explicit per-pair 2x2 rotation, locking in the
+    ``angle = -tau * 2pi * (1/theta^(2k/rotary_dim))`` formula, the sign, and the GPT-J pairing."""
+    dim, theta, tau = 4, 100.0, 0.5
+    rote = RotaryTimeEmbedding(dim, theta=theta, rotary_fraction=1.0)
+
+    x = torch.tensor([[[1.0, 2.0, 3.0, 4.0]]])  # (B=1, T=1, C=4)
+    times = torch.tensor([[tau]])
+    out = rote(x, times)
+
+    # Two channel pairs, each rotated by angle_k = -tau * 2pi / theta^(2k/dim), k in {0, 1}.
+    angles = [-tau * 2.0 * math.pi / (theta ** (2 * k / dim)) for k in range(dim // 2)]
+    expected = torch.empty_like(x)
+    for k, a in enumerate(angles):
+        xa, xb = x[0, 0, 2 * k], x[0, 0, 2 * k + 1]
+        c, s = math.cos(a), math.sin(a)
+        # [[c, -s], [s, c]] @ [xa, xb]
+        expected[0, 0, 2 * k] = xa * c - xb * s
+        expected[0, 0, 2 * k + 1] = xa * s + xb * c
+
+    assert torch.allclose(out, expected, atol=1e-6)
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16])
+def test_shape_dtype_preserved(dtype):
+    """Output keeps the input shape and dtype (the internal angle math runs in fp32, so the
+    passthrough channels round-trip exactly for fp32 and the lower-precision dtypes)."""
+    dim = 80
+    rote = RotaryTimeEmbedding(dim)
+    x = torch.randn(2, 5, dim, dtype=dtype)
+    times = torch.arange(5, dtype=torch.float32).unsqueeze(0).expand(2, -1)
+    out = rote(x, times)
+    assert out.shape == x.shape
+    assert out.dtype == dtype
+    # Partial rotation: channels beyond rotary_dim pass through unchanged.
+    assert torch.equal(out[..., rote.rotary_dim :], x[..., rote.rotary_dim :])

tests/collections/speechlm2/test_salm.py

@@ -0,0 +1,453 @@
+# Copyright (c) 2025, 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.
+import os
+
+import pytest
+import torch
+from lhotse import CutSet, SupervisionSegment
+from lhotse.testing.dummies import dummy_cut, dummy_recording
+from omegaconf import DictConfig
+from transformers import GenerationConfig
+
+from nemo.collections.common.data.lhotse import NeMoMultimodalConversation
+from nemo.collections.common.data.lhotse.text_adapters import AudioTurn, TextTurn
+from nemo.collections.common.data.utils import move_data_to_device
+from nemo.collections.common.prompts import PromptFormatter
+from nemo.collections.speechlm2.data import SALMDataset
+from nemo.collections.speechlm2.models import SALM
+from tests.collections.speechlm2._chunking_helpers import (
+    ChunkingTestPerception,
+    ChunkingTestTokenizer,
+    chunking_test_devices,
+)
+
+if torch.cuda.is_available():
+    torch.set_default_device('cuda')
+
+
+def resolve_pretrained_models():
+    if os.path.exists("/home/TestData/speechlm/pretrained_models"):
+        # CI pre-cached paths:
+        return {
+            "pretrained_llm": "/home/TestData/speechlm/pretrained_models/TinyLlama--TinyLlama_v1.1",
+            "pretrained_asr": "/home/TestData/speechlm/pretrained_models/canary-1b-flash.nemo",
+        }
+    else:
+        # HF URLs:
+        return {
+            "pretrained_asr": "nvidia/canary-1b-flash",
+            "pretrained_llm": "TinyLlama/TinyLlama_v1.1",
+        }
+
+
+AUDIO_LOCATOR_TAG = "<|audioplaceholder|>"
+PROMPT = "llama2"
+
+
+@pytest.fixture(scope="session")
+def model():
+    cfg = {
+        **resolve_pretrained_models(),
+        "pretrained_weights": False,
+        "prompt_format": PROMPT,
+        "audio_locator_tag": AUDIO_LOCATOR_TAG,
+        "perception": {
+            "target": "nemo.collections.speechlm2.modules.perception.AudioPerceptionModule",
+            "output_dim": 2048,
+            "encoder": {
+                "_target_": "nemo.collections.asr.modules.ConformerEncoder",
+                "att_context_size": [-1, -1],
+                "causal_downsampling": False,
+                "conv_context_size": None,
+                "conv_kernel_size": 9,
+                "conv_norm_type": "batch_norm",
+                "d_model": 1024,
+                "dropout": 0.1,
+                "dropout_att": 0.1,
+                "dropout_emb": 0.0,
+                "dropout_pre_encoder": 0.1,
+                "feat_in": 128,
+                "feat_out": -1,
+                "ff_expansion_factor": 4,
+                "n_heads": 8,
+                "n_layers": 2,
+                "pos_emb_max_len": 5000,
+                "self_attention_model": "rel_pos",
+                "subsampling": "dw_striding",
+                "subsampling_conv_channels": 256,
+                "subsampling_factor": 8,
+            },
+            "modality_adapter": {
+                "_target_": "nemo.collections.speechlm2.modules.perception.IdentityConnector",
+                "d_model": 1024,
+            },
+            "preprocessor": {
+                "_target_": "nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor",
+                "dither": 1e-05,
+                "features": 128,
+                "frame_splicing": 1,
+                "log": True,
+                "n_fft": 512,
+                "normalize": "per_feature",
+                "pad_to": 0,
+                "pad_value": 0.0,
+                "sample_rate": 16000,
+                "window": "hann",
+                "window_size": 0.025,
+                "window_stride": 0.01,
+            },
+        },
+        "optimizer": {"_target_": "torch.optim.AdamW"},
+    }
+    model = SALM(cfg)
+    if torch.cuda.is_available():
+        model.to("cuda")
+    return model
+
+
+@pytest.fixture(scope="session")
+def dataset(model):
+    return SALMDataset(model.tokenizer)
+
+
+@pytest.fixture(scope="session")
+def prompt_formatter(model):
+    return PromptFormatter.resolve(PROMPT)(model.tokenizer)
+
+
+@pytest.fixture(scope="session")
+def training_cutset_batch():
+    cut = dummy_cut(0, recording=dummy_recording(0, with_data=True))
+    cut.supervisions = [
+        SupervisionSegment(
+            id=cut.id, recording_id=cut.recording_id, start=0, duration=1.0, text='Some text transcription.'
+        )
+    ]
+    return CutSet(
+        [
+            NeMoMultimodalConversation(
+                id="example-0",
+                turns=[
+                    TextTurn(role="user", value="Repeat after me:"),
+                    AudioTurn(role="user", cut=cut, audio_locator_tag=AUDIO_LOCATOR_TAG),
+                    TextTurn(role="assistant", value=cut.supervisions[0].text),
+                ],
+                token_equivalent_duration=0.08,
+            )
+        ]
+    )
+
+
+def test_salm_dataset(dataset, prompt_formatter, training_cutset_batch):
+    # This first step pre-tokenizes the examples, usually handled within `get_lhotse_dataloder_from_config`.
+    training_cutset_batch = training_cutset_batch.map(lambda c: c.apply_prompt_format(prompt_formatter), apply_fn=None)
+    # fmt: off
+    tokenized = training_cutset_batch[0].input_ids
+    assert (
+        prompt_formatter.tokenizer.tokenizer.decode(tokenized) ==
+        f"<s> [INST] Repeat after me: {AUDIO_LOCATOR_TAG} [/INST] Some text transcription. </s>"
+    )
+    # fmt: on
+    batch = dataset[training_cutset_batch]
+    for key in ("audios", "audio_lens", "input_ids", "loss_mask"):
+        assert key in batch
+        assert torch.is_tensor(batch[key])
+
+
+def test_salm_training_step(model, dataset, prompt_formatter, training_cutset_batch):
+    training_cutset_batch = training_cutset_batch.map(lambda c: c.apply_prompt_format(prompt_formatter), apply_fn=None)
+    batch = dataset[training_cutset_batch]
+    batch = move_data_to_device(batch, device=model.device)
+    results = model.training_step(batch, batch_idx=0)
+    assert torch.is_tensor(results["loss"])
+    assert not torch.isnan(results["loss"])
+    assert results["loss"] > 0
+
+
+def test_salm_validation_step(model, dataset, prompt_formatter, training_cutset_batch):
+    model.on_validation_epoch_start()
+    training_cutset_batch = training_cutset_batch.map(lambda c: c.apply_prompt_format(prompt_formatter), apply_fn=None)
+    batch = dataset[training_cutset_batch]
+    batch = move_data_to_device(batch, device=model.device)
+    results = model.validation_step({"dummy_val_set": batch}, batch_idx=0)
+    assert results is None
+
+
+def test_salm_generation(model):
+    answer = model.generate(
+        prompts=[
+            [
+                {"role": "user", "slots": {"message": f"Repeat after me: {AUDIO_LOCATOR_TAG}"}},
+            ]
+        ],
+        audios=torch.randn(1, 16000),
+        audio_lens=torch.tensor([16000]),
+        max_new_tokens=4,
+    )
+    assert answer.shape == (1, 4)
+    assert answer.dtype == torch.long
+    assert (answer >= 0).all()
+    assert (answer < model.text_vocab_size).all()
+
+
+@pytest.mark.parametrize(
+    ("enable_thinking", "expected_formatter_kwargs"),
+    [
+        (False, {"enable_thinking": False}),
+        (None, {}),
+    ],
+)
+def test_salm_generation_passes_enable_thinking(model, monkeypatch, enable_thinking, expected_formatter_kwargs):
+    seen = {}
+
+    class _FakeFormatter:
+        def __init__(self, tokenizer):
+            pass
+
+        def encode_dialog(self, turns, **kwargs):
+            seen["turns"] = turns
+            seen["formatter_kwargs"] = kwargs
+            return {"input_ids": torch.tensor([1, 2], dtype=torch.long)}
+
+    def fake_generate(*, input_ids, attention_mask, generation_config, **kwargs):
+        seen["input_ids"] = input_ids
+        seen["attention_mask"] = attention_mask
+        max_new_tokens = kwargs["max_new_tokens"]
+        return torch.zeros((input_ids.shape[0], max_new_tokens), dtype=torch.long, device=input_ids.device)
+
+    monkeypatch.setattr(PromptFormatter, "resolve", staticmethod(lambda name: _FakeFormatter))
+    monkeypatch.setattr(model.llm, "generate", fake_generate, raising=False)
+
+    answer = model.generate(
+        prompts=[[{"role": "user", "slots": {"message": "test"}}]],
+        enable_thinking=enable_thinking,
+        max_new_tokens=3,
+    )
+
+    assert seen["formatter_kwargs"] == expected_formatter_kwargs
+    assert seen["turns"] == [{"role": "user", "slots": {"message": "test"}}]
+    assert seen["input_ids"].shape == (1, 2)
+    assert torch.equal(seen["attention_mask"], torch.ones_like(seen["input_ids"], dtype=torch.bool))
+    assert answer.shape == (1, 3)
+
+
+def test_salm_generation_audios_via_prompt(model, tmp_path):
+    audio_path = tmp_path / "audio.wav"
+    dummy_cut(0, with_data=True).save_audio(audio_path)
+
+    answer = model.generate(
+        prompts=[
+            [{"role": "user", "content": f"Repeat after me: {AUDIO_LOCATOR_TAG}", "audio": [audio_path]}],
+            [
+                {
+                    "role": "user",
+                    "content": f"Repeat after me: {AUDIO_LOCATOR_TAG} and {AUDIO_LOCATOR_TAG}",
+                    "audio": [audio_path, audio_path],
+                }
+            ],
+        ],
+        generation_config=GenerationConfig(max_new_tokens=4),
+    )
+    assert answer.shape == (2, 4)
+    assert answer.dtype == torch.long
+    assert (answer >= 0).all()
+    assert (answer < model.text_vocab_size).all()
+
+
+def test_salm_generation_prompts_as_tensor(model):
+    answer = model.generate(
+        prompts=torch.tensor([[1, 2, 3, 4, 5, 6, 7, model.audio_locator_tag_id]]),
+        audios=torch.randn(1, 16000),
+        audio_lens=torch.tensor([16000]),
+        max_new_tokens=4,
+    )
+    assert answer.shape == (1, 4)
+    assert answer.dtype == torch.long
+    assert (answer >= 0).all()
+    assert (answer < model.text_vocab_size).all()
+
+
+@pytest.mark.parametrize("device", chunking_test_devices())
+def test_salm_prepare_inputs_chunks_long_audio(device):
+    model = _make_chunking_test_model(encoder_chunk_size_seconds=1.0, sampling_rate=2, device=device)
+    batch = {
+        "audios": torch.tensor([[1.0, 2.0, 3.0, 4.0, 5.0]], device=device),
+        "audio_lens": torch.tensor([5], dtype=torch.long, device=device),
+        "input_ids": torch.tensor([[model.audio_locator_tag_id, 10]], dtype=torch.long, device=device),
+        "loss_mask": torch.tensor([[False, True]], dtype=torch.bool, device=device),
+    }
+
+    inputs = model.prepare_inputs(batch)
+
+    chunked_signal, chunked_lens = model.perception.calls[0]
+    assert chunked_signal.shape == (2, 3)
+    assert torch.equal(chunked_lens, torch.tensor([2, 3], dtype=torch.long, device=device))
+    assert torch.equal(inputs["input_embeds"][0, :, 0], torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0], device=device))
+    assert torch.equal(inputs["attention_mask"], torch.ones((1, 5), dtype=torch.bool, device=device))
+
+
+@pytest.mark.parametrize("device", chunking_test_devices())
+def test_salm_prepare_inputs_passes_chunk_time_offsets(device):
+    # sampling_rate=8, chunk=0.5s -> 4 samples/chunk; a 12-sample audio splits cleanly into
+    # 3 chunks starting at samples 0, 4, 8, i.e. time offsets 0.0, 0.5, 1.0 seconds.
+    model = _make_chunking_test_model(encoder_chunk_size_seconds=0.5, sampling_rate=8, hop_length=2, device=device)
+    batch = {
+        "audios": torch.arange(1.0, 13.0, device=device).unsqueeze(0),
+        "audio_lens": torch.tensor([12], dtype=torch.long, device=device),
+        "input_ids": torch.tensor([[model.audio_locator_tag_id, 10]], dtype=torch.long, device=device),
+        "loss_mask": torch.tensor([[False, True]], dtype=torch.bool, device=device),
+    }
+
+    model.prepare_inputs(batch)
+
+    chunked_signal, chunked_lens = model.perception.calls[0]
+    assert chunked_signal.shape == (3, 4)
+    assert torch.equal(chunked_lens, torch.tensor([4, 4, 4], dtype=torch.long, device=device))
+    time_offset = model.perception.time_offsets[0]
+    assert time_offset is not None
+    # chunk N start_sample / sampling_rate: [0/8, 4/8, 8/8]
+    assert torch.equal(time_offset, torch.tensor([0.0, 0.5, 1.0], device=device))
+
+
+@pytest.mark.parametrize("device", chunking_test_devices())
+def test_salm_prepare_inputs_merges_short_tail_chunk(device):
+    model = _make_chunking_test_model(encoder_chunk_size_seconds=0.5, sampling_rate=8, hop_length=2, device=device)
+    batch = {
+        "audios": torch.tensor([[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]], device=device),
+        "audio_lens": torch.tensor([9], dtype=torch.long, device=device),
+        "input_ids": torch.tensor([[model.audio_locator_tag_id, 10]], dtype=torch.long, device=device),
+        "loss_mask": torch.tensor([[False, True]], dtype=torch.bool, device=device),
+    }
+
+    inputs = model.prepare_inputs(batch)
+
+    chunked_signal, chunked_lens = model.perception.calls[0]
+    assert chunked_signal.shape == (2, 5)
+    assert torch.equal(chunked_lens, torch.tensor([4, 5], dtype=torch.long, device=device))
+    assert torch.equal(
+        inputs["input_embeds"][0, :, 0],
+        torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0], device=device),
+    )
+
+
+@pytest.mark.parametrize("device", chunking_test_devices())
+def test_salm_prepare_inputs_skips_chunking_when_size_is_null(device):
+    model = _make_chunking_test_model(encoder_chunk_size_seconds=None, sampling_rate=2, device=device)
+    batch = {
+        "audios": torch.tensor([[1.0, 2.0, 3.0, 4.0, 5.0]], device=device),
+        "audio_lens": torch.tensor([5], dtype=torch.long, device=device),
+        "input_ids": torch.tensor([[model.audio_locator_tag_id, 10]], dtype=torch.long, device=device),
+        "loss_mask": torch.tensor([[False, True]], dtype=torch.bool, device=device),
+    }
+
+    model.prepare_inputs(batch)
+
+    input_signal, input_signal_lens = model.perception.calls[0]
+    assert input_signal.shape == (1, 5)
+    assert torch.equal(input_signal_lens, torch.tensor([5], dtype=torch.long, device=device))
+
+
+@pytest.mark.parametrize("device", chunking_test_devices())
+def test_salm_prepare_inputs_skips_chunking_when_key_absent(device):
+    """Backwards compat: missing encoder_chunk_size_seconds key matches legacy non-chunking path."""
+    model = _make_chunking_test_model(encoder_chunk_size_seconds=None, sampling_rate=2, device=device)
+    model.cfg = DictConfig({})
+    batch = {
+        "audios": torch.tensor([[1.0, 2.0, 3.0, 4.0, 5.0]], device=device),
+        "audio_lens": torch.tensor([5], dtype=torch.long, device=device),
+        "input_ids": torch.tensor([[model.audio_locator_tag_id, 10]], dtype=torch.long, device=device),
+        "loss_mask": torch.tensor([[False, True]], dtype=torch.bool, device=device),
+    }
+
+    model.prepare_inputs(batch)
+
+    input_signal, input_signal_lens = model.perception.calls[0]
+    assert input_signal.shape == (1, 5)
+    assert torch.equal(input_signal_lens, torch.tensor([5], dtype=torch.long, device=device))
+
+
+@pytest.mark.parametrize("device", chunking_test_devices())
+def test_salm_prepare_inputs_preserves_chunked_audio_order(device):
+    model = _make_chunking_test_model(encoder_chunk_size_seconds=1.0, sampling_rate=2, device=device)
+    batch = {
+        "audios": torch.tensor(
+            [
+                [1.0, 2.0, 3.0, 0.0, 0.0],
+                [10.0, 11.0, 12.0, 13.0, 14.0],
+            ],
+            device=device,
+        ),
+        "audio_lens": torch.tensor([3, 5], dtype=torch.long, device=device),
+        "input_ids": torch.tensor(
+            [[model.audio_locator_tag_id, model.audio_locator_tag_id, 10]], dtype=torch.long, device=device
+        ),
+        "loss_mask": torch.tensor([[False, False, True]], dtype=torch.bool, device=device),
+    }
+
+    inputs = model.prepare_inputs(batch)
+
+    assert torch.equal(
+        inputs["input_embeds"][0, :, 0],
+        torch.tensor([1.0, 2.0, 3.0, 10.0, 11.0, 12.0, 13.0, 14.0], device=device),
+    )
+
+
+@pytest.mark.parametrize("device", chunking_test_devices())
+def test_salm_generate_chunks_audio_before_llm(device):
+    model = _make_chunking_test_model(encoder_chunk_size_seconds=1.0, sampling_rate=2, device=device)
+
+    answer = model.generate(
+        prompts=torch.tensor([[model.audio_locator_tag_id, 10]], dtype=torch.long, device=device),
+        audios=torch.tensor([[1.0, 2.0, 3.0, 4.0, 5.0]], device=device),
+        audio_lens=torch.tensor([5], dtype=torch.long, device=device),
+        max_new_tokens=3,
+    )
+
+    chunked_signal, chunked_lens = model.perception.calls[0]
+    assert chunked_signal.shape == (2, 3)
+    assert torch.equal(chunked_lens, torch.tensor([2, 3], dtype=torch.long, device=device))
+    assert torch.equal(
+        model.llm.generate_kwargs["inputs_embeds"][0, :5, 0],
+        torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0], device=device),
+    )
+    assert answer.shape == (1, 3)
+
+
+def _make_chunking_test_model(encoder_chunk_size_seconds, sampling_rate, device, hop_length=1):
+    model = SALM.__new__(SALM)
+    torch.nn.Module.__init__(model)
+    model.cfg = DictConfig({"encoder_chunk_size_seconds": encoder_chunk_size_seconds})
+    model.audio_locator_tag = AUDIO_LOCATOR_TAG
+    model.tokenizer = ChunkingTestTokenizer(AUDIO_LOCATOR_TAG)
+    model.embed_tokens = torch.nn.Embedding(128, 1, device=device)
+    with torch.no_grad():
+        model.embed_tokens.weight.zero_()
+    model.llm = _SalmChunkingTestLLM()
+    model.perception = ChunkingTestPerception(sampling_rate=sampling_rate, hop_length=hop_length)
+    model._use_tp = False
+    return model
+
+
+class _SalmChunkingTestLLM(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        # move_embedding() writes model.embed_tokens into llm.model.embed_tokens then deletes it.
+        self.model = torch.nn.Module()
+        self.generate_kwargs = None
+
+    def generate(self, **kwargs):
+        self.generate_kwargs = kwargs
+        batch_size = kwargs["inputs_embeds"].shape[0]
+        max_new_tokens = kwargs["max_new_tokens"]
+        return torch.zeros((batch_size, max_new_tokens), dtype=torch.long, device=kwargs["inputs_embeds"].device)

tests/collections/speechlm2/test_salm_asr_decoder.py

@@ -0,0 +1,274 @@
+# Copyright (c) 2025, 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.
+import os
+
+import pytest
+import torch
+from lhotse import CutSet, SupervisionSegment
+from lhotse.testing.dummies import dummy_cut, dummy_recording
+from transformers import GenerationConfig
+
+from nemo.collections.common.data.lhotse import NeMoMultimodalConversation
+from nemo.collections.common.data.lhotse.text_adapters import AudioTurn, TextTurn
+from nemo.collections.common.data.utils import move_data_to_device
+from nemo.collections.common.prompts import PromptFormatter
+from nemo.collections.speechlm2.data import SALMDataset
+from nemo.collections.speechlm2.models.salm_asr_decoder import SALMWithAsrDecoder
+
+if torch.cuda.is_available():
+    torch.set_default_device('cuda')
+
+
+def resolve_pretrained_models():
+    if os.path.exists("/home/TestData/speechlm/pretrained_models"):
+        # CI pre-cached paths:
+        return {
+            "pretrained_llm": "/home/TestData/speechlm/pretrained_models/TinyLlama--TinyLlama_v1.1",
+            "pretrained_asr": "/home/TestData/speechlm/pretrained_models/parakeet-tdt-0.6b-v2.nemo",
+        }
+    else:
+        # HF URLs:
+        return {
+            "pretrained_asr": "nvidia/parakeet-tdt-0.6b-v2",
+            "pretrained_llm": "TinyLlama/TinyLlama_v1.1",
+        }
+
+
+AUDIO_LOCATOR_TAG = "<|audioplaceholder|>"
+PROMPT = "llama2"
+
+
+@pytest.fixture(scope="session")
+def model():
+    cfg = {
+        **resolve_pretrained_models(),
+        "pretrained_weights": True,
+        "prompt_format": PROMPT,
+        "audio_locator_tag": AUDIO_LOCATOR_TAG,
+        "perception": {
+            "target": "nemo.collections.speechlm2.modules.perception.AudioTranscriptionPerceptionModule",
+            "output_dim": 2048,
+            "asr": {
+                "encoder": {
+                    "_target_": "nemo.collections.asr.modules.ConformerEncoder",
+                    "att_context_size": [-1, -1],
+                    "causal_downsampling": False,
+                    "conv_context_size": None,
+                    "conv_kernel_size": 9,
+                    "conv_norm_type": "batch_norm",
+                    "d_model": 1024,
+                    "dropout": 0.1,
+                    "dropout_att": 0.1,
+                    "dropout_emb": 0.0,
+                    "dropout_pre_encoder": 0.1,
+                    "feat_in": 128,
+                    "feat_out": -1,
+                    "ff_expansion_factor": 4,
+                    "n_heads": 8,
+                    "n_layers": 2,
+                    "pos_emb_max_len": 5000,
+                    "self_attention_model": "rel_pos",
+                    "subsampling": "dw_striding",
+                    "subsampling_conv_channels": 256,
+                    "subsampling_factor": 8,
+                },
+                "preprocessor": {
+                    "_target_": "nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor",
+                    "dither": 1e-05,
+                    "features": 128,
+                    "frame_splicing": 1,
+                    "log": True,
+                    "n_fft": 512,
+                    "normalize": "per_feature",
+                    "pad_to": 0,
+                    "pad_value": 0.0,
+                    "sample_rate": 16000,
+                    "window": "hann",
+                    "window_size": 0.025,
+                    "window_stride": 0.01,
+                },
+            },
+            "modality_adapter": {
+                "_target_": "nemo.collections.speechlm2.modules.perception.IdentityConnector",
+                "d_model": 1024,
+            },
+        },
+        "optimizer": {"_target_": "torch.optim.AdamW"},
+    }
+    model = SALMWithAsrDecoder(cfg)
+    if torch.cuda.is_available():
+        model.to("cuda")
+    return model
+
+
+@pytest.fixture(scope="session")
+def dataset(model):
+    return SALMDataset(model.tokenizer)
+
+
+@pytest.fixture(scope="session")
+def prompt_formatter(model):
+    return PromptFormatter.resolve(PROMPT)(model.tokenizer)
+
+
+@pytest.fixture(scope="session")
+def training_cutset_batch():
+    cut = dummy_cut(0, recording=dummy_recording(0, with_data=True))
+    cut.supervisions = [
+        SupervisionSegment(
+            id=cut.id, recording_id=cut.recording_id, start=0, duration=1.0, text='Some text transcription.'
+        )
+    ]
+    return CutSet(
+        [
+            NeMoMultimodalConversation(
+                id="example-0",
+                turns=[
+                    TextTurn(role="user", value="Repeat after me:"),
+                    AudioTurn(role="user", cut=cut, audio_locator_tag=AUDIO_LOCATOR_TAG),
+                    TextTurn(role="assistant", value=cut.supervisions[0].text),
+                ],
+                token_equivalent_duration=0.08,
+            )
+        ]
+    )
+
+
+def test_salm_dataset(dataset, prompt_formatter, training_cutset_batch):
+    # This first step pre-tokenizes the examples, usually handled within `get_lhotse_dataloder_from_config`.
+    training_cutset_batch = training_cutset_batch.map(lambda c: c.apply_prompt_format(prompt_formatter), apply_fn=None)
+    # fmt: off
+    tokenized = training_cutset_batch[0].input_ids
+    assert (
+        prompt_formatter.tokenizer.tokenizer.decode(tokenized) ==
+        f"<s> [INST] Repeat after me: {AUDIO_LOCATOR_TAG} [/INST] Some text transcription. </s>"
+    )
+    # fmt: on
+    batch = dataset[training_cutset_batch]
+    for key in ("audios", "audio_lens", "input_ids", "loss_mask"):
+        assert key in batch
+        assert torch.is_tensor(batch[key])
+
+
+def test_salm_training_step(model, dataset, prompt_formatter, training_cutset_batch):
+    training_cutset_batch = training_cutset_batch.map(lambda c: c.apply_prompt_format(prompt_formatter), apply_fn=None)
+    batch = dataset[training_cutset_batch]
+    batch = move_data_to_device(batch, device=model.device)
+    results = model.training_step(batch, batch_idx=0)
+    assert torch.is_tensor(results["loss"])
+    assert not torch.isnan(results["loss"])
+    assert results["loss"] > 0
+
+
+def test_salm_validation_step(model, dataset, prompt_formatter, training_cutset_batch):
+    model.on_validation_epoch_start()
+    training_cutset_batch = training_cutset_batch.map(lambda c: c.apply_prompt_format(prompt_formatter), apply_fn=None)
+    batch = dataset[training_cutset_batch]
+    batch = move_data_to_device(batch, device=model.device)
+    results = model.validation_step({"dummy_val_set": batch}, batch_idx=0)
+    assert results is None
+
+
+def test_salm_generation(model):
+    answer = model.generate(
+        prompts=[
+            [
+                {"role": "user", "slots": {"message": f"Repeat after me: {AUDIO_LOCATOR_TAG}"}},
+            ]
+        ],
+        audios=torch.randn(1, 16000),
+        audio_lens=torch.tensor([16000]),
+        max_new_tokens=4,
+    )
+    assert answer.shape == (1, 4)
+    assert answer.dtype == torch.long
+    assert (answer >= 0).all()
+    assert (answer < model.text_vocab_size).all()
+
+
+@pytest.mark.parametrize(
+    ("enable_thinking", "expected_formatter_kwargs"),
+    [
+        (False, {"enable_thinking": False}),
+        (None, {}),
+    ],
+)
+def test_salm_generation_passes_enable_thinking(model, monkeypatch, enable_thinking, expected_formatter_kwargs):
+    seen = {}
+
+    class _FakeFormatter:
+        def __init__(self, tokenizer):
+            pass
+
+        def encode_dialog(self, turns, **kwargs):
+            seen["turns"] = turns
+            seen["formatter_kwargs"] = kwargs
+            return {"input_ids": torch.tensor([1, 2], dtype=torch.long)}
+
+    def fake_generate(*, inputs_embeds, attention_mask, generation_config, **kwargs):
+        seen["inputs_embeds"] = inputs_embeds
+        seen["attention_mask"] = attention_mask
+        max_new_tokens = kwargs["max_new_tokens"]
+        return torch.zeros((inputs_embeds.shape[0], max_new_tokens), dtype=torch.long, device=inputs_embeds.device)
+
+    monkeypatch.setattr(PromptFormatter, "resolve", staticmethod(lambda name: _FakeFormatter))
+    monkeypatch.setattr(model.llm, "generate", fake_generate, raising=False)
+
+    answer = model.generate(
+        prompts=[[{"role": "user", "slots": {"message": "test"}}]],
+        enable_thinking=enable_thinking,
+        max_new_tokens=3,
+    )
+
+    assert seen["formatter_kwargs"] == expected_formatter_kwargs
+    assert seen["turns"] == [{"role": "user", "slots": {"message": "test"}}]
+    assert seen["inputs_embeds"].shape[:2] == (1, 2)
+    assert torch.equal(seen["attention_mask"], torch.ones((1, 2), dtype=torch.bool, device=model.device))
+    assert answer.shape == (1, 3)
+
+
+def test_salm_generation_audios_via_prompt(model, tmp_path):
+    audio_path = tmp_path / "audio.wav"
+    dummy_cut(0, with_data=True).save_audio(audio_path)
+
+    answer = model.generate(
+        prompts=[
+            [{"role": "user", "content": f"Repeat after me: {AUDIO_LOCATOR_TAG}", "audio": [audio_path]}],
+            [
+                {
+                    "role": "user",
+                    "content": f"Repeat after me: {AUDIO_LOCATOR_TAG} and {AUDIO_LOCATOR_TAG}",
+                    "audio": [audio_path, audio_path],
+                }
+            ],
+        ],
+        generation_config=GenerationConfig(max_new_tokens=4),
+    )
+    assert answer.shape == (2, 4)
+    assert answer.dtype == torch.long
+    assert (answer >= 0).all()
+    assert (answer < model.text_vocab_size).all()
+
+
+def test_salm_generation_prompts_as_tensor(model):
+    answer = model.generate(
+        prompts=torch.tensor([[1, 2, 3, 4, 5, 6, 7, model.audio_locator_tag_id]]),
+        audios=torch.randn(1, 16000),
+        audio_lens=torch.tensor([16000]),
+        max_new_tokens=4,
+    )
+    assert answer.shape == (1, 4)
+    assert answer.dtype == torch.long
+    assert (answer >= 0).all()
+    assert (answer < model.text_vocab_size).all()

tests/collections/speechlm2/test_salm_asr_decoder_multilayerproj.py

@@ -0,0 +1,235 @@
+# Copyright (c) 2025, 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.
+import os
+
+import pytest
+import torch
+from lhotse import CutSet, SupervisionSegment
+from lhotse.testing.dummies import dummy_cut, dummy_recording
+from transformers import GenerationConfig
+
+from nemo.collections.common.data.lhotse import NeMoMultimodalConversation
+from nemo.collections.common.data.lhotse.text_adapters import AudioTurn, TextTurn
+from nemo.collections.common.data.utils import move_data_to_device
+from nemo.collections.common.prompts import PromptFormatter
+from nemo.collections.speechlm2.data import SALMDataset
+from nemo.collections.speechlm2.models.salm_asr_decoder import SALMWithAsrDecoder
+
+if torch.cuda.is_available():
+    torch.set_default_device('cuda')
+
+
+def resolve_pretrained_models():
+    if os.path.exists("/home/TestData/speechlm/pretrained_models"):
+        # CI pre-cached paths:
+        return {
+            "pretrained_llm": "/home/TestData/speechlm/pretrained_models/TinyLlama--TinyLlama_v1.1",
+            "pretrained_asr": "/home/TestData/speechlm/pretrained_models/parakeet-tdt-0.6b-v2.nemo",
+        }
+    else:
+        # HF URLs:
+        return {
+            "pretrained_asr": "nvidia/parakeet-tdt-0.6b-v2",
+            "pretrained_llm": "TinyLlama/TinyLlama_v1.1",
+        }
+
+
+AUDIO_LOCATOR_TAG = "<|audioplaceholder|>"
+PROMPT = "llama2"
+
+
+@pytest.fixture(scope="session")
+def model():
+    cfg = {
+        **resolve_pretrained_models(),
+        "pretrained_weights": True,
+        "prompt_format": PROMPT,
+        "audio_locator_tag": AUDIO_LOCATOR_TAG,
+        "perception": {
+            "target": "nemo.collections.speechlm2.modules.perception.AudioTranscriptionPerceptionModule",
+            "output_dim": 2048,
+            "asr": {
+                "encoder": {
+                    "_target_": "nemo.collections.asr.modules.ConformerEncoder",
+                    "att_context_size": [-1, -1],
+                    "causal_downsampling": False,
+                    "conv_context_size": None,
+                    "conv_kernel_size": 9,
+                    "conv_norm_type": "batch_norm",
+                    "d_model": 1024,
+                    "dropout": 0.1,
+                    "dropout_att": 0.1,
+                    "dropout_emb": 0.0,
+                    "dropout_pre_encoder": 0.1,
+                    "feat_in": 128,
+                    "feat_out": -1,
+                    "ff_expansion_factor": 4,
+                    "n_heads": 8,
+                    "n_layers": 2,
+                    "pos_emb_max_len": 5000,
+                    "self_attention_model": "rel_pos",
+                    "subsampling": "dw_striding",
+                    "subsampling_conv_channels": 256,
+                    "subsampling_factor": 8,
+                },
+                "preprocessor": {
+                    "_target_": "nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor",
+                    "dither": 1e-05,
+                    "features": 128,
+                    "frame_splicing": 1,
+                    "log": True,
+                    "n_fft": 512,
+                    "normalize": "per_feature",
+                    "pad_to": 0,
+                    "pad_value": 0.0,
+                    "sample_rate": 16000,
+                    "window": "hann",
+                    "window_size": 0.025,
+                    "window_stride": 0.01,
+                },
+            },
+            "modality_adapter": {
+                "_target_": "nemo.collections.speechlm2.modules.perception.MultiLayerProjectionConnector",
+                "target_layer_ids": [7, 15, 19, 23],
+                "input_dim": 1024,
+                "output_dim": 2048,
+            },
+        },
+        "optimizer": {"_target_": "torch.optim.AdamW"},
+    }
+    model = SALMWithAsrDecoder(cfg)
+    if torch.cuda.is_available():
+        model.to("cuda")
+    return model
+
+
+@pytest.fixture(scope="session")
+def dataset(model):
+    return SALMDataset(model.tokenizer)
+
+
+@pytest.fixture(scope="session")
+def prompt_formatter(model):
+    return PromptFormatter.resolve(PROMPT)(model.tokenizer)
+
+
+@pytest.fixture(scope="session")
+def training_cutset_batch():
+    cut = dummy_cut(0, recording=dummy_recording(0, with_data=True))
+    cut.supervisions = [
+        SupervisionSegment(
+            id=cut.id, recording_id=cut.recording_id, start=0, duration=1.0, text='Some text transcription.'
+        )
+    ]
+    return CutSet(
+        [
+            NeMoMultimodalConversation(
+                id="example-0",
+                turns=[
+                    TextTurn(role="user", value="Repeat after me:"),
+                    AudioTurn(role="user", cut=cut, audio_locator_tag=AUDIO_LOCATOR_TAG),
+                    TextTurn(role="assistant", value=cut.supervisions[0].text),
+                ],
+                token_equivalent_duration=0.08,
+            )
+        ]
+    )
+
+
+def test_salm_dataset(dataset, prompt_formatter, training_cutset_batch):
+    # This first step pre-tokenizes the examples, usually handled within `get_lhotse_dataloder_from_config`.
+    training_cutset_batch = training_cutset_batch.map(lambda c: c.apply_prompt_format(prompt_formatter), apply_fn=None)
+    # fmt: off
+    tokenized = training_cutset_batch[0].input_ids
+    assert (
+        prompt_formatter.tokenizer.tokenizer.decode(tokenized) ==
+        f"<s> [INST] Repeat after me: {AUDIO_LOCATOR_TAG} [/INST] Some text transcription. </s>"
+    )
+    # fmt: on
+    batch = dataset[training_cutset_batch]
+    for key in ("audios", "audio_lens", "input_ids", "loss_mask"):
+        assert key in batch
+        assert torch.is_tensor(batch[key])
+
+
+def test_salm_training_step(model, dataset, prompt_formatter, training_cutset_batch):
+    training_cutset_batch = training_cutset_batch.map(lambda c: c.apply_prompt_format(prompt_formatter), apply_fn=None)
+    batch = dataset[training_cutset_batch]
+    batch = move_data_to_device(batch, device=model.device)
+    results = model.training_step(batch, batch_idx=0)
+    assert torch.is_tensor(results["loss"])
+    assert not torch.isnan(results["loss"])
+    assert results["loss"] > 0
+
+
+def test_salm_validation_step(model, dataset, prompt_formatter, training_cutset_batch):
+    model.on_validation_epoch_start()
+    training_cutset_batch = training_cutset_batch.map(lambda c: c.apply_prompt_format(prompt_formatter), apply_fn=None)
+    batch = dataset[training_cutset_batch]
+    batch = move_data_to_device(batch, device=model.device)
+    results = model.validation_step({"dummy_val_set": batch}, batch_idx=0)
+    assert results is None
+
+
+def test_salm_generation(model):
+    answer = model.generate(
+        prompts=[
+            [
+                {"role": "user", "slots": {"message": f"Repeat after me: {AUDIO_LOCATOR_TAG}"}},
+            ]
+        ],
+        audios=torch.randn(1, 16000),
+        audio_lens=torch.tensor([16000]),
+        max_new_tokens=4,
+    )
+    assert answer.shape == (1, 4)
+    assert answer.dtype == torch.long
+    assert (answer >= 0).all()
+    assert (answer < model.text_vocab_size).all()
+
+
+def test_salm_generation_audios_via_prompt(model, tmp_path):
+    audio_path = tmp_path / "audio.wav"
+    dummy_cut(0, with_data=True).save_audio(audio_path)
+
+    answer = model.generate(
+        prompts=[
+            [{"role": "user", "content": f"Repeat after me: {AUDIO_LOCATOR_TAG}", "audio": [audio_path]}],
+            [
+                {
+                    "role": "user",
+                    "content": f"Repeat after me: {AUDIO_LOCATOR_TAG} and {AUDIO_LOCATOR_TAG}",
+                    "audio": [audio_path, audio_path],
+                }
+            ],
+        ],
+        generation_config=GenerationConfig(max_new_tokens=4),
+    )
+    assert answer.shape == (2, 4)
+    assert answer.dtype == torch.long
+    assert (answer >= 0).all()
+    assert (answer < model.text_vocab_size).all()
+
+
+def test_salm_generation_prompts_as_tensor(model):
+    answer = model.generate(
+        prompts=torch.tensor([[1, 2, 3, 4, 5, 6, 7, model.audio_locator_tag_id]]),
+        audios=torch.randn(1, 16000),
+        audio_lens=torch.tensor([16000]),
+        max_new_tokens=4,
+    )
+    assert answer.shape == (1, 4)
+    assert answer.dtype == torch.long
+    assert (answer >= 0).all()
+    assert (answer < model.text_vocab_size).all()

tests/collections/speechlm2/test_salm_automodel.py

@@ -0,0 +1,480 @@
+# Copyright (c) 2025, 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.
+import os
+
+import pytest
+import torch
+from lhotse import CutSet, SupervisionSegment
+from lhotse.testing.dummies import dummy_cut, dummy_recording
+from transformers import GenerationConfig
+
+from nemo.collections.common.data.lhotse import NeMoMultimodalConversation
+from nemo.collections.common.data.lhotse.text_adapters import AudioTurn, TextTurn
+from nemo.collections.common.data.utils import move_data_to_device
+from nemo.collections.common.prompts import PromptFormatter
+from nemo.collections.speechlm2.data import SALMDataset
+from nemo.collections.speechlm2.models import SALMAutomodel
+from tests.collections.speechlm2._chunking_helpers import (
+    ChunkingTestPerception,
+    ChunkingTestTokenizer,
+    chunking_test_devices,
+)
+
+requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="SALMAutomodel requires CUDA")
+
+
+@pytest.fixture(autouse=True, scope="module")
+def _default_device_cuda():
+    """Run this module's tests on CUDA by default, but scope the change so it does
+    not leak. ``torch.set_default_device`` is a global, process-wide mutation; setting
+    it at import time bleeds into other modules collected in the same pytest session
+    (e.g. the CPU tests in tests/collections/asr/test_parallel_expert_encoder.py),
+    causing spurious cuda/cpu device-mismatch failures. The previous default device
+    is always restored on teardown.
+    """
+    if not torch.cuda.is_available():
+        yield
+        return
+    prev = torch.get_default_device()
+    torch.set_default_device('cuda')
+    try:
+        yield
+    finally:
+        torch.set_default_device(prev)
+
+
+def resolve_pretrained_models():
+    if os.path.exists("/home/TestData/speechlm/pretrained_models"):
+        # CI pre-cached paths:
+        return {
+            "pretrained_llm": "/home/TestData/speechlm/pretrained_models/Qwen--Qwen3-1.7B",
+            "pretrained_asr": "/home/TestData/speechlm/pretrained_models/canary-1b-flash.nemo",
+        }
+    else:
+        # HF URLs:
+        return {
+            "pretrained_asr": "nvidia/canary-1b-flash",
+            "pretrained_llm": "Qwen/Qwen3-1.7B",
+        }
+
+
+AUDIO_LOCATOR_TAG = "<|audioplaceholder|>"
+PROMPT = "qwen"
+
+
+@pytest.fixture(scope="session")
+def model():
+    if not torch.cuda.is_available():
+        pytest.skip("SALMAutomodel requires CUDA")
+    cfg = {
+        **resolve_pretrained_models(),
+        "pretrained_weights": False,
+        "prompt_format": PROMPT,
+        "audio_locator_tag": AUDIO_LOCATOR_TAG,
+        "perception": {
+            "target": "nemo.collections.speechlm2.modules.perception.AudioPerceptionModule",
+            "output_dim": 2048,
+            "encoder": {
+                "_target_": "nemo.collections.asr.modules.ConformerEncoder",
+                "att_context_size": [-1, -1],
+                "causal_downsampling": False,
+                "conv_context_size": None,
+                "conv_kernel_size": 9,
+                "conv_norm_type": "batch_norm",
+                "d_model": 1024,
+                "dropout": 0.1,
+                "dropout_att": 0.1,
+                "dropout_emb": 0.0,
+                "dropout_pre_encoder": 0.1,
+                "feat_in": 128,
+                "feat_out": -1,
+                "ff_expansion_factor": 4,
+                "n_heads": 8,
+                "n_layers": 2,
+                "pos_emb_max_len": 5000,
+                "self_attention_model": "rel_pos",
+                "subsampling": "dw_striding",
+                "subsampling_conv_channels": 256,
+                "subsampling_factor": 8,
+            },
+            "modality_adapter": {
+                "_target_": "nemo.collections.speechlm2.modules.perception.IdentityConnector",
+                "d_model": 1024,
+            },
+            "preprocessor": {
+                "_target_": "nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor",
+                "dither": 1e-05,
+                "features": 128,
+                "frame_splicing": 1,
+                "log": True,
+                "n_fft": 512,
+                "normalize": "per_feature",
+                "pad_to": 0,
+                "pad_value": 0.0,
+                "sample_rate": 16000,
+                "window": "hann",
+                "window_size": 0.025,
+                "window_stride": 0.01,
+            },
+        },
+        "optimizer": {"_target_": "torch.optim.AdamW"},
+        "torch_dtype": "bfloat16",
+    }
+    model = SALMAutomodel(cfg)
+    model.configure_model()
+    model.to("cuda")
+    return model
+
+
+@pytest.fixture(scope="session")
+def dataset(model):
+    return SALMDataset(model.tokenizer)
+
+
+@pytest.fixture(scope="session")
+def prompt_formatter(model):
+    return PromptFormatter.resolve(PROMPT)(model.tokenizer)
+
+
+@pytest.fixture(scope="session")
+def training_cutset_batch():
+    cut = dummy_cut(0, recording=dummy_recording(0, with_data=True))
+    cut.supervisions = [
+        SupervisionSegment(
+            id=cut.id, recording_id=cut.recording_id, start=0, duration=1.0, text='Some text transcription.'
+        )
+    ]
+    return CutSet(
+        [
+            NeMoMultimodalConversation(
+                id="example-0",
+                turns=[
+                    TextTurn(role="user", value="Repeat after me:"),
+                    AudioTurn(role="user", cut=cut, audio_locator_tag=AUDIO_LOCATOR_TAG),
+                    TextTurn(role="assistant", value=cut.supervisions[0].text),
+                ],
+                token_equivalent_duration=0.08,
+            )
+        ]
+    )
+
+
+@requires_cuda
+def test_salm_automodel_dataset(dataset, prompt_formatter, training_cutset_batch):
+    # This first step pre-tokenizes the examples, usually handled within `get_lhotse_dataloder_from_config`.
+    training_cutset_batch = training_cutset_batch.map(lambda c: c.apply_prompt_format(prompt_formatter), apply_fn=None)
+    # fmt: off
+    tokenized = training_cutset_batch[0].input_ids
+    assert (
+        prompt_formatter.tokenizer.tokenizer.decode(tokenized) ==
+        f"<|im_start|>user\nRepeat after me: {AUDIO_LOCATOR_TAG}<|im_end|>\n<|im_start|>assistant\nSome text transcription.<|im_end|>\n"
+    )
+    # fmt: on
+    batch = dataset[training_cutset_batch]
+    for key in ("audios", "audio_lens", "input_ids", "loss_mask"):
+        assert key in batch
+        assert torch.is_tensor(batch[key])
+
+
+@requires_cuda
+def test_salm_automodel_training_step(model, dataset, prompt_formatter, training_cutset_batch):
+    training_cutset_batch = training_cutset_batch.map(lambda c: c.apply_prompt_format(prompt_formatter), apply_fn=None)
+    batch = dataset[training_cutset_batch]
+    batch = move_data_to_device(batch, device=model.device)
+    results = model.training_step(batch, batch_idx=0)
+    assert torch.is_tensor(results["loss"])
+    assert not torch.isnan(results["loss"])
+    assert results["loss"] > 0
+
+
+@requires_cuda
+def test_salm_automodel_validation_step(model, dataset, prompt_formatter, training_cutset_batch):
+    model.on_validation_epoch_start()
+    training_cutset_batch = training_cutset_batch.map(lambda c: c.apply_prompt_format(prompt_formatter), apply_fn=None)
+    batch = dataset[training_cutset_batch]
+    batch = move_data_to_device(batch, device=model.device)
+    results = model.validation_step({"dummy_val_set": batch}, batch_idx=0)
+    assert results is None
+
+
+def test_salm_automodel_validation_epoch_end_uses_token_weighted_metrics():
+    model = SALMAutomodel.__new__(SALMAutomodel)
+    torch.nn.Module.__init__(model)
+    model.on_validation_epoch_start()
+    model._get_moe_dp_group = lambda: None
+
+    model._partial_val_loss_sums["dummy"].extend([torch.tensor(2.0), torch.tensor(18.0)])
+    model._partial_val_corrects["dummy"].extend([torch.tensor(1.0), torch.tensor(3.0)])
+    model._partial_val_num_frames["dummy"].extend([torch.tensor(1.0), torch.tensor(9.0)])
+
+    logged = {}
+
+    def fake_log(name, value, **kwargs):
+        logged[name] = value.detach().cpu()
+
+    model.log = fake_log
+    model.on_validation_epoch_end()
+
+    assert logged["val_loss_dummy"].item() == pytest.approx(2.0)
+    assert logged["val_acc_dummy"].item() == pytest.approx(0.4)
+    assert logged["val_loss"].item() == pytest.approx(2.0)
+    assert logged["val_acc"].item() == pytest.approx(0.4)
+    assert not model._partial_val_loss_sums
+    assert not model._partial_val_corrects
+    assert not model._partial_val_num_frames
+
+
+@requires_cuda
+def test_salm_automodel_generation(model):
+    answer = model.generate(
+        prompts=[
+            [
+                {"role": "user", "slots": {"message": f"Repeat after me: {AUDIO_LOCATOR_TAG}"}},
+            ]
+        ],
+        audios=torch.randn(1, 16000),
+        audio_lens=torch.tensor([16000]),
+        max_new_tokens=4,
+    )
+    assert answer.shape == (1, 4)
+    assert answer.dtype == torch.long
+    assert (answer >= 0).all()
+    assert (answer < model.text_vocab_size).all()
+
+
+@requires_cuda
+@pytest.mark.parametrize(
+    ("enable_thinking", "expected_formatter_kwargs"),
+    [
+        (False, {"enable_thinking": False}),
+        (None, {}),
+    ],
+)
+def test_salm_automodel_generation_passes_enable_thinking(
+    model, monkeypatch, enable_thinking, expected_formatter_kwargs
+):
+    seen = {}
+
+    class _FakeFormatter:
+        def __init__(self, tokenizer):
+            pass
+
+        def encode_dialog(self, turns, **kwargs):
+            seen["turns"] = turns
+            seen["formatter_kwargs"] = kwargs
+            return {"input_ids": torch.tensor([1, 2], dtype=torch.long)}
+
+    def fake_generate(*, input_ids, attention_mask, generation_config, **kwargs):
+        seen["input_ids"] = input_ids
+        seen["attention_mask"] = attention_mask
+        max_new_tokens = kwargs["max_new_tokens"]
+        return torch.zeros((input_ids.shape[0], max_new_tokens), dtype=torch.long, device=input_ids.device)
+
+    monkeypatch.setattr(PromptFormatter, "resolve", staticmethod(lambda name: _FakeFormatter))
+    monkeypatch.setattr(model.llm, "generate", fake_generate, raising=False)
+
+    answer = model.generate(
+        prompts=[[{"role": "user", "slots": {"message": "test"}}]],
+        enable_thinking=enable_thinking,
+        max_new_tokens=3,
+    )
+
+    assert seen["formatter_kwargs"] == expected_formatter_kwargs
+    assert seen["turns"] == [{"role": "user", "slots": {"message": "test"}}]
+    assert seen["input_ids"].shape == (1, 2)
+    assert torch.equal(seen["attention_mask"], torch.ones_like(seen["input_ids"], dtype=torch.bool))
+    assert answer.shape == (1, 3)
+
+
+@requires_cuda
+def test_salm_automodel_generation_audios_via_prompt(model, tmp_path):
+    audio_path = tmp_path / "audio.wav"
+    dummy_cut(0, with_data=True).save_audio(audio_path)
+
+    answer = model.generate(
+        prompts=[
+            [{"role": "user", "content": f"Repeat after me: {AUDIO_LOCATOR_TAG}", "audio": [audio_path]}],
+            [
+                {
+                    "role": "user",
+                    "content": f"Repeat after me: {AUDIO_LOCATOR_TAG} and {AUDIO_LOCATOR_TAG}",
+                    "audio": [audio_path, audio_path],
+                }
+            ],
+        ],
+        generation_config=GenerationConfig(max_new_tokens=4),
+    )
+    assert answer.shape == (2, 4)
+    assert answer.dtype == torch.long
+    assert (answer >= 0).all()
+    assert (answer < model.text_vocab_size).all()
+
+
+@requires_cuda
+def test_salm_automodel_generation_prompts_as_tensor(model):
+    answer = model.generate(
+        prompts=torch.tensor([[1, 2, 3, 4, 5, 6, 7, model.audio_locator_tag_id]]),
+        audios=torch.randn(1, 16000),
+        audio_lens=torch.tensor([16000]),
+        max_new_tokens=4,
+    )
+    assert answer.shape == (1, 4)
+    assert answer.dtype == torch.long
+    assert (answer >= 0).all()
+    assert (answer < model.text_vocab_size).all()
+
+
+@pytest.mark.parametrize("device", chunking_test_devices())
+def test_salm_automodel_prepare_inputs_chunks_long_audio(device):
+    model = _make_chunking_test_model(encoder_chunk_size_seconds=1.0, sampling_rate=2, device=device)
+    spk_targets = torch.arange(10, dtype=torch.float32, device=device).reshape(1, 5, 2)
+    batch = {
+        "audios": torch.tensor([[1.0, 2.0, 3.0, 4.0, 5.0]], device=device),
+        "audio_lens": torch.tensor([5], dtype=torch.long, device=device),
+        "input_ids": torch.tensor([[model.audio_locator_tag_id, 10]], dtype=torch.long, device=device),
+        "loss_mask": torch.tensor([[False, True]], dtype=torch.bool, device=device),
+        "spk_targets": spk_targets,
+    }
+
+    inputs = model.prepare_inputs(batch)
+
+    chunked_signal, chunked_lens = model.perception.calls[0]
+    assert chunked_signal.shape == (2, 3)
+    assert torch.equal(chunked_lens, torch.tensor([2, 3], dtype=torch.long, device=device))
+    assert torch.equal(
+        model.perception.spk_targets_calls[0],
+        torch.stack([torch.cat([spk_targets[0, :2], spk_targets[0, 1:2]]), spk_targets[0, 2:5]]),
+    )
+    assert torch.equal(inputs["input_embeds"][0, :, 0], torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0], device=device))
+    assert torch.equal(inputs["attention_mask"], torch.ones((1, 5), dtype=torch.bool, device=device))
+
+
+@pytest.mark.parametrize("device", chunking_test_devices())
+def test_salm_automodel_prepare_inputs_merges_short_tail_chunk(device):
+    model = _make_chunking_test_model(encoder_chunk_size_seconds=0.5, sampling_rate=8, hop_length=2, device=device)
+    batch = {
+        "audios": torch.tensor([[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]], device=device),
+        "audio_lens": torch.tensor([9], dtype=torch.long, device=device),
+        "input_ids": torch.tensor([[model.audio_locator_tag_id, 10]], dtype=torch.long, device=device),
+        "loss_mask": torch.tensor([[False, True]], dtype=torch.bool, device=device),
+    }
+
+    inputs = model.prepare_inputs(batch)
+
+    chunked_signal, chunked_lens = model.perception.calls[0]
+    assert chunked_signal.shape == (2, 5)
+    assert torch.equal(chunked_lens, torch.tensor([4, 5], dtype=torch.long, device=device))
+    assert model.perception.spk_targets_calls[0] is None
+    assert torch.equal(
+        inputs["input_embeds"][0, :, 0],
+        torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0], device=device),
+    )
+
+
+@pytest.mark.parametrize("device", chunking_test_devices())
+def test_salm_automodel_prepare_inputs_skips_chunking_when_size_is_null(device):
+    model = _make_chunking_test_model(encoder_chunk_size_seconds=None, sampling_rate=2, device=device)
+    batch = {
+        "audios": torch.tensor([[1.0, 2.0, 3.0, 4.0, 5.0]], device=device),
+        "audio_lens": torch.tensor([5], dtype=torch.long, device=device),
+        "input_ids": torch.tensor([[model.audio_locator_tag_id, 10]], dtype=torch.long, device=device),
+        "loss_mask": torch.tensor([[False, True]], dtype=torch.bool, device=device),
+    }
+
+    model.prepare_inputs(batch)
+
+    input_signal, input_signal_lens = model.perception.calls[0]
+    assert input_signal.shape == (1, 5)
+    assert torch.equal(input_signal_lens, torch.tensor([5], dtype=torch.long, device=device))
+
+
+@pytest.mark.parametrize("device", chunking_test_devices())
+def test_salm_automodel_prepare_inputs_preserves_chunked_audio_order(device):
+    model = _make_chunking_test_model(encoder_chunk_size_seconds=1.0, sampling_rate=2, device=device)
+    batch = {
+        "audios": torch.tensor(
+            [
+                [1.0, 2.0, 3.0, 0.0, 0.0],
+                [10.0, 11.0, 12.0, 13.0, 14.0],
+            ],
+            device=device,
+        ),
+        "audio_lens": torch.tensor([3, 5], dtype=torch.long, device=device),
+        "input_ids": torch.tensor(
+            [[model.audio_locator_tag_id, model.audio_locator_tag_id, 10]], dtype=torch.long, device=device
+        ),
+        "loss_mask": torch.tensor([[False, False, True]], dtype=torch.bool, device=device),
+    }
+
+    inputs = model.prepare_inputs(batch)
+
+    assert torch.equal(
+        inputs["input_embeds"][0, :, 0],
+        torch.tensor([1.0, 2.0, 3.0, 10.0, 11.0, 12.0, 13.0, 14.0], device=device),
+    )
+
+
+@pytest.mark.parametrize("device", chunking_test_devices())
+def test_salm_automodel_generate_chunks_audio_before_llm(device):
+    model = _make_chunking_test_model(encoder_chunk_size_seconds=1.0, sampling_rate=2, device=device)
+    spk_targets = torch.arange(10, dtype=torch.float32, device=device).reshape(1, 5, 2)
+
+    answer = model.generate(
+        prompts=torch.tensor([[model.audio_locator_tag_id, 10]], dtype=torch.long, device=device),
+        audios=torch.tensor([[1.0, 2.0, 3.0, 4.0, 5.0]], device=device),
+        audio_lens=torch.tensor([5], dtype=torch.long, device=device),
+        spk_targets=spk_targets,
+        max_new_tokens=3,
+    )
+
+    chunked_signal, chunked_lens = model.perception.calls[0]
+    assert chunked_signal.shape == (2, 3)
+    assert torch.equal(chunked_lens, torch.tensor([2, 3], dtype=torch.long, device=device))
+    assert torch.equal(
+        model.perception.spk_targets_calls[0],
+        torch.stack([torch.cat([spk_targets[0, :2], spk_targets[0, 1:2]]), spk_targets[0, 2:5]]),
+    )
+    assert torch.equal(
+        model.llm.generate_kwargs["inputs_embeds"][0, :5, 0],
+        torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0], device=device),
+    )
+    assert answer.shape == (1, 3)
+
+
+def _make_chunking_test_model(encoder_chunk_size_seconds, sampling_rate, device, hop_length=1):
+    model = SALMAutomodel.__new__(SALMAutomodel)
+    torch.nn.Module.__init__(model)
+    model.cfg = {"encoder_chunk_size_seconds": encoder_chunk_size_seconds}
+    model.audio_locator_tag = AUDIO_LOCATOR_TAG
+    model.tokenizer = ChunkingTestTokenizer(AUDIO_LOCATOR_TAG)
+    model.llm = _AutomodelChunkingTestLLM(device=device)
+    model.perception = ChunkingTestPerception(sampling_rate=sampling_rate, hop_length=hop_length)
+    model._use_tp = False
+    return model
+
+
+class _AutomodelChunkingTestLLM(torch.nn.Module):
+    def __init__(self, device):
+        super().__init__()
+        self.model = torch.nn.Module()
+        self.model.embed_tokens = torch.nn.Embedding(128, 1, device=device)
+        self.generate_kwargs = None
+        with torch.no_grad():
+            self.model.embed_tokens.weight.zero_()
+
+    def generate(self, **kwargs):
+        self.generate_kwargs = kwargs
+        batch_size = kwargs["inputs_embeds"].shape[0]
+        max_new_tokens = kwargs["max_new_tokens"]
+        return torch.zeros((batch_size, max_new_tokens), dtype=torch.long, device=kwargs["inputs_embeds"].device)

tests/collections/speechlm2/test_salm_automodel_lora.py

@@ -0,0 +1,369 @@
+# Copyright (c) 2025, 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.
+"""Tests for SALM + Automodel LoRA integration."""
+import os
+import re
+
+import pytest
+import torch
+from lhotse import CutSet, SupervisionSegment
+from lhotse.testing.dummies import dummy_cut, dummy_recording
+from omegaconf import DictConfig, OmegaConf
+
+from nemo.collections.common.data.lhotse import NeMoMultimodalConversation
+from nemo.collections.common.data.lhotse.text_adapters import AudioTurn, TextTurn
+from nemo.collections.common.data.utils import move_data_to_device
+from nemo.collections.common.prompts import PromptFormatter
+from nemo.collections.speechlm2.data import SALMDataset
+from nemo.collections.speechlm2.models import SALMAutomodel
+from nemo.collections.speechlm2.parts.automodel_lora import (
+    LORA_PARAM_PATTERN,
+    ensure_lora_trainable,
+    make_peft_config,
+    maybe_install_lora,
+)
+
+if torch.cuda.is_available():
+    torch.set_default_device('cuda')
+
+
+def resolve_pretrained_models():
+    if os.path.exists("/home/TestData/speechlm/pretrained_models"):
+        return {
+            "pretrained_llm": "/home/TestData/speechlm/pretrained_models/Qwen--Qwen3-1.7B",
+            "pretrained_asr": "/home/TestData/speechlm/pretrained_models/canary-1b-flash.nemo",
+        }
+    else:
+        return {
+            "pretrained_asr": "nvidia/canary-1b-flash",
+            "pretrained_llm": "Qwen/Qwen3-1.7B",
+        }
+
+
+AUDIO_LOCATOR_TAG = "<|audioplaceholder|>"
+PROMPT = "qwen"
+
+LORA_CFG = {
+    "dim": 8,
+    "alpha": 16,
+    "dropout": 0.0,
+    "target_modules": ["q_proj", "v_proj"],
+}
+
+BASE_CFG = {
+    **resolve_pretrained_models(),
+    "pretrained_weights": False,
+    "prompt_format": PROMPT,
+    "audio_locator_tag": AUDIO_LOCATOR_TAG,
+    "freeze_params": ["^llm\\..+$", "^perception\\.preprocessor\\..+$", "^perception\\.encoder\\..+$"],
+    "prevent_freeze_params": [],
+    "lora": LORA_CFG,
+    "perception": {
+        "target": "nemo.collections.speechlm2.modules.perception.AudioPerceptionModule",
+        "output_dim": 2048,
+        "encoder": {
+            "_target_": "nemo.collections.asr.modules.ConformerEncoder",
+            "att_context_size": [-1, -1],
+            "causal_downsampling": False,
+            "conv_context_size": None,
+            "conv_kernel_size": 9,
+            "conv_norm_type": "batch_norm",
+            "d_model": 1024,
+            "dropout": 0.1,
+            "dropout_att": 0.1,
+            "dropout_emb": 0.0,
+            "dropout_pre_encoder": 0.1,
+            "feat_in": 128,
+            "feat_out": -1,
+            "ff_expansion_factor": 4,
+            "n_heads": 8,
+            "n_layers": 2,
+            "pos_emb_max_len": 5000,
+            "self_attention_model": "rel_pos",
+            "subsampling": "dw_striding",
+            "subsampling_conv_channels": 256,
+            "subsampling_factor": 8,
+        },
+        "modality_adapter": {
+            "_target_": "nemo.collections.speechlm2.modules.perception.IdentityConnector",
+            "d_model": 1024,
+        },
+        "preprocessor": {
+            "_target_": "nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor",
+            "dither": 1e-05,
+            "features": 128,
+            "frame_splicing": 1,
+            "log": True,
+            "n_fft": 512,
+            "normalize": "per_feature",
+            "pad_to": 0,
+            "pad_value": 0.0,
+            "sample_rate": 16000,
+            "window": "hann",
+            "window_size": 0.025,
+            "window_stride": 0.01,
+        },
+    },
+    "optimizer": {"_target_": "torch.optim.AdamW"},
+    "torch_dtype": "bfloat16",
+}
+
+
+# ---------------------------------------------------------------------------
+# Unit tests for automodel_lora helpers (no model needed)
+# ---------------------------------------------------------------------------
+
+
+class TestMakePeftConfig:
+    def test_basic(self):
+        cfg = DictConfig({"dim": 32, "alpha": 64, "target_modules": ["q_proj"]})
+        pc = make_peft_config(cfg)
+        assert pc.dim == 32
+        assert pc.alpha == 64
+        assert pc.target_modules == ["*.q_proj"]
+
+    def test_defaults(self):
+        cfg = DictConfig({})
+        pc = make_peft_config(cfg)
+        assert pc is None
+
+    def test_all_fields(self):
+        cfg = DictConfig(
+            {
+                "dim": 16,
+                "alpha": 32,
+                "dropout": 0.1,
+                "target_modules": ["q_proj", "v_proj"],
+                "exclude_modules": ["lm_head"],
+                "match_all_linear": False,
+                "use_dora": True,
+                "dropout_position": "pre",
+                "lora_A_init": "kaiming_uniform",
+                "use_triton": False,
+            }
+        )
+        pc = make_peft_config(cfg)
+        assert pc.dim == 16
+        assert pc.use_dora is True
+        assert pc.dropout_position == "pre"
+        assert pc.lora_A_init == "kaiming_uniform"
+        assert pc.exclude_modules == ["*.lm_head"]
+
+    def test_short_names_get_wildcard_prefix(self):
+        """Short leaf names like 'q_proj' should be auto-prefixed with '*.' so
+        that automodel's ModuleMatcher matches them against full dotted paths."""
+        cfg = DictConfig({"target_modules": ["q_proj", "v_proj"], "exclude_modules": ["lm_head"]})
+        pc = make_peft_config(cfg)
+        assert pc.target_modules == ["*.q_proj", "*.v_proj"]
+        assert pc.exclude_modules == ["*.lm_head"]
+
+    def test_already_qualified_names_unchanged(self):
+        """Names that already contain '*' or '.' should not be double-prefixed."""
+        cfg = DictConfig({"target_modules": ["*.q_proj", "model.layers.*.self_attn.v_proj"]})
+        pc = make_peft_config(cfg)
+        assert pc.target_modules == ["*.q_proj", "model.layers.*.self_attn.v_proj"]
+
+
+class TestEnsureLoraTrainable:
+    def _model_stub(self, cfg_dict):
+        class Stub:
+            pass
+
+        s = Stub()
+        s.cfg = DictConfig(cfg_dict)
+        return s
+
+    def test_creates_prevent_freeze_params(self):
+        model = self._model_stub({})
+        ensure_lora_trainable(model)
+        assert LORA_PARAM_PATTERN in model.cfg.prevent_freeze_params
+
+    def test_appends_to_existing(self):
+        model = self._model_stub({"prevent_freeze_params": ["^some_other$"]})
+        ensure_lora_trainable(model)
+        assert len(model.cfg.prevent_freeze_params) == 2
+        assert LORA_PARAM_PATTERN in model.cfg.prevent_freeze_params
+
+    def test_idempotent(self):
+        model = self._model_stub({"prevent_freeze_params": []})
+        ensure_lora_trainable(model)
+        ensure_lora_trainable(model)
+        assert model.cfg.prevent_freeze_params.count(LORA_PARAM_PATTERN) == 1
+
+
+class TestLoraParamPattern:
+    """Verify the regex matches typical LoRA parameter names."""
+
+    @pytest.mark.parametrize(
+        "name",
+        [
+            "llm.model.layers.0.self_attn.q_proj.lora_A.weight",
+            "llm.model.layers.0.self_attn.v_proj.lora_B.weight",
+            "llm.model.layers.31.mlp.gate_proj.lora_A.weight",
+        ],
+    )
+    def test_matches_lora_params(self, name):
+        assert re.compile(LORA_PARAM_PATTERN).match(name)
+
+    @pytest.mark.parametrize(
+        "name",
+        [
+            "llm.model.layers.0.self_attn.q_proj.weight",
+            "perception.encoder.layers.0.weight",
+            "llm.lm_head.weight",
+        ],
+    )
+    def test_does_not_match_base_params(self, name):
+        assert re.compile(LORA_PARAM_PATTERN).match(name) is None
+
+
+# ---------------------------------------------------------------------------
+# Integration tests (require CUDA — Automodel's from_config needs a GPU)
+# ---------------------------------------------------------------------------
+
+requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="Automodel requires CUDA")
+
+
+@pytest.fixture(scope="module")
+def salm_with_lora():
+    if not torch.cuda.is_available():
+        pytest.skip("Automodel requires CUDA")
+    model = SALMAutomodel(BASE_CFG)
+    model.configure_model()
+    model.to("cuda")
+    return model
+
+
+@pytest.fixture(scope="module")
+def salm_without_lora():
+    if not torch.cuda.is_available():
+        pytest.skip("Automodel requires CUDA")
+    cfg = {k: v for k, v in BASE_CFG.items() if k != "lora"}
+    model = SALMAutomodel(cfg)
+    model.configure_model()
+    model.to("cuda")
+    return model
+
+
+@pytest.fixture(scope="module")
+def dataset(salm_with_lora):
+    return SALMDataset(salm_with_lora.tokenizer)
+
+
+@pytest.fixture(scope="module")
+def prompt_formatter(salm_with_lora):
+    return PromptFormatter.resolve(PROMPT)(salm_with_lora.tokenizer)
+
+
+@pytest.fixture(scope="module")
+def training_cutset_batch():
+    cut = dummy_cut(0, recording=dummy_recording(0, with_data=True))
+    cut.supervisions = [
+        SupervisionSegment(
+            id=cut.id, recording_id=cut.recording_id, start=0, duration=1.0, text='Some text transcription.'
+        )
+    ]
+    return CutSet(
+        [
+            NeMoMultimodalConversation(
+                id="example-0",
+                turns=[
+                    TextTurn(role="user", value="Repeat after me:"),
+                    AudioTurn(role="user", cut=cut, audio_locator_tag=AUDIO_LOCATOR_TAG),
+                    TextTurn(role="assistant", value=cut.supervisions[0].text),
+                ],
+                token_equivalent_duration=0.08,
+            )
+        ]
+    )
+
+
+@requires_cuda
+def test_lora_params_exist(salm_with_lora):
+    """After configure_model with lora config, LLM should have lora_A/lora_B parameters."""
+    lora_params = {n for n, _ in salm_with_lora.named_parameters() if ".lora_" in n}
+    assert len(lora_params) > 0, "No LoRA parameters found"
+    # We targeted q_proj and v_proj, so each should have lora_A.weight and lora_B.weight
+    assert any("q_proj.lora_A" in n for n in lora_params)
+    assert any("q_proj.lora_B" in n for n in lora_params)
+    assert any("v_proj.lora_A" in n for n in lora_params)
+    assert any("v_proj.lora_B" in n for n in lora_params)
+
+
+@requires_cuda
+def test_no_lora_params_without_config(salm_without_lora):
+    """Without lora config, no LoRA parameters should exist."""
+    lora_params = [n for n, _ in salm_without_lora.named_parameters() if ".lora_" in n]
+    assert len(lora_params) == 0
+
+
+@requires_cuda
+def test_lora_prevent_freeze_pattern_set(salm_with_lora):
+    """The prevent_freeze_params list should include the LoRA pattern."""
+    assert LORA_PARAM_PATTERN in salm_with_lora.cfg.prevent_freeze_params
+
+
+@requires_cuda
+def test_lora_params_are_trainable(salm_with_lora):
+    """LoRA parameters should have requires_grad=True."""
+    for name, param in salm_with_lora.named_parameters():
+        if ".lora_" in name:
+            assert param.requires_grad, f"LoRA param {name} is not trainable"
+
+
+@requires_cuda
+def test_base_llm_params_not_patched_elsewhere(salm_with_lora):
+    """Only targeted modules (q_proj, v_proj) should have LoRA, not others like k_proj."""
+    lora_params = {n for n, _ in salm_with_lora.named_parameters() if ".lora_" in n}
+    for n in lora_params:
+        assert "q_proj" in n or "v_proj" in n, f"Unexpected LoRA param on non-targeted module: {n}"
+
+
+@requires_cuda
+def test_lora_training_step(salm_with_lora, dataset, prompt_formatter, training_cutset_batch):
+    training_cutset_batch = training_cutset_batch.map(lambda c: c.apply_prompt_format(prompt_formatter), apply_fn=None)
+    batch = dataset[training_cutset_batch]
+    batch = move_data_to_device(batch, device=salm_with_lora.device)
+    results = salm_with_lora.training_step(batch, batch_idx=0)
+    assert torch.is_tensor(results["loss"])
+    assert not torch.isnan(results["loss"])
+    assert results["loss"] > 0
+
+
+@requires_cuda
+def test_lora_generation(salm_with_lora):
+    answer = salm_with_lora.generate(
+        prompts=[
+            [{"role": "user", "slots": {"message": f"Repeat after me: {AUDIO_LOCATOR_TAG}"}}],
+        ],
+        audios=torch.randn(1, 16000),
+        audio_lens=torch.tensor([16000]),
+        max_new_tokens=4,
+    )
+    assert answer.shape == (1, 4)
+    assert answer.dtype == torch.long
+
+
+@requires_cuda
+def test_lora_match_all_linear():
+    """Test match_all_linear=True applies LoRA to all linear layers in the LLM."""
+    cfg = {**BASE_CFG}
+    cfg["lora"] = {"dim": 8, "alpha": 16, "match_all_linear": True}
+    model = SALMAutomodel(cfg)
+    model.configure_model()
+
+    lora_params = {n for n, _ in model.named_parameters() if ".lora_" in n}
+    # Should have LoRA on more than just q_proj/v_proj
+    modules_with_lora = {n.rsplit(".lora_", 1)[0] for n in lora_params}
+    assert len(modules_with_lora) > 2, f"match_all_linear should patch many modules, got: {modules_with_lora}"

tests/collections/speechlm2/test_salm_automodel_pee.py

@@ -0,0 +1,391 @@
+# 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.
+"""SALMAutomodel tests for the Parallel Expert Encoder (PEE) recipe.
+
+Mirrors ``test_salm_automodel.py`` but swaps the plain ConformerEncoder for a
+``ParallelExpertEncoder`` (streaming Sortformer diarizer + ASR Conformer, fused).
+The tiny-but-real PE encoder is the same dummy bundle built in
+``tests/collections/asr/test_parallel_expert_encoder.py`` (``build_toy_pe_encoder``);
+this file mounts it onto ``model.perception.encoder`` exactly the way
+``nemo.collections.speechlm2.parts.pretrained.setup_parallel_expert_encoder`` does
+and exercises the SALM training / validation / generation path end to end, plus the
+``spk_targets -> spk_targets`` routing that is unique to the PEE recipe.
+"""
+import importlib.util
+import os
+from types import SimpleNamespace
+
+import pytest
+import torch
+from lhotse import CutSet, SupervisionSegment
+from lhotse.testing.dummies import dummy_cut, dummy_recording
+from omegaconf import DictConfig
+
+from nemo.collections.asr.modules.parallel_expert_encoder import ParallelExpertEncoder
+from nemo.collections.common.data.lhotse import NeMoMultimodalConversation
+from nemo.collections.common.data.lhotse.text_adapters import AudioTurn, TextTurn
+from nemo.collections.common.data.utils import move_data_to_device
+from nemo.collections.common.prompts import PromptFormatter
+from nemo.collections.speechlm2.data import SALMDataset
+from nemo.collections.speechlm2.models import SALMAutomodel
+
+# Reuse the toy PE encoder (and its dimensions) defined for the standalone
+# ParallelExpertEncoder tests, so both suites share one dummy bundle definition.
+from tests.collections.asr.test_parallel_expert_encoder import (
+    _ASR_D_MODEL,
+    _MEL_FEATURES,
+    _N_SPK,
+    _SUBSAMPLING_FACTOR,
+    build_toy_pe_encoder,
+)
+
+# SALMAutomodel.configure_model() pulls in the (gitignored) nemo_automodel package,
+# so the full-model tests need both CUDA and that dependency to be importable.
+# Build a precise skip reason that names only the piece that is actually missing
+# (so a CUDA-equipped box without nemo_automodel doesn't get a misleading "requires CUDA").
+_HAS_CUDA = torch.cuda.is_available()
+_HAS_AUTOMODEL = importlib.util.find_spec("nemo_automodel") is not None
+_MISSING = [name for name, ok in (("CUDA", _HAS_CUDA), ("the nemo_automodel package", _HAS_AUTOMODEL)) if not ok]
+requires_cuda = pytest.mark.skipif(
+    bool(_MISSING),
+    reason="SALMAutomodel needs " + " and ".join(_MISSING) if _MISSING else "",
+)
+
+# NOTE: deliberately do NOT call torch.set_default_device('cuda') here. It is a
+# global, process-wide mutation that leaks into other test modules collected in
+# the same pytest session (e.g. the device-agnostic CPU tests in
+# tests/collections/asr/test_parallel_expert_encoder.py), causing spurious
+# cuda/cpu device-mismatch failures. CUDA tests below place their tensors on
+# model.device explicitly instead.
+
+
+def resolve_pretrained_models():
+    if os.path.exists("/home/TestData/speechlm/pretrained_models"):
+        # CI pre-cached paths:
+        return {
+            "pretrained_llm": "/home/TestData/speechlm/pretrained_models/Qwen--Qwen3-1.7B",
+            "pretrained_asr": "/home/TestData/speechlm/pretrained_models/canary-1b-flash.nemo",
+        }
+    else:
+        # HF URLs:
+        return {
+            "pretrained_asr": "nvidia/canary-1b-flash",
+            "pretrained_llm": "Qwen/Qwen3-1.7B",
+        }
+
+
+AUDIO_LOCATOR_TAG = "<|audioplaceholder|>"
+PROMPT = "qwen"
+
+# Speaker-activity (SOT) target settings; num_speakers matches the PE encoder's n_spk.
+SOT_CFG = {
+    "num_speakers": _N_SPK,
+    "sample_rate": 16000,
+    "window_stride": 0.01,
+    "subsampling_factor": _SUBSAMPLING_FACTOR,
+    "no_rttm_to_ones": True,
+}
+
+
+def mount_dummy_pe_encoder(model: SALMAutomodel) -> ParallelExpertEncoder:
+    """Mount a dummy ParallelExpertEncoder onto ``model.perception.encoder``.
+
+    Replicates ``setup_parallel_expert_encoder`` without needing an on-disk
+    ``.nemo`` bundle: it swaps the perception encoder for a tiny-but-real PE
+    encoder, disables the outer preprocessor normalization (the PE encoder
+    re-applies ASR normalization internally), and matches the model dtype/device.
+    """
+    pe_encoder = build_toy_pe_encoder()
+    # The PE encoder consumes un-normalised mels; turn off the perception
+    # preprocessor's normalization just like the real mounting helper does.
+    model.perception.preprocessor.featurizer.normalize = None
+    model.perception.encoder = pe_encoder
+    target_dtype = next(model.llm.parameters()).dtype
+    model.perception.to(device=model.device, dtype=target_dtype)
+    return pe_encoder
+
+
+@pytest.fixture(scope="session")
+def model():
+    if _MISSING:
+        pytest.skip("SALMAutomodel needs " + " and ".join(_MISSING))
+    cfg = {
+        **resolve_pretrained_models(),
+        "pretrained_weights": False,
+        "prompt_format": PROMPT,
+        "audio_locator_tag": AUDIO_LOCATOR_TAG,
+        "perception": {
+            "target": "nemo.collections.speechlm2.modules.perception.AudioPerceptionModule",
+            "output_dim": 2048,
+            # Placeholder encoder; its only job is to define a d_model that matches
+            # the dummy PE encoder we mount in its place (mount validation checks this).
+            "encoder": {
+                "_target_": "nemo.collections.asr.modules.ConformerEncoder",
+                "att_context_size": [-1, -1],
+                "causal_downsampling": False,
+                "conv_context_size": None,
+                "conv_kernel_size": 9,
+                "conv_norm_type": "batch_norm",
+                "d_model": _ASR_D_MODEL,
+                "dropout": 0.0,
+                "dropout_att": 0.0,
+                "dropout_emb": 0.0,
+                "dropout_pre_encoder": 0.0,
+                "feat_in": _MEL_FEATURES,
+                "feat_out": -1,
+                "ff_expansion_factor": 4,
+                "n_heads": 4,
+                "n_layers": 1,
+                "pos_emb_max_len": 5000,
+                "self_attention_model": "rel_pos",
+                "subsampling": "dw_striding",
+                "subsampling_conv_channels": 16,
+                "subsampling_factor": _SUBSAMPLING_FACTOR,
+            },
+            "modality_adapter": {
+                "_target_": "nemo.collections.speechlm2.modules.perception.IdentityConnector",
+                "d_model": _ASR_D_MODEL,
+            },
+            "preprocessor": {
+                "_target_": "nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor",
+                "dither": 1e-05,
+                "features": _MEL_FEATURES,
+                "frame_splicing": 1,
+                "log": True,
+                "n_fft": 512,
+                "normalize": "per_feature",
+                "pad_to": 0,
+                "pad_value": 0.0,
+                "sample_rate": 16000,
+                "window": "hann",
+                "window_size": 0.025,
+                "window_stride": 0.01,
+            },
+        },
+        "optimizer": {"_target_": "torch.optim.AdamW"},
+        "torch_dtype": "bfloat16",
+    }
+    model = SALMAutomodel(cfg)
+    model.configure_model()
+    model.to("cuda")
+    mount_dummy_pe_encoder(model)
+    return model
+
+
+@pytest.fixture(scope="session")
+def dataset(model):
+    # SOT mode: each batch additionally carries RTTM-derived `spk_targets`.
+    return SALMDataset(model.tokenizer, multispeaker_cfg=SOT_CFG)
+
+
+@pytest.fixture(scope="session")
+def prompt_formatter(model):
+    return PromptFormatter.resolve(PROMPT)(model.tokenizer)
+
+
+@pytest.fixture(scope="session")
+def training_cutset_batch():
+    cut = dummy_cut(0, recording=dummy_recording(0, with_data=True))
+    cut.supervisions = [
+        SupervisionSegment(
+            id=cut.id, recording_id=cut.recording_id, start=0, duration=1.0, text='Some text transcription.'
+        )
+    ]
+    return CutSet(
+        [
+            NeMoMultimodalConversation(
+                id="example-0",
+                turns=[
+                    TextTurn(role="user", value="Repeat after me:"),
+                    AudioTurn(role="user", cut=cut, audio_locator_tag=AUDIO_LOCATOR_TAG),
+                    TextTurn(role="assistant", value=cut.supervisions[0].text),
+                ],
+                token_equivalent_duration=0.08,
+            )
+        ]
+    )
+
+
+@requires_cuda
+def test_salm_automodel_pee_uses_parallel_expert_encoder(model):
+    assert isinstance(model.perception.encoder, ParallelExpertEncoder)
+    assert model._uses_parallel_expert_encoder()
+    # The mounted PE encoder is a drop-in: its d_model drives the perception output projection.
+    assert model.perception.encoder.d_model == _ASR_D_MODEL
+    assert model.perception.encoder.n_spk == _N_SPK
+    # PE encoder consumes un-normalised mels; the outer preprocessor normalization is disabled.
+    assert model.perception.preprocessor.featurizer.normalize is None
+
+
+@requires_cuda
+def test_salm_automodel_pee_dataset_emits_speaker_targets(dataset, prompt_formatter, training_cutset_batch):
+    training_cutset_batch = training_cutset_batch.map(lambda c: c.apply_prompt_format(prompt_formatter), apply_fn=None)
+    batch = dataset[training_cutset_batch]
+    for key in ("audios", "audio_lens", "input_ids", "loss_mask", "spk_targets"):
+        assert key in batch
+        assert torch.is_tensor(batch[key])
+    # spk_targets: (B, T_target, num_speakers) -> last dim must match the PE encoder n_spk.
+    assert batch["spk_targets"].ndim == 3
+    assert batch["spk_targets"].shape[0] == batch["audios"].shape[0]
+    assert batch["spk_targets"].shape[-1] == _N_SPK
+
+
+@requires_cuda
+def test_salm_automodel_pee_training_step(model, dataset, prompt_formatter, training_cutset_batch):
+    training_cutset_batch = training_cutset_batch.map(lambda c: c.apply_prompt_format(prompt_formatter), apply_fn=None)
+    batch = dataset[training_cutset_batch]
+    assert "spk_targets" in batch  # injected as spk_targets into the PE encoder during training
+    batch = move_data_to_device(batch, device=model.device)
+    results = model.training_step(batch, batch_idx=0)
+    assert torch.is_tensor(results["loss"])
+    assert not torch.isnan(results["loss"])
+    assert results["loss"] > 0
+
+
+@requires_cuda
+def test_salm_automodel_pee_validation_step(model, dataset, prompt_formatter, training_cutset_batch):
+    model.on_validation_epoch_start()
+    training_cutset_batch = training_cutset_batch.map(lambda c: c.apply_prompt_format(prompt_formatter), apply_fn=None)
+    batch = dataset[training_cutset_batch]
+    batch = move_data_to_device(batch, device=model.device)
+    # Validation ignores spk_targets and lets the PE encoder run its embedded Sortformer.
+    results = model.validation_step({"dummy_val_set": batch}, batch_idx=0)
+    assert results is None
+
+
+@requires_cuda
+def test_salm_automodel_pee_generation(model):
+    # No spk_targets at inference -> the PE encoder predicts diarization internally.
+    answer = model.generate(
+        prompts=[
+            [
+                {"role": "user", "slots": {"message": f"Repeat after me: {AUDIO_LOCATOR_TAG}"}},
+            ]
+        ],
+        audios=torch.randn(1, 16000, device=model.device),
+        audio_lens=torch.tensor([16000], device=model.device),
+        max_new_tokens=4,
+    )
+    assert answer.shape == (1, 4)
+    assert answer.dtype == torch.long
+    assert (answer >= 0).all()
+    assert (answer < model.text_vocab_size).all()
+
+
+# ----------------------------------------------------------------------------- #
+# CPU-only: spk_targets -> spk_targets routing in prepare_inputs.
+#
+# Uses a bare SALMAutomodel (no LLM/ASR download) with a stub perception whose
+# `.encoder` is the real dummy ParallelExpertEncoder, so `_uses_parallel_expert_encoder`
+# takes the PEE branch and we can assert how speaker targets are forwarded.
+# ----------------------------------------------------------------------------- #
+class _PEETestTokenizer:
+    pad = 0
+    unk_id = None
+    bos_id = 1
+    eos_id = 2
+
+    def __init__(self, audio_locator_tag):
+        self._audio_locator_tag = audio_locator_tag
+
+    def token_to_id(self, token):
+        assert token == self._audio_locator_tag
+        return 99
+
+
+class _PEETestLLM(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.model = torch.nn.Module()
+        self.model.embed_tokens = torch.nn.Embedding(128, 1)
+        with torch.no_grad():
+            self.model.embed_tokens.weight.zero_()
+
+
+class _PEETestPerception(torch.nn.Module):
+    """Stub perception that exposes a real PE encoder and records the spk_targets it receives."""
+
+    def __init__(self, pe_encoder: ParallelExpertEncoder):
+        super().__init__()
+        self.encoder = pe_encoder  # real ParallelExpertEncoder -> drives the PEE branch
+        self.preprocessor = SimpleNamespace(featurizer=SimpleNamespace(sample_rate=16000, hop_length=160))
+        self.spk_targets_calls = []
+
+    def forward(self, input_signal=None, input_signal_length=None, spk_targets=None):
+        self.spk_targets_calls.append(spk_targets)
+        max_len = int(input_signal_length.max().item())
+        return input_signal[:, :max_len].unsqueeze(-1), input_signal_length.clone()
+
+
+def _make_pee_routing_test_model(pe_encoder, cfg=None):
+    model = SALMAutomodel.__new__(SALMAutomodel)
+    torch.nn.Module.__init__(model)
+    model.cfg = DictConfig(cfg or {"encoder_chunk_size_seconds": None})
+    model.audio_locator_tag = AUDIO_LOCATOR_TAG
+    model.tokenizer = _PEETestTokenizer(AUDIO_LOCATOR_TAG)
+    model.llm = _PEETestLLM()
+    model.perception = _PEETestPerception(pe_encoder)
+    model._use_tp = False
+    return model
+
+
+@pytest.fixture(scope="module")
+def dummy_pe_encoder():
+    return build_toy_pe_encoder().eval()
+
+
+@pytest.mark.unit
+def test_pee_prepare_inputs_detects_parallel_expert_encoder(dummy_pe_encoder):
+    model = _make_pee_routing_test_model(dummy_pe_encoder)
+    assert model._uses_parallel_expert_encoder()
+
+
+@pytest.mark.unit
+def test_pee_prepare_inputs_routes_spk_targets_as_spk_targets(dummy_pe_encoder):
+    model = _make_pee_routing_test_model(dummy_pe_encoder)
+    spk_targets = torch.rand(1, 4, _N_SPK)
+    batch = {
+        "audios": torch.tensor([[1.0, 2.0, 3.0, 4.0, 5.0]]),
+        "audio_lens": torch.tensor([5], dtype=torch.long),
+        "input_ids": torch.tensor([[model.audio_locator_tag_id, 10]], dtype=torch.long),
+        "loss_mask": torch.tensor([[False, True]], dtype=torch.bool),
+        "spk_targets": spk_targets,
+    }
+
+    # RTTM-derived spk_targets are forwarded when present in the batch.
+    model.prepare_inputs(batch)
+    assert model.perception.spk_targets_calls[-1] is spk_targets
+
+    # No spk_targets key means the embedded Sortformer predicts speaker activity.
+    batch_without_spk_targets = {k: v for k, v in batch.items() if k != "spk_targets"}
+    model.prepare_inputs(batch_without_spk_targets)
+    assert model.perception.spk_targets_calls[-1] is None
+
+
+@pytest.mark.unit
+def test_pee_prepare_inputs_warns_for_experimental_inference_options(dummy_pe_encoder):
+    model = _make_pee_routing_test_model(
+        dummy_pe_encoder,
+        cfg={"pe_encoder_path": "/tmp/pee.nemo", "encoder_chunk_size_seconds": 30.0},
+    )
+    batch = {
+        "audios": torch.tensor([[1.0, 2.0, 3.0, 4.0, 5.0]]),
+        "audio_lens": torch.tensor([5], dtype=torch.long),
+        "input_ids": torch.tensor([[model.audio_locator_tag_id, 10]], dtype=torch.long),
+        "loss_mask": torch.tensor([[False, True]], dtype=torch.bool),
+    }
+
+    with pytest.warns(UserWarning, match="ParallelExpertEncoder inference path.*encoder_chunk_size_seconds"):
+        model.prepare_inputs(batch)
+    assert model.perception.spk_targets_calls[-1] is None

tests/collections/speechlm2/test_salm_cp_helpers.py

@@ -0,0 +1,201 @@
+# Copyright (c) 2025, 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.
+"""CPU-only tests for the CP-helper module.
+
+The ``cp_size > 1`` path in ``encode_audio_with_cp_distribution`` requires
+a real ``torch.distributed`` process group; it's exercised by the 2-GPU
+smoke. These tests cover the fallback contracts that run on every machine
+(``cp_mesh is None``, ``B_aud == 0``).
+"""
+import pytest
+import torch
+
+from nemo.collections.speechlm2.parts.cp_helpers import encode_audio_with_cp_distribution, get_cp_mesh
+
+
+def test_get_cp_mesh_none():
+    assert get_cp_mesh(None) == (None, 1, 0)
+
+
+class _DummyCpDim:
+    """Stand-in for ``device_mesh['cp']`` whose ``.size()`` is 1 (CP inactive)."""
+
+    def size(self):
+        return 1
+
+
+class _DummyDeviceMesh:
+    """Minimal ``DeviceMesh``-like object exposing only the bits ``get_cp_mesh`` reads."""
+
+    def __init__(self, cp_size: int = 1, has_cp: bool = True):
+        self.mesh_dim_names = ("dp", "cp", "tp") if has_cp else ("dp", "tp")
+        self._cp_size = cp_size
+
+    def __getitem__(self, key):
+        if key == "cp":
+
+            class _Dim:
+                def __init__(self, size):
+                    self._size = size
+
+                def size(self):
+                    return self._size
+
+            return _Dim(self._cp_size)
+        raise KeyError(key)
+
+
+def test_get_cp_mesh_cp_size_one():
+    assert get_cp_mesh(_DummyDeviceMesh(cp_size=1)) == (None, 1, 0)
+
+
+def test_get_cp_mesh_no_cp_dim():
+    assert get_cp_mesh(_DummyDeviceMesh(has_cp=False)) == (None, 1, 0)
+
+
+class _PerceptionStub:
+    """Stand-in for ``self.perception``: returns a deterministic embedding per audio."""
+
+    def __init__(self, hidden_size: int = 4):
+        self.hidden_size = hidden_size
+        self.spk_targets_calls = []
+
+    def __call__(self, *, input_signal, input_signal_length, spk_targets=None):
+        self.spk_targets_calls.append(None if spk_targets is None else spk_targets.detach().clone())
+        # Pretend each audio of length L produces L // 2 frames of embeddings;
+        # encode the row index into the first column so we can verify ordering.
+        B, T = input_signal.shape
+        if B == 0:
+            return torch.zeros(0, 0, self.hidden_size, dtype=torch.float32), input_signal_length
+        # Frame count per row scales with audio_lens.
+        out_lens = (input_signal_length // 2).clamp(min=1)
+        max_out = int(out_lens.max().item())
+        embs = torch.zeros(B, max_out, self.hidden_size, dtype=torch.float32)
+        for i in range(B):
+            embs[i, : int(out_lens[i].item()), 0] = float(i)  # marker
+        return embs, out_lens
+
+
+def test_encode_audio_no_cp_returns_unpadded_list():
+    perception = _PerceptionStub(hidden_size=4)
+    audios = torch.zeros(3, 1600, dtype=torch.float32)
+    audio_lens = torch.tensor([800, 1200, 1600], dtype=torch.long)
+    embs = encode_audio_with_cp_distribution(
+        perception,
+        audios,
+        audio_lens,
+        chunk_size_seconds=None,
+        sampling_rate=16000,
+        cp_mesh=None,
+    )
+    # 3 audios → 3 embedding tensors with row-specific lengths.
+    assert len(embs) == 3
+    assert perception.spk_targets_calls[-1] is None
+    expected_lens = [400, 600, 800]
+    for i, e in enumerate(embs):
+        assert e.shape == (expected_lens[i], 4)
+        # Marker preserved.
+        assert torch.all(e[:, 0] == float(i))
+
+
+def test_encode_audio_empty_batch_returns_empty():
+    perception = _PerceptionStub()
+    audios = torch.zeros(0, 1600, dtype=torch.float32)
+    audio_lens = torch.zeros(0, dtype=torch.long)
+    embs = encode_audio_with_cp_distribution(
+        perception,
+        audios,
+        audio_lens,
+        chunk_size_seconds=None,
+        sampling_rate=16000,
+        cp_mesh=None,
+    )
+    assert embs == []
+
+
+def test_encode_audio_no_cp_forwards_spk_targets():
+    perception = _PerceptionStub(hidden_size=4)
+    audios = torch.zeros(2, 1600, dtype=torch.float32)
+    audio_lens = torch.tensor([800, 1600], dtype=torch.long)
+    spk_targets = torch.arange(16, dtype=torch.float32).reshape(2, 4, 2)
+    encode_audio_with_cp_distribution(
+        perception,
+        audios,
+        audio_lens,
+        chunk_size_seconds=None,
+        sampling_rate=16000,
+        cp_mesh=None,
+        spk_targets=spk_targets,
+    )
+    assert torch.equal(perception.spk_targets_calls[-1], spk_targets)
+
+
+class _FakeCpMesh:
+    def size(self):
+        return 2
+
+    def get_group(self):
+        return "fake-cp-group"
+
+
+class _TrainablePerceptionStub(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.scale = torch.nn.Parameter(torch.tensor(2.0))
+        self.spk_targets_calls = []
+
+    def forward(self, *, input_signal, input_signal_length, spk_targets=None):
+        self.spk_targets_calls.append(None if spk_targets is None else spk_targets.detach().clone())
+        B = input_signal.shape[0]
+        embs = input_signal[:, :2].unsqueeze(-1) * self.scale
+        lens = torch.full((B,), 2, dtype=input_signal_length.dtype, device=input_signal_length.device)
+        return embs, lens
+
+
+def test_encode_audio_cp_distribution_preserves_local_autograd(monkeypatch):
+    perception = _TrainablePerceptionStub()
+    audios = torch.tensor([[1.0, 2.0, 0.0], [3.0, 4.0, 0.0]])
+    audio_lens = torch.tensor([3, 3], dtype=torch.long)
+
+    def fake_all_gather(local_stack, group):
+        assert group == "fake-cp-group"
+        remote_stack = torch.zeros_like(local_stack)
+        return (local_stack, remote_stack)
+
+    def fake_lens_all_gather(gathered_lens, local_lens, group):
+        assert group == "fake-cp-group"
+        gathered_lens[0].copy_(local_lens)
+        gathered_lens[1].fill_(2)
+
+    monkeypatch.setattr("nemo.collections.speechlm2.parts.cp_helpers.dist.get_rank", lambda group: 0)
+    monkeypatch.setattr("nemo.collections.speechlm2.parts.cp_helpers.dist.all_reduce", lambda *args, **kwargs: None)
+    monkeypatch.setattr("nemo.collections.speechlm2.parts.cp_helpers.dist.all_gather", fake_lens_all_gather)
+    monkeypatch.setattr("nemo.collections.speechlm2.parts.cp_helpers.differentiable_all_gather", fake_all_gather)
+    spk_targets = torch.arange(8, dtype=torch.float32).reshape(2, 2, 2)
+
+    embs = encode_audio_with_cp_distribution(
+        perception,
+        audios,
+        audio_lens,
+        chunk_size_seconds=None,
+        sampling_rate=16000,
+        cp_mesh=_FakeCpMesh(),
+        spk_targets=spk_targets,
+    )
+
+    assert torch.equal(perception.spk_targets_calls[-1], spk_targets[:1])
+    assert embs[0].requires_grad
+    embs[0].sum().backward()
+    assert perception.scale.grad is not None
+    assert perception.scale.grad.item() == pytest.approx(3.0)

tests/collections/speechlm2/test_salm_dataset.py

@@ -0,0 +1,98 @@
+# 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.
+
+import pytest
+import torch
+from lhotse import CutSet, SupervisionSegment
+from lhotse.testing.dummies import dummy_cut, dummy_recording
+
+import nemo.collections.speechlm2.data.salm_dataset as salm_dataset_module
+from nemo.collections.common.data.lhotse import NeMoMultimodalConversation
+from nemo.collections.common.data.lhotse.text_adapters import AudioTurn, TextTurn
+
+
+class _Tokenizer:
+    pad = 0
+    unk_id = 1
+
+
+@pytest.mark.unit
+def test_salm_dataset_builds_aligned_sot_targets(monkeypatch):
+    text = "<spk:0> hello world <spk:1> yes now"
+    cut = dummy_cut(0, duration=0.04, recording=dummy_recording(0, duration=0.04, with_data=True))
+    cut.custom = {}
+    cut.supervisions = [
+        SupervisionSegment(id=cut.id, recording_id=cut.recording_id, start=0.0, duration=0.04, text=text)
+    ]
+    conversation = NeMoMultimodalConversation(
+        id="example-0",
+        turns=[
+            AudioTurn(role="user", cut=cut, audio_locator_tag="<|audio|>", text=text),
+            TextTurn(role="assistant", value=text),
+        ],
+        token_equivalent_duration=0.01,
+    )
+    conversation.input_ids = torch.tensor([7, 8, 9], dtype=torch.long)
+    conversation.mask = torch.tensor([False, True, True])
+    conversations = CutSet([conversation])
+
+    def fake_audio_collate(conversations, *args, **kwargs):
+        return torch.zeros(1, 640), torch.tensor([640], dtype=torch.long), conversations
+
+    def fake_speaker_activity_from_cut(cut, **kwargs):
+        assert cut.supervisions[0].text == text
+        return torch.tensor(
+            [
+                [0.0, 1.0],
+                [0.0, 1.0],
+                [1.0, 0.0],
+                [1.0, 0.0],
+            ]
+        )
+
+    monkeypatch.setattr(salm_dataset_module, "collate_conversation_audio_fault_tolerant", fake_audio_collate)
+    monkeypatch.setattr(salm_dataset_module, "speaker_activity_from_cut", fake_speaker_activity_from_cut)
+
+    dataset = salm_dataset_module.SALMDataset(
+        tokenizer=_Tokenizer(),
+        multispeaker_cfg={
+            "num_speakers": 2,
+            "sample_rate": 16000,
+            "window_stride": 0.01,
+            "subsampling_factor": 1,
+        },
+    )
+    assert dataset.multispeaker_cfg == salm_dataset_module.MultiSpeakerConfig(
+        num_speakers=2,
+        num_sample_per_mel_frame=160,
+        num_mel_frame_per_target_frame=1,
+    )
+    assert dataset.multispeaker_cfg.num_sample_per_mel_frame == 160
+
+    batch = dataset[conversations]
+
+    assert torch.equal(
+        batch["spk_targets"],
+        torch.tensor(
+            [
+                [
+                    [1.0, 0.0],
+                    [1.0, 0.0],
+                    [0.0, 1.0],
+                    [0.0, 1.0],
+                ]
+            ]
+        ),
+    )
+    assert torch.equal(batch["spk_target_length"], torch.tensor([4]))

tests/collections/speechlm2/test_salm_lora.py

@@ -0,0 +1,74 @@
+# Copyright (c) 2025, 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 types import SimpleNamespace
+
+import pytest
+from omegaconf import DictConfig
+from transformers import Qwen3Config, Qwen3ForCausalLM
+
+
+def _make_qwen3_stub():
+    llm = Qwen3ForCausalLM(
+        Qwen3Config(
+            vocab_size=32,
+            hidden_size=16,
+            intermediate_size=32,
+            num_hidden_layers=1,
+            num_attention_heads=2,
+            num_key_value_heads=2,
+            max_position_embeddings=32,
+        )
+    )
+    return SimpleNamespace(
+        cfg=DictConfig(
+            {
+                "prevent_freeze_params": [],
+                "lora": {
+                    "r": 4,
+                    "lora_alpha": 8,
+                    "lora_dropout": 0.0,
+                    "target_modules": ["q_proj", "v_proj"],
+                    "task_type": "CAUSAL_LM",
+                },
+            }
+        ),
+        llm=llm,
+        embed_tokens=llm.model.embed_tokens,
+    )
+
+
+def _import_peft_lora_dependencies():
+    """Import peft-dependent symbols, skipping when CI has an older torchao stack."""
+    try:
+        from peft import PeftModel
+        from peft.tuners.lora.torchao import dispatch_torchao  # noqa: F401
+        from nemo.collections.speechlm2.parts.lora import maybe_install_lora
+    except ImportError as e:
+        if "torchao" in str(e).lower():
+            pytest.skip(f"peft requires a newer torchao: {e}")
+        raise
+    return PeftModel, maybe_install_lora
+
+
+def test_maybe_install_lora_restores_qwen3_input_embeddings_temporarily():
+    PeftModel, maybe_install_lora = _import_peft_lora_dependencies()
+    model = _make_qwen3_stub()
+    del model.llm.model.embed_tokens
+
+    maybe_install_lora(model)
+
+    assert isinstance(model.llm, PeftModel)
+    assert hasattr(model.llm.base_model.model.model.layers[0].self_attn.q_proj, "lora_A")
+    assert model.cfg.prevent_freeze_params == [r"^.+\.lora_.+$"]
+    assert not hasattr(model.llm.base_model.model.model, "embed_tokens")

tests/collections/speechlm2/test_salm_packed_sequences.py

@@ -0,0 +1,552 @@
+# Copyright (c) 2025, 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.
+import ast
+from pathlib import Path
+
+import pytest
+import torch
+
+from nemo.collections.speechlm2.parts.packed_sequences import pack_audio_into_text_embeds, prepare_packed_llm_inputs
+
+PAD = 0
+AUDIO = 100
+REPO_ROOT = Path(__file__).parents[3]
+
+
+def test_packed_sequences_does_not_import_speechlm2_models_globally():
+    source = REPO_ROOT / "nemo/collections/speechlm2/parts/packed_sequences.py"
+    tree = ast.parse(source.read_text())
+    bad_imports = []
+    for node in tree.body:
+        if (
+            isinstance(node, ast.ImportFrom)
+            and node.module
+            and node.module.startswith("nemo.collections.speechlm2.models")
+        ):
+            bad_imports.append((node.lineno, node.module))
+        elif isinstance(node, ast.Import):
+            for alias in node.names:
+                if alias.name.startswith("nemo.collections.speechlm2.models"):
+                    bad_imports.append((node.lineno, alias.name))
+    assert bad_imports == []
+
+
+def _basic_batch():
+    """Mirrors `test_audio_placeholders.py::test_replace_placeholders`.
+
+    Two utterances, three audio replacements of length 4, 3, 2.
+    """
+    input_ids = torch.tensor(
+        [
+            [7, AUDIO, 1, 2, AUDIO, 1],
+            [PAD, PAD, 3, AUDIO, 4, 5],  # left-padded
+        ]
+    )
+    loss_mask = torch.tensor(
+        [
+            [False, False, False, False, False, True],
+            [False, False, False, False, True, True],
+        ]
+    )
+    embeds = torch.ones(2, 6, 2)
+    embeds[1, :2] = 0  # zero left-pad slots
+    replacements = [
+        torch.full((4, 2), fill_value=2.0),
+        torch.full((3, 2), fill_value=3.0),
+        torch.full((2, 2), fill_value=4.0),
+    ]
+    target_ids = input_ids.where(loss_mask, -100)
+    return input_ids, embeds, target_ids, replacements
+
+
+def test_basic_pack_shapes_and_cu_seqlens():
+    input_ids, embeds, target_ids, replacements = _basic_batch()
+    out = pack_audio_into_text_embeds(
+        input_ids=input_ids,
+        embeds=embeds,
+        target_ids=target_ids,
+        replacements=replacements,
+        padding_id=PAD,
+        placeholder_id=AUDIO,
+    )
+    # Real per-utterance flat lengths:
+    #   utt0: 1 text + 4 audio + 2 text + 3 audio + 1 text = 11
+    #   utt1: 1 text + 2 audio + 2 text                    = 5  (left-pad of 2 stripped)
+    assert out["seq_lens"].squeeze(-1).tolist() == [11, 5]
+    # cp_size=1, tp_size=1 ⇒ no rounding
+    assert out["seq_lens_padded"].squeeze(-1).tolist() == [11, 5]
+    # cu_seqlens = [0] + cumsum(seq_lens_padded)
+    assert out["cu_seqlens"].dtype == torch.int32
+    assert out["cu_seqlens"].tolist() == [0, 11, 16]
+    assert out["max_seqlen"].item() == 11
+    assert out["qkv_format"] == "thd"
+
+    T_total = 11 + 5
+    assert out["inputs_embeds"].shape == (T_total, 2)
+    assert out["labels"].shape == (T_total,)
+    assert out["position_ids"].shape == (T_total,)
+
+
+def test_position_ids_reset_per_utt():
+    input_ids, embeds, target_ids, replacements = _basic_batch()
+    out = pack_audio_into_text_embeds(
+        input_ids=input_ids,
+        embeds=embeds,
+        target_ids=target_ids,
+        replacements=replacements,
+        padding_id=PAD,
+        placeholder_id=AUDIO,
+    )
+    pos = out["position_ids"]
+    cu = out["cu_seqlens"].tolist()
+    for start, end in zip(cu[:-1], cu[1:]):
+        assert pos[start].item() == 0
+        assert torch.equal(pos[start:end], torch.arange(end - start, dtype=torch.long))
+
+
+def test_audio_frame_labels_are_ignored():
+    """Audio-frame slots must be -100 in `labels` regardless of loss_mask."""
+    input_ids, embeds, target_ids, replacements = _basic_batch()
+    out = pack_audio_into_text_embeds(
+        input_ids=input_ids,
+        embeds=embeds,
+        target_ids=target_ids,
+        replacements=replacements,
+        padding_id=PAD,
+        placeholder_id=AUDIO,
+    )
+    labels = out["labels"]
+    cu = out["cu_seqlens"].tolist()
+    # Utterance 0 layout: [t, a, a, a, a, t, t, a, a, a, t]
+    #                   pos 0  1  2  3  4  5  6  7  8  9  10
+    # Audio slots before shift: 1..4 and 7..9. After per-utt next-token shift,
+    # the *previous* slot's label becomes the audio target → also ignored once
+    # the original slot's label was -100. Verify all original audio slots map
+    # to -100 in the shifted output (audio at t means lab_shift[t-1] gets
+    # what was at t, which is -100 from the audio fill).
+    utt0 = labels[cu[0] : cu[1]]
+    # Shifted: original audio at positions 1-4 → label[0..3] should be -100;
+    # audio at 7-9 → label[6..8] should be -100; last slot (10) is -100.
+    assert (utt0[0:4] == -100).all()
+    assert (utt0[6:9] == -100).all()
+    assert utt0[-1].item() == -100
+
+
+def test_labels_shifted_per_utt():
+    """`labels[t]` should equal the original `target_ids` at position t+1
+    *within the utterance*, with the last slot set to -100."""
+    input_ids, embeds, target_ids, replacements = _basic_batch()
+    out = pack_audio_into_text_embeds(
+        input_ids=input_ids,
+        embeds=embeds,
+        target_ids=target_ids,
+        replacements=replacements,
+        padding_id=PAD,
+        placeholder_id=AUDIO,
+    )
+    labels = out["labels"]
+    cu = out["cu_seqlens"].tolist()
+    # Utt 0 last slot is the trailing text "1" with loss_mask=True (target=1),
+    # which after per-utt shift becomes the label of position 9 (the last
+    # audio frame). Since the shift puts orig[t+1] at slot t, the second-to-
+    # last slot of utt0 holds target_ids of the trailing "1".
+    utt0 = labels[cu[0] : cu[1]]
+    assert utt0[-2].item() == 1  # trailing text token "1" with loss_mask=True
+    assert utt0[-1].item() == -100  # last position of every utterance is -100
+
+    # Utt 1 last two text tokens (4, 5) had loss_mask=True. After per-utt
+    # shift, label[L-3] = 4, label[L-2] = 5, label[L-1] = -100.
+    utt1 = labels[cu[1] : cu[2]]
+    assert utt1[-3].item() == 4
+    assert utt1[-2].item() == 5
+    assert utt1[-1].item() == -100
+
+
+def test_no_audio_utterance():
+    """Utterance without any audio placeholders still packs correctly."""
+    input_ids = torch.tensor([[1, 2, 3, 4, 5]])
+    loss_mask = torch.tensor([[False, False, True, True, True]])
+    embeds = torch.full((1, 5, 2), 1.0)
+    target_ids = input_ids.where(loss_mask, -100)
+    out = pack_audio_into_text_embeds(
+        input_ids=input_ids,
+        embeds=embeds,
+        target_ids=target_ids,
+        replacements=[],
+        padding_id=PAD,
+        placeholder_id=AUDIO,
+    )
+    assert out["seq_lens"].squeeze(-1).tolist() == [5]
+    assert out["seq_lens_padded"].squeeze(-1).tolist() == [5]
+    assert out["cu_seqlens"].tolist() == [0, 5]
+    labels = out["labels"]
+    # Original target_ids = [-100, -100, 3, 4, 5]; after per-utt shift:
+    # [-100, 3, 4, 5, -100]
+    assert labels.tolist() == [-100, 3, 4, 5, -100]
+
+
+def test_b_one():
+    """Single-utterance batch produces valid `cu_seqlens=[0, L]`."""
+    input_ids = torch.tensor([[1, AUDIO, 2]])
+    loss_mask = torch.tensor([[False, False, True]])
+    embeds = torch.full((1, 3, 2), 1.0)
+    target_ids = input_ids.where(loss_mask, -100)
+    replacements = [torch.full((3, 2), 7.0)]
+    out = pack_audio_into_text_embeds(
+        input_ids=input_ids,
+        embeds=embeds,
+        target_ids=target_ids,
+        replacements=replacements,
+        padding_id=PAD,
+        placeholder_id=AUDIO,
+    )
+    assert out["seq_lens"].squeeze(-1).tolist() == [5]  # 1 + 3 audio + 1
+    assert out["cu_seqlens"].tolist() == [0, 5]
+    assert out["inputs_embeds"].shape == (5, 2)
+
+
+@pytest.mark.parametrize("cp_size", [2, 4])
+def test_cp_divisibility(cp_size):
+    """Each per-utt padded length is a multiple of 2*cp_size."""
+    input_ids, embeds, target_ids, replacements = _basic_batch()
+    out = pack_audio_into_text_embeds(
+        input_ids=input_ids,
+        embeds=embeds,
+        target_ids=target_ids,
+        replacements=replacements,
+        padding_id=PAD,
+        placeholder_id=AUDIO,
+        cp_size=cp_size,
+    )
+    mult = 2 * cp_size
+    for L in out["seq_lens_padded"].squeeze(-1).tolist():
+        assert L % mult == 0
+
+
+def test_tp_divisibility():
+    """`T_total` is a multiple of tp_size (last utterance gets the bump)."""
+    input_ids, embeds, target_ids, replacements = _basic_batch()
+    # real_lens = [11, 5], total = 16; pick tp_size=3 to force a bump.
+    out = pack_audio_into_text_embeds(
+        input_ids=input_ids,
+        embeds=embeds,
+        target_ids=target_ids,
+        replacements=replacements,
+        padding_id=PAD,
+        placeholder_id=AUDIO,
+        tp_size=3,
+    )
+    T_total = out["seq_lens_padded"].sum().item()
+    assert T_total % 3 == 0
+    # First utterance untouched (only the last gets the TP bump).
+    assert out["seq_lens_padded"].squeeze(-1).tolist()[0] == 11
+    # Last utterance bumped from 5 → 7 (next multiple of 3 after 16 is 18).
+    assert out["seq_lens_padded"].squeeze(-1).tolist()[1] == 7
+
+
+def test_tp_and_cp_combined():
+    input_ids, embeds, target_ids, replacements = _basic_batch()
+    out = pack_audio_into_text_embeds(
+        input_ids=input_ids,
+        embeds=embeds,
+        target_ids=target_ids,
+        replacements=replacements,
+        padding_id=PAD,
+        placeholder_id=AUDIO,
+        cp_size=2,
+        tp_size=8,
+    )
+    padded = out["seq_lens_padded"].squeeze(-1).tolist()
+    real = out["seq_lens"].squeeze(-1).tolist()
+    # Every padded length ≥ real and divisible by 4.
+    for r, p in zip(real, padded):
+        assert p >= r
+        assert p % 4 == 0
+    # Total divisible by tp_size.
+    assert sum(padded) % 8 == 0
+
+
+def test_tp_bump_preserves_cp_alignment_when_tp_is_not_cp_multiple():
+    input_ids = torch.tensor([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]])
+    loss_mask = torch.ones_like(input_ids, dtype=torch.bool)
+    embeds = torch.ones(2, 5, 2)
+    target_ids = input_ids.where(loss_mask, -100)
+    out = pack_audio_into_text_embeds(
+        input_ids=input_ids,
+        embeds=embeds,
+        target_ids=target_ids,
+        replacements=[],
+        padding_id=PAD,
+        placeholder_id=AUDIO,
+        cp_size=2,
+        tp_size=6,
+    )
+    padded = out["seq_lens_padded"].squeeze(-1).tolist()
+    assert padded == [8, 16]
+    for L in padded:
+        assert L % 4 == 0
+    assert sum(padded) % 6 == 0
+
+
+def test_cu_seqlens_matches_padded_cumsum():
+    input_ids, embeds, target_ids, replacements = _basic_batch()
+    out = pack_audio_into_text_embeds(
+        input_ids=input_ids,
+        embeds=embeds,
+        target_ids=target_ids,
+        replacements=replacements,
+        padding_id=PAD,
+        placeholder_id=AUDIO,
+        cp_size=2,
+        tp_size=8,
+    )
+    expected = [0]
+    for L in out["seq_lens_padded"].squeeze(-1).tolist():
+        expected.append(expected[-1] + L)
+    assert out["cu_seqlens"].tolist() == expected
+    assert out["max_seqlen"].item() == max(out["seq_lens_padded"].squeeze(-1).tolist())
+
+
+def test_loss_mask_propagates_to_minus_100():
+    """Positions where loss_mask=False end up as -100 in the shifted labels."""
+    input_ids = torch.tensor([[1, 2, 3, 4]])
+    loss_mask = torch.tensor([[False, False, False, False]])  # nothing supervised
+    embeds = torch.full((1, 4, 2), 1.0)
+    target_ids = input_ids.where(loss_mask, -100)
+    out = pack_audio_into_text_embeds(
+        input_ids=input_ids,
+        embeds=embeds,
+        target_ids=target_ids,
+        replacements=[],
+        padding_id=PAD,
+        placeholder_id=AUDIO,
+    )
+    assert (out["labels"] == -100).all()
+
+
+def test_prepare_packed_llm_inputs_attention_kwargs_reach_te_preprocessor():
+    """End-to-end contract: ``prepare_packed_llm_inputs`` → Automodel's TE
+    attention preprocessor must yield ``max_seqlen_q``/``max_seqlen_kv`` and
+    ``cu_seqlens_q``/``cu_seqlens_kv`` populated for the THD path.
+
+    Regression for: ``prepare_packed_llm_inputs`` previously emitted
+    ``max_seqlen_q``/``max_seqlen_kv`` (plural), but the preprocessor only
+    inspects the singular ``max_seqlen`` key in the ``cu_seqlens`` branch
+    (``Automodel/.../attention/utils.py``). The plural keys were silently
+    dropped, TE fell back to a degenerate varlen path, and step-1 backward
+    produced NaN gradients that poisoned every subsequent step.
+    """
+    automodel_attn = pytest.importorskip(
+        "nemo_automodel.components.attention.utils",
+        reason="Automodel attention preprocessor required for the contract test.",
+    )
+    input_ids, embeds, target_ids, replacements = _basic_batch()
+    out = prepare_packed_llm_inputs(
+        input_ids=input_ids,
+        text_embs=embeds,
+        audio_embs=replacements,
+        target_ids=target_ids,
+        padding_id=PAD,
+        placeholder_id=AUDIO,
+        device_mesh=None,  # CP=1, TP=1 path
+    )
+    llm_kwargs = out["llm_kwargs"]
+    assert out["attention_mask"] is None
+    assert llm_kwargs["qkv_format"] == "thd"
+    assert "cu_seqlens" in llm_kwargs
+    assert "max_seqlen" in llm_kwargs, (
+        "Automodel's preprocessor only checks the singular `max_seqlen` key in the "
+        "cu_seqlens THD branch; pre-split `max_seqlen_q`/`max_seqlen_kv` would be dropped."
+    )
+    assert "cu_seqlens_padded" not in llm_kwargs, (
+        "Standard Automodel pipeline emits only `cu_seqlens`, never both `cu_seqlens` "
+        "and `cu_seqlens_padded`. Passing both activates the `pad_between_seqs=True` "
+        "branch in Automodel/.../attention/utils.py, routing TE down a different path."
+    )
+
+    # Run the LLM kwargs through the preprocessor exactly as the attention
+    # layer does: 4D BSHD-shaped Q/K/V plus attention_mask=None plus the
+    # llm_kwargs splatted in. ``input_embeds`` is now 2D ``[T, H]`` per the
+    # canonical Automodel THD shape contract.
+    B, T, H = 1, int(out["input_embeds"].shape[0]), 2
+    nh, hd = 2, 4
+    q = torch.zeros(B, T, nh, hd)
+    k = torch.zeros(B, T, nh, hd)
+    v = torch.zeros(B, T, nh, hd)
+    _, _, _, te_attn_kwargs = automodel_attn.preprocess_args_and_kwargs_for_attn(
+        q, k, v, attention_mask=None, attn_impl="te", **llm_kwargs
+    )
+    assert te_attn_kwargs.get("qkv_format") == "thd"
+    assert te_attn_kwargs.get("attn_mask_type") == "padding_causal"
+    assert "cu_seqlens_q" in te_attn_kwargs and "cu_seqlens_kv" in te_attn_kwargs
+    assert "max_seqlen_q" in te_attn_kwargs and "max_seqlen_kv" in te_attn_kwargs, (
+        "TE DotProductAttention requires max_seqlen_q/kv for qkv_format='thd'; "
+        "missing keys cause silent degenerate-path fallback and NaN gradients."
+    )
+    assert te_attn_kwargs["max_seqlen_q"] == llm_kwargs["max_seqlen"]
+    assert te_attn_kwargs["max_seqlen_kv"] == llm_kwargs["max_seqlen"]
+
+
+def _bshd_supervised_pairs(input_ids, embeds, target_ids, replacements):
+    """Run the BSHD path (``replace_placeholders_and_build_targets`` + the
+    ``[:-1] / [1:]`` next-token shift used in
+    ``SALMAutomodel.prepare_inputs``) and return the ordered list of
+    supervised ``(input_embedding, target_token_id)`` pairs.
+    """
+    from nemo.collections.speechlm2.models.salm import replace_placeholders_and_build_targets
+
+    bshd_embs, bshd_targets, _ = replace_placeholders_and_build_targets(
+        input_ids=input_ids,
+        embeds=embeds,
+        padding_id=PAD,
+        placeholder_id=AUDIO,
+        replacements=[r.clone() for r in replacements],
+        target_ids=target_ids,
+    )
+    bshd_embs = bshd_embs[:, :-1]
+    bshd_targets = bshd_targets[:, 1:]
+
+    pairs = []
+    B, T = bshd_targets.shape
+    for b in range(B):
+        for t in range(T):
+            tgt = bshd_targets[b, t].item()
+            if tgt != -100:
+                pairs.append((bshd_embs[b, t].clone(), tgt))
+    return pairs
+
+
+def _thd_supervised_pairs(input_ids, embeds, target_ids, replacements):
+    """Run the THD path (``pack_audio_into_text_embeds`` with the per-utt
+    next-token shift) and return the ordered list of supervised
+    ``(input_embedding, target_token_id)`` pairs.
+    """
+    out = pack_audio_into_text_embeds(
+        input_ids=input_ids,
+        embeds=embeds,
+        target_ids=target_ids,
+        replacements=[r.clone() for r in replacements],
+        padding_id=PAD,
+        placeholder_id=AUDIO,
+    )
+    embs = out["inputs_embeds"]  # [T_total, H]
+    labels = out["labels"]  # [T_total]
+    pairs = []
+    for t in range(labels.shape[0]):
+        tgt = labels[t].item()
+        if tgt != -100:
+            pairs.append((embs[t].clone(), tgt))
+    return pairs
+
+
+def _assert_pairs_equivalent(bshd_pairs, thd_pairs, *, atol=1e-6):
+    assert len(bshd_pairs) == len(thd_pairs), (
+        f"BSHD has {len(bshd_pairs)} supervised pairs, THD has {len(thd_pairs)}. "
+        f"Both paths must yield the same ordered set of (input, target) pairs."
+    )
+    for i, ((e_b, t_b), (e_t, t_t)) in enumerate(zip(bshd_pairs, thd_pairs)):
+        assert t_b == t_t, (
+            f"Pair {i}: target_id mismatch — BSHD={t_b}, THD={t_t}. "
+            f"Per-utt next-token shift must align with global [:-1]/[1:] shift."
+        )
+        assert torch.allclose(e_b, e_t, atol=atol), (
+            f"Pair {i} (target={t_b}): input embedding mismatch between BSHD and THD. " f"BSHD={e_b}, THD={e_t}"
+        )
+
+
+def test_thd_and_bshd_supervised_pairs_match_basic():
+    """First-principles invariant: BSHD and THD are different *layouts* of the
+    same data. The set of supervised ``(input_embedding, target_token_id)``
+    pairs that contribute to the cross-entropy loss must be identical between
+    paths. Any divergence in this set means the THD path is feeding the model
+    something the BSHD path is not (or vice-versa).
+    """
+    input_ids, embeds, target_ids, replacements = _basic_batch()
+    bshd_pairs = _bshd_supervised_pairs(input_ids, embeds, target_ids, replacements)
+    thd_pairs = _thd_supervised_pairs(input_ids, embeds, target_ids, replacements)
+    _assert_pairs_equivalent(bshd_pairs, thd_pairs)
+
+
+def test_thd_and_bshd_supervised_pairs_match_no_audio_utt():
+    """Pure-text utterance (no audio_locator)."""
+    input_ids = torch.tensor([[1, 2, 3, 4, 5]])
+    loss_mask = torch.tensor([[False, False, True, True, True]])
+    embeds = torch.randn(1, 5, 4)
+    target_ids = input_ids.where(loss_mask, -100)
+    bshd_pairs = _bshd_supervised_pairs(input_ids, embeds, target_ids, replacements=[])
+    thd_pairs = _thd_supervised_pairs(input_ids, embeds, target_ids, replacements=[])
+    _assert_pairs_equivalent(bshd_pairs, thd_pairs)
+
+
+def test_thd_and_bshd_supervised_pairs_match_left_padded():
+    """Left-padded utterances must yield the same supervised pairs."""
+    input_ids = torch.tensor(
+        [
+            [PAD, PAD, PAD, 1, 2, AUDIO, 3, 4],
+            [PAD, 5, 6, AUDIO, 7, AUDIO, 8, 9],
+        ]
+    )
+    loss_mask = torch.tensor(
+        [
+            [False, False, False, False, False, True, True, True],
+            [False, False, False, True, True, True, True, True],
+        ]
+    )
+    embeds = torch.randn(2, 8, 4)
+    embeds[0, :3] = 0  # zero left-pad slots
+    embeds[1, :1] = 0
+    target_ids = input_ids.where(loss_mask, -100)
+    replacements = [
+        torch.randn(3, 4),  # utt0 audio
+        torch.randn(2, 4),  # utt1 first audio
+        torch.randn(4, 4),  # utt1 second audio
+    ]
+    bshd_pairs = _bshd_supervised_pairs(input_ids, embeds, target_ids, replacements)
+    thd_pairs = _thd_supervised_pairs(input_ids, embeds, target_ids, replacements)
+    _assert_pairs_equivalent(bshd_pairs, thd_pairs)
+
+
+def test_thd_and_bshd_supervised_pairs_match_b1():
+    """Single-utterance batch."""
+    input_ids = torch.tensor([[1, AUDIO, 2, 3, AUDIO, 4]])
+    loss_mask = torch.tensor([[False, False, True, True, True, True]])
+    embeds = torch.randn(1, 6, 4)
+    target_ids = input_ids.where(loss_mask, -100)
+    replacements = [torch.randn(2, 4), torch.randn(5, 4)]
+    bshd_pairs = _bshd_supervised_pairs(input_ids, embeds, target_ids, replacements)
+    thd_pairs = _thd_supervised_pairs(input_ids, embeds, target_ids, replacements)
+    _assert_pairs_equivalent(bshd_pairs, thd_pairs)
+
+
+def test_padded_slots_have_zero_embed_and_ignored_label():
+    """Inter-utt padding (added for cp_size rounding) gets zero embedding,
+    -100 label, and contiguous position_ids."""
+    input_ids, embeds, target_ids, replacements = _basic_batch()
+    out = pack_audio_into_text_embeds(
+        input_ids=input_ids,
+        embeds=embeds,
+        target_ids=target_ids,
+        replacements=replacements,
+        padding_id=PAD,
+        placeholder_id=AUDIO,
+        cp_size=2,  # rounds 11→12 and 6→6
+    )
+    embs = out["inputs_embeds"]
+    labels = out["labels"]
+    pos = out["position_ids"]
+    # Utt 0 had real_len=11 and padded to 12 (next multiple of 4). Slot 11 is
+    # the pad slot.
+    assert torch.equal(embs[11], torch.zeros(2))
+    assert labels[11].item() == -100
+    assert pos[11].item() == 11  # contiguous with the utt's real positions

tests/collections/speechlm2/test_salm_parallelism_validation.py

@@ -0,0 +1,154 @@
+# Copyright (c) 2025, 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.
+"""Unit tests for ``validate_parallelism_compatibility``.
+
+Pure-function tests — no Lightning, no model, no device mesh required.
+"""
+import warnings
+
+import pytest
+
+from nemo.collections.speechlm2.parts.parallel import validate_parallelism_compatibility
+
+
+# Combinations that must pass without raising or warning.
+
+
+def test_bshd_cp1_te_passes():
+    validate_parallelism_compatibility(
+        packed_sequences=False,
+        cp_size=1,
+        attn_backend="te",
+        nvte_fused_attn=None,
+        device_capability=(9, 0),  # H100
+    )
+
+
+def test_bshd_cp1_sdpa_passes():
+    validate_parallelism_compatibility(
+        packed_sequences=False,
+        cp_size=1,
+        attn_backend="sdpa",
+        nvte_fused_attn=None,
+        device_capability=(9, 0),
+    )
+
+
+def test_thd_cp1_te_with_fused_attn_off_passes():
+    validate_parallelism_compatibility(
+        packed_sequences=True,
+        cp_size=1,
+        attn_backend="te",
+        nvte_fused_attn="0",
+        device_capability=(12, 0),  # sm_120 — still fine because env is set
+    )
+
+
+def test_thd_cp2_te_with_fused_attn_off_passes():
+    validate_parallelism_compatibility(
+        packed_sequences=True,
+        cp_size=2,
+        attn_backend="te",
+        nvte_fused_attn="0",
+        device_capability=(12, 0),
+    )
+
+
+# BSHD + CP > 1 — hard error regardless of other knobs.
+
+
+@pytest.mark.parametrize("cp_size", [2, 4, 8])
+def test_bshd_with_cp_raises(cp_size):
+    with pytest.raises(ValueError, match="BSHD .* incompatible with cp_size > 1"):
+        validate_parallelism_compatibility(
+            packed_sequences=False,
+            cp_size=cp_size,
+            attn_backend="te",
+            nvte_fused_attn="0",
+            device_capability=(9, 0),
+        )
+
+
+# THD + non-TE attention — hard error.
+
+
+@pytest.mark.parametrize("attn_backend", ["sdpa", "flex"])
+def test_thd_with_non_te_attn_raises(attn_backend):
+    with pytest.raises(ValueError, match=r"THD.*requires.*attn=te"):
+        validate_parallelism_compatibility(
+            packed_sequences=True,
+            cp_size=1,
+            attn_backend=attn_backend,
+            nvte_fused_attn="0",
+            device_capability=(9, 0),
+        )
+
+
+# THD + TE + NVTE_FUSED_ATTN unset — warns on non-sm_120, raises on sm_120.
+
+
+@pytest.mark.parametrize("nvte_fused_attn", [None, "", "1", "true"])
+def test_thd_te_without_fused_attn_off_warns_on_other_archs(nvte_fused_attn):
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        validate_parallelism_compatibility(
+            packed_sequences=True,
+            cp_size=1,
+            attn_backend="te",
+            nvte_fused_attn=nvte_fused_attn,
+            device_capability=(9, 0),  # H100
+        )
+    assert len(caught) == 1
+    assert "NVTE_FUSED_ATTN" in str(caught[0].message)
+
+
+@pytest.mark.parametrize("nvte_fused_attn", [None, "", "1", "true"])
+def test_thd_te_without_fused_attn_off_raises_on_sm120(nvte_fused_attn):
+    with pytest.raises(ValueError, match="NVTE_FUSED_ATTN"):
+        validate_parallelism_compatibility(
+            packed_sequences=True,
+            cp_size=1,
+            attn_backend="te",
+            nvte_fused_attn=nvte_fused_attn,
+            device_capability=(12, 0),  # sm_120
+        )
+
+
+def test_thd_te_with_fused_attn_off_does_not_warn_on_sm120():
+    """The escape-hatch case: user has the env set, no warning fires."""
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        validate_parallelism_compatibility(
+            packed_sequences=True,
+            cp_size=1,
+            attn_backend="te",
+            nvte_fused_attn="0",
+            device_capability=(12, 0),
+        )
+    assert len(caught) == 0
+
+
+def test_unknown_device_capability_warns_not_raises():
+    """``device_capability=None`` (CPU-only env) treats the THD/TE check as a warning."""
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        validate_parallelism_compatibility(
+            packed_sequences=True,
+            cp_size=1,
+            attn_backend="te",
+            nvte_fused_attn=None,
+            device_capability=None,
+        )
+    assert len(caught) == 1
+    assert "NVTE_FUSED_ATTN" in str(caught[0].message)

tests/collections/speechlm2/test_to_hf.py

@@ -0,0 +1,260 @@
+# Copyright (c) 2025, 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.
+"""Unit tests for ``examples/speechlm2/to_hf.py::prepare_for_vllm``.
+
+The script lives under ``examples/`` (not an importable package), so we load
+it via ``importlib`` and patch ``AutoTokenizer`` / ``_detect_vllm_architecture``
+to avoid any network or real-model dependencies.
+"""
+import importlib.util
+import json
+from pathlib import Path
+from unittest.mock import patch
+
+import pytest
+import torch
+from safetensors.torch import load_file
+
+_TO_HF_PATH = Path(__file__).parents[3] / "examples" / "speechlm2" / "to_hf.py"
+_spec = importlib.util.spec_from_file_location("to_hf_for_test", _TO_HF_PATH)
+to_hf = importlib.util.module_from_spec(_spec)
+_spec.loader.exec_module(to_hf)
+
+AUDIO_TOKEN = "<|audio|>"
+CHAT_TEMPLATE_INLINE = "{% for msg in messages %}{{msg.content}}{% endfor %}"
+CHAT_TEMPLATE_LARGE = "{% for msg in messages %}" + "X" * 4096 + "{{msg.content}}{% endfor %}"
+
+
+class _FakeTokenizer:
+    """Minimal stand-in for an HF ``AutoTokenizer`` instance.
+
+    Mimics only the surface used by ``prepare_for_vllm``:
+      * ``get_vocab`` / ``add_special_tokens``
+      * ``save_pretrained`` (writes tokenizer_config.json, optionally
+        splitting a large chat_template into a separate .jinja file)
+      * ``eos_token_id``
+    """
+
+    def __init__(
+        self,
+        *,
+        vocab_tokens=(),
+        chat_template=CHAT_TEMPLATE_INLINE,
+        split_chat_template=False,
+        tokenizer_class="Qwen2Tokenizer",
+        eos_token_id=42,
+    ):
+        self._vocab = {tok: i for i, tok in enumerate(vocab_tokens)}
+        self._chat_template = chat_template
+        self._split_chat_template = split_chat_template
+        self._tokenizer_class = tokenizer_class
+        self.eos_token_id = eos_token_id
+        self.add_special_tokens_calls = []
+
+    def get_vocab(self):
+        return dict(self._vocab)
+
+    def add_special_tokens(self, special_tokens_dict):
+        added = special_tokens_dict.get("additional_special_tokens", [])
+        for tok in added:
+            if tok not in self._vocab:
+                self._vocab[tok] = len(self._vocab)
+        self.add_special_tokens_calls.append(added)
+        return len(added)
+
+    def save_pretrained(self, output_dir):
+        output_dir = Path(output_dir)
+        # Transformers writes extra_special_tokens as a LIST (not dict); mimic
+        # that here so the dict-form coercion in to_hf.py is exercised.
+        tok_cfg = {
+            "tokenizer_class": self._tokenizer_class,
+            "extra_special_tokens": [AUDIO_TOKEN],
+        }
+        if self._split_chat_template:
+            (output_dir / "chat_template.jinja").write_text(self._chat_template)
+        else:
+            tok_cfg["chat_template"] = self._chat_template
+        (output_dir / "tokenizer_config.json").write_text(json.dumps(tok_cfg))
+        (output_dir / "tokenizer.json").write_text('{"fake": true}')
+
+
+def _seed_output_dir(tmp_path, llm_arch="Qwen2ForCausalLM"):
+    """Pre-populate ``output_dir`` with what ``save_hf_checkpoint`` would write."""
+    (tmp_path / "config.json").write_text(
+        json.dumps(
+            {
+                "architectures": [llm_arch],
+                "hidden_size": 2048,
+                "num_hidden_layers": 24,
+            }
+        )
+    )
+    return tmp_path
+
+
+class _FakeLLMConfig:
+    def save_pretrained(self, output_dir):
+        output_dir = Path(output_dir)
+        output_dir.mkdir(parents=True, exist_ok=True)
+        (output_dir / "config.json").write_text(
+            json.dumps(
+                {
+                    "model_type": "qwen2",
+                    "architectures": ["Qwen2ForCausalLM"],
+                    "hidden_size": 2048,
+                }
+            )
+        )
+
+
+class _FakeExportModel:
+    cfg = {
+        "pretrained_llm": "fake-model",
+        "pretrained_asr": "fake-asr",
+        "pretrained_weights": False,
+        "dtype": "bf16",
+        "torch_dtype": "bf16",
+        "audio_locator_tag": AUDIO_TOKEN,
+    }
+    llm = type("_FakeLLM", (), {"config": _FakeLLMConfig()})()
+
+
+def test_save_hf_checkpoint_writes_llm_backbone_config(tmp_path):
+    cfg = to_hf.HfExportConfig(
+        class_path="fake.Class",
+        ckpt_path="fake.ckpt",
+        ckpt_config="fake.yaml",
+        output_dir=str(tmp_path),
+        dtype="bfloat16",
+    )
+    to_hf.save_hf_checkpoint(_FakeExportModel(), {"weight": torch.zeros(1)}, cfg)
+
+    root_cfg = json.loads((tmp_path / "config.json").read_text())
+    llm_cfg = json.loads((tmp_path / "llm_backbone" / "config.json").read_text())
+
+    assert "llm_config" not in root_cfg
+    assert root_cfg["pretrained_llm"] == "fake-model"
+    assert root_cfg["dtype"] == "bfloat16"
+    assert root_cfg["torch_dtype"] == "bfloat16"
+    assert llm_cfg["model_type"] == "qwen2"
+    assert llm_cfg["architectures"] == ["Qwen2ForCausalLM"]
+    assert _FakeExportModel.cfg["dtype"] == "bf16"
+
+
+def test_save_hf_checkpoint_accepts_bf16_export_dtype(tmp_path):
+    cfg = to_hf.HfExportConfig(
+        class_path="fake.Class",
+        ckpt_path="fake.ckpt",
+        ckpt_config="fake.yaml",
+        output_dir=str(tmp_path),
+        dtype="bf16",
+    )
+    to_hf.save_hf_checkpoint(_FakeExportModel(), {"weight": torch.zeros(1)}, cfg)
+
+    root_cfg = json.loads((tmp_path / "config.json").read_text())
+    state_dict = load_file(tmp_path / "model.safetensors")
+
+    assert root_cfg["dtype"] == "bfloat16"
+    assert root_cfg["torch_dtype"] == "bfloat16"
+    assert state_dict["weight"].dtype == torch.bfloat16
+
+
+# ──────────────────────────────────────────────────────────────────────
+# Error paths (no mocking required — checks run before any HF calls)
+# ──────────────────────────────────────────────────────────────────────
+
+
+def test_prepare_for_vllm_missing_pretrained_llm(tmp_path):
+    with pytest.raises(ValueError, match="pretrained_llm"):
+        to_hf.prepare_for_vllm(str(tmp_path), {"audio_locator_tag": AUDIO_TOKEN})
+
+
+def test_prepare_for_vllm_missing_audio_locator_tag(tmp_path):
+    with pytest.raises(ValueError, match="audio_locator_tag"):
+        to_hf.prepare_for_vllm(str(tmp_path), {"pretrained_llm": "fake-model"})
+
+
+# ──────────────────────────────────────────────────────────────────────
+# Happy paths (mock AutoTokenizer + _detect_vllm_architecture)
+# ──────────────────────────────────────────────────────────────────────
+
+
+def _run_prepare(tmp_path, fake_tok, arch="NeMoSpeechLMForConditionalGeneration", llm_arch="Qwen2ForCausalLM"):
+    output_dir = _seed_output_dir(tmp_path, llm_arch=llm_arch)
+    with (
+        patch.object(to_hf, "_detect_vllm_architecture", return_value=arch),
+        patch("transformers.AutoTokenizer.from_pretrained", return_value=fake_tok),
+    ):
+        to_hf.prepare_for_vllm(
+            str(output_dir),
+            {"pretrained_llm": "fake-model", "audio_locator_tag": AUDIO_TOKEN},
+        )
+    return output_dir
+
+
+def test_prepare_for_vllm_patches_config_json(tmp_path):
+    """config.json gets model_type, architectures, and audio_locator_tag."""
+    output_dir = _run_prepare(tmp_path, _FakeTokenizer())
+    cfg = json.loads((output_dir / "config.json").read_text())
+    assert cfg["model_type"] == "nemo_speechlm"
+    assert cfg["architectures"] == ["NeMoSpeechLMForConditionalGeneration"]
+    assert cfg["audio_locator_tag"] == AUDIO_TOKEN
+    # Original LLM fields are preserved.
+    assert cfg["hidden_size"] == 2048
+
+
+def test_prepare_for_vllm_adds_audio_token_to_vocab(tmp_path):
+    """Audio token is registered via add_special_tokens when not already in vocab."""
+    fake_tok = _FakeTokenizer(vocab_tokens=["<|im_start|>", "<|im_end|>"])
+    _run_prepare(tmp_path, fake_tok)
+    assert fake_tok.add_special_tokens_calls == [[AUDIO_TOKEN]]
+    assert AUDIO_TOKEN in fake_tok.get_vocab()
+
+
+def test_prepare_for_vllm_skips_add_if_audio_token_already_in_vocab(tmp_path):
+    """Avoid re-adding a token that was already present in the backbone vocab."""
+    fake_tok = _FakeTokenizer(vocab_tokens=["<|im_start|>", "<|im_end|>", AUDIO_TOKEN])
+    _run_prepare(tmp_path, fake_tok)
+    assert fake_tok.add_special_tokens_calls == []
+
+
+def test_prepare_for_vllm_tokenizer_config_normalized(tmp_path):
+    """tokenizer_config.json has dict-form extra_special_tokens + forced tokenizer_class."""
+    output_dir = _run_prepare(tmp_path, _FakeTokenizer(tokenizer_class="TokenizersBackend"))
+    tok_cfg = json.loads((output_dir / "tokenizer_config.json").read_text())
+    assert tok_cfg["tokenizer_class"] == "PreTrainedTokenizerFast"
+    assert tok_cfg["extra_special_tokens"] == {"audio_token": AUDIO_TOKEN}
+
+
+def test_prepare_for_vllm_preserves_inline_chat_template_verbatim(tmp_path):
+    """No enable_thinking patching: chat_template is byte-identical after prep."""
+    output_dir = _run_prepare(tmp_path, _FakeTokenizer(chat_template=CHAT_TEMPLATE_INLINE))
+    tok_cfg = json.loads((output_dir / "tokenizer_config.json").read_text())
+    assert tok_cfg["chat_template"] == CHAT_TEMPLATE_INLINE
+
+
+def test_prepare_for_vllm_rescues_chat_template_jinja_file(tmp_path):
+    """Qwen3-style: chat_template split to .jinja file → inlined + file removed."""
+    fake_tok = _FakeTokenizer(chat_template=CHAT_TEMPLATE_LARGE, split_chat_template=True)
+    output_dir = _run_prepare(tmp_path, fake_tok)
+    tok_cfg = json.loads((output_dir / "tokenizer_config.json").read_text())
+    assert tok_cfg["chat_template"] == CHAT_TEMPLATE_LARGE
+    assert not (output_dir / "chat_template.jinja").exists()
+
+
+def test_prepare_for_vllm_generation_config(tmp_path):
+    """generation_config.json gets written with the tokenizer's eos_token_id."""
+    output_dir = _run_prepare(tmp_path, _FakeTokenizer(eos_token_id=99))
+    gen_cfg = json.loads((output_dir / "generation_config.json").read_text())
+    assert gen_cfg == {"eos_token_id": [99]}

tests/collections/speechlm2/test_vllm_audio_token_estimator.py

@@ -0,0 +1,169 @@
+# 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.
+
+"""Drift check for the vLLM plugin's audio-token estimator.
+
+The plugin hand-rolls
+``NeMoSpeechLMProcessingInfo._estimate_audio_tokens`` in pure Python for
+speed (~90x faster than the equivalent tensor-ops path). Accuracy
+matters: the result drives how many ``<|audio|>`` placeholders the
+prompt processor inserts AND the encoder's real output frame count at
+forward time; a mismatch breaks placeholder-to-feature alignment.
+
+This test locks the hand-rolled math to NeMo's ``calc_length`` on a
+canonical set of audio lengths. If FastConformer's subsampling ever
+changes upstream (different kernel/stride/repeat), this test fails and
+forces an update to the estimator.
+"""
+
+import pytest
+import torch
+
+pytest.importorskip("vllm")
+
+from nemo.collections.asr.parts.submodules.subsampling import calc_length
+from nemo.collections.speechlm2.vllm.salm.audio import (
+    _DUMMY_AUDIO_MAX_DURATION_S,
+    _MIN_CHUNK_SIZE_SAMPLES,
+    _SAMPLING_RATE,
+    NeMoSpeechLMProcessingInfo,
+)
+
+
+def _reference(audio_length_samples: int) -> int:
+    """Reference impl using NeMo's calc_length for the conv chain."""
+    n_fft = 512
+    hop_length = 160
+    stft_pad = n_fft // 2
+    fbank_len = (audio_length_samples + 2 * stft_pad - n_fft) // hop_length
+    out = calc_length(
+        torch.tensor([fbank_len], dtype=torch.float),
+        all_paddings=2,
+        kernel_size=3,
+        stride=2,
+        ceil_mode=False,
+        repeat_num=3,
+    )
+    return max(1, int(out.item()))
+
+
+@pytest.mark.parametrize(
+    "samples",
+    [
+        1_600,  # 0.1 s
+        16_000,  # 1 s
+        80_000,  # 5 s
+        160_000,  # 10 s
+        320_000,  # 20 s
+        640_000,  # 40 s, the typical max
+        12_345,  # arbitrary small
+        54_321,  # arbitrary mid
+        100_001,  # arbitrary (odd)
+    ],
+)
+def test_estimator_matches_calc_length(samples: int) -> None:
+    ours = NeMoSpeechLMProcessingInfo._estimate_audio_tokens(samples)
+    ref = _reference(samples)
+    assert ours == ref, (
+        f"audio_token estimator diverged from NeMo calc_length for "
+        f"samples={samples}: ours={ours}, ref={ref}. "
+        f"Check if FastConformer's subsampling stack changed upstream."
+    )
+
+
+def test_estimator_min_one() -> None:
+    """Even for very short audio the estimator must return at least 1."""
+    assert NeMoSpeechLMProcessingInfo._estimate_audio_tokens(1) >= 1
+
+
+def test_estimator_chunking_disabled_matches_single_pass() -> None:
+    """``chunk_size_seconds=None`` must match the legacy single-pass estimate."""
+    samples = 30 * 16_000
+    assert NeMoSpeechLMProcessingInfo._estimate_audio_tokens(
+        samples, chunk_size_seconds=None
+    ) == NeMoSpeechLMProcessingInfo._estimate_audio_tokens_single_pass(samples)
+
+
+def test_estimator_short_audio_falls_back_to_single_pass() -> None:
+    """Audio shorter than the chunk size collapses to a single forward."""
+    samples = 5 * 16_000
+    assert NeMoSpeechLMProcessingInfo._estimate_audio_tokens(
+        samples, chunk_size_seconds=30.0
+    ) == NeMoSpeechLMProcessingInfo._estimate_audio_tokens_single_pass(samples)
+
+
+def test_estimator_chunked_sums_per_chunk_frames() -> None:
+    """Long audio is split into chunks and per-chunk frame counts are summed,
+    matching ``encode_audio_with_optional_chunking``'s concat behavior."""
+    samples = 90 * 16_000
+    chunk_size_seconds = 30.0
+    chunk_samples = int(round(chunk_size_seconds * 16_000))
+    expected = sum(
+        NeMoSpeechLMProcessingInfo._estimate_audio_tokens_single_pass(min(chunk_samples, samples - i))
+        for i in range(0, samples, chunk_samples)
+    )
+    assert (
+        NeMoSpeechLMProcessingInfo._estimate_audio_tokens(samples, chunk_size_seconds=chunk_size_seconds) == expected
+    )
+
+
+def test_estimator_chunked_tail_folded_into_previous_chunk() -> None:
+    """A tiny tail (< min chunk size) is folded into the previous chunk so
+    the total token count matches the runtime helper instead of producing a
+    spurious single-frame chunk that the audio preprocessor would reject."""
+    chunk_size_seconds = 30.0
+    chunk_samples = int(round(chunk_size_seconds * 16_000))
+    samples = chunk_samples + 100  # 100 sample tail < min_chunk_size_samples (320)
+    # Folded: one chunk of `samples` samples (no split).
+    expected = NeMoSpeechLMProcessingInfo._estimate_audio_tokens_single_pass(samples)
+    assert (
+        NeMoSpeechLMProcessingInfo._estimate_audio_tokens(samples, chunk_size_seconds=chunk_size_seconds) == expected
+    )
+
+
+def test_estimator_clamps_tiny_chunk_size_to_min_samples() -> None:
+    assert _MIN_CHUNK_SIZE_SAMPLES == 320
+
+    chunk_size_seconds = 1 / _SAMPLING_RATE
+    samples = 2 * _MIN_CHUNK_SIZE_SAMPLES + 100
+    expected = NeMoSpeechLMProcessingInfo._estimate_audio_tokens_single_pass(
+        _MIN_CHUNK_SIZE_SAMPLES
+    ) + NeMoSpeechLMProcessingInfo._estimate_audio_tokens_single_pass(_MIN_CHUNK_SIZE_SAMPLES + 100)
+
+    assert (
+        NeMoSpeechLMProcessingInfo._estimate_audio_tokens(samples, chunk_size_seconds=chunk_size_seconds) == expected
+    )
+
+
+def test_estimator_negative_chunk_size_raises() -> None:
+    with pytest.raises(ValueError, match="encoder_chunk_size_seconds"):
+        NeMoSpeechLMProcessingInfo._estimate_audio_tokens(16_000, chunk_size_seconds=-1.0)
+
+
+@pytest.mark.parametrize("chunk_size_seconds", [None, 30.0])
+def test_samples_for_audio_tokens_returns_minimum_sample_count(chunk_size_seconds: float | None) -> None:
+    target_tokens = 17
+
+    samples = NeMoSpeechLMProcessingInfo._samples_for_audio_tokens(target_tokens, chunk_size_seconds)
+
+    assert NeMoSpeechLMProcessingInfo._estimate_audio_tokens(samples, chunk_size_seconds) >= target_tokens
+    assert NeMoSpeechLMProcessingInfo._estimate_audio_tokens(samples - 1, chunk_size_seconds) < target_tokens
+
+
+def test_samples_for_audio_tokens_rejects_unreachable_target() -> None:
+    max_samples = int(_DUMMY_AUDIO_MAX_DURATION_S * _SAMPLING_RATE)
+    max_tokens = NeMoSpeechLMProcessingInfo._estimate_audio_tokens(max_samples)
+
+    with pytest.raises(ValueError, match="Cannot produce"):
+        NeMoSpeechLMProcessingInfo._samples_for_audio_tokens(max_tokens + 1)

tests/collections/speechlm2/test_vllm_plugin.py

@@ -0,0 +1,605 @@
+# 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.
+
+"""Unit tests for the vLLM NeMo Speech LM (SALM) plugin.
+
+Covers plugin registration, config loading + escape-hatch wiring, special
+token handling, and backend selection -- without requiring GPU or model
+weights.
+"""
+
+import importlib.util
+from types import SimpleNamespace
+
+import pytest
+
+try:
+    from nemo.collections.speechlm2.vllm.salm import config as _config_module
+
+    NeMoSpeechLMConfig = _config_module.NeMoSpeechLMConfig
+
+    _HAS_CONFIG = True
+except (ImportError, RuntimeError):
+    _HAS_CONFIG = False
+
+_HAS_VLLM = importlib.util.find_spec("vllm") is not None
+_DEFAULT_CONFIG_KWARGS = {
+    "pretrained_llm": "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16",
+    "pretrained_asr": "nvidia/canary-1b-v2",
+    "audio_locator_tag": "<|audio|>",
+    "prompt_format": "nemotron-nano-v3",
+    "pretrained_weights": True,
+}
+
+
+@pytest.mark.skipif(not _HAS_CONFIG, reason="NeMoSpeechLMConfig not available")
+class TestNeMoSpeechLMConfig:
+    """Tests for NeMoSpeechLMConfig."""
+
+    @pytest.fixture(autouse=True)
+    def mock_backbone_config(self, monkeypatch):
+        def from_pretrained(model_name: str, trust_remote_code: bool = True):
+            if "Nemotron" in model_name:
+                return SimpleNamespace(
+                    architectures=["NemotronHybridForCausalLM"],
+                    hidden_size=2048,
+                    vocab_size=131072,
+                    num_hidden_layers=4,
+                    num_key_value_heads=2,
+                    layer_norm_epsilon=1e-5,
+                )
+            return SimpleNamespace(
+                architectures=["Qwen3ForCausalLM"],
+                hidden_size=2048,
+                vocab_size=151936,
+                num_hidden_layers=4,
+                rms_norm_eps=1e-6,
+            )
+
+        monkeypatch.setattr(_config_module.AutoConfig, "from_pretrained", from_pretrained)
+
+    def test_model_type(self):
+        assert NeMoSpeechLMConfig.model_type == "nemo_speechlm"
+
+    def test_default_construction_for_hf_serialization(self):
+        """HF internally constructs a no-arg config when serializing configs."""
+        cfg = NeMoSpeechLMConfig()
+        assert cfg.pretrained_llm is None
+        assert cfg.pretrained_asr is None
+        assert cfg.audio_locator_tag is None
+        assert cfg.prompt_format is None
+        assert cfg.pretrained_weights is None
+        assert cfg.llm_architectures == []
+        assert cfg.get_text_config() is cfg.text_config
+
+    def test_loads_text_config(self):
+        """Config should load a text_config from the pretrained LLM."""
+        cfg = NeMoSpeechLMConfig(**_DEFAULT_CONFIG_KWARGS)
+        assert cfg.text_config is not None
+        assert hasattr(cfg.text_config, "hidden_size")
+        assert cfg.get_text_config() is cfg.text_config
+
+    def test_hybrid_backbone_aliases_for_vllm(self):
+        cfg = NeMoSpeechLMConfig(**_DEFAULT_CONFIG_KWARGS)
+        assert cfg.is_hybrid is True
+        assert cfg.llm_architectures == ["NemotronHForCausalLM"]
+        assert cfg.text_config.total_num_kv_heads == cfg.text_config.num_key_value_heads
+        assert cfg.text_config.rms_norm_eps == cfg.text_config.layer_norm_epsilon
+
+    @pytest.mark.parametrize(
+        "architectures, expected_is_hybrid",
+        [
+            (["NemotronHForCausalLM"], True),
+            (["NemotronHybridForCausalLM"], True),
+            (["Qwen3ForCausalLM"], False),
+            (["LlamaForCausalLM"], False),
+            (["Qwen2ForCausalLM"], False),
+        ],
+    )
+    def test_is_hybrid_backend_helper(self, architectures, expected_is_hybrid):
+        """``_is_hybrid_backend`` should match the documented hybrid allow-list."""
+        from nemo.collections.speechlm2.vllm.salm.config import _is_hybrid_backend
+
+        assert _is_hybrid_backend(architectures) is expected_is_hybrid
+
+    @pytest.mark.parametrize(
+        "backbone_archs, expected_is_hybrid",
+        [
+            (["NemotronHForCausalLM"], True),
+            (["NemotronHybridForCausalLM"], True),
+            (["Qwen3ForCausalLM"], False),
+        ],
+    )
+    def test_is_hybrid_set_from_backbone_architectures(self, monkeypatch, backbone_archs, expected_is_hybrid):
+        """``cfg.is_hybrid`` is driven by the backbone HF config's ``architectures``."""
+
+        def from_pretrained(model_name: str, trust_remote_code: bool = True):
+            kwargs = dict(
+                architectures=backbone_archs,
+                hidden_size=2048,
+                vocab_size=131072,
+                num_hidden_layers=4,
+            )
+            if expected_is_hybrid:
+                kwargs.update(num_key_value_heads=2, layer_norm_epsilon=1e-5)
+            else:
+                kwargs.update(rms_norm_eps=1e-6)
+            return SimpleNamespace(**kwargs)
+
+        monkeypatch.setattr(_config_module.AutoConfig, "from_pretrained", from_pretrained)
+
+        cfg = NeMoSpeechLMConfig(**_DEFAULT_CONFIG_KWARGS)
+        assert cfg.is_hybrid is expected_is_hybrid
+
+    def test_hybrid_backbone_does_not_set_layer_types_shim(self):
+        """Hybrid backbones must NOT have layer_types overridden -- the runtime
+        is_hybrid escape hatch only fires when every layer is 'attention'."""
+        cfg = NeMoSpeechLMConfig(**_DEFAULT_CONFIG_KWARGS)
+        assert cfg.is_hybrid is True
+        assert getattr(cfg.text_config, "layer_types", None) is None
+
+    def test_transformer_backbone_engages_layer_types_shim(self):
+        """Non-hybrid backbones get layer_types=['attention']*N so vLLM's
+        ModelConfig.is_hybrid property returns False at runtime even though
+        the model class declares IsHybrid (needed for NemotronH path)."""
+        cfg = NeMoSpeechLMConfig(
+            **{
+                **_DEFAULT_CONFIG_KWARGS,
+                "pretrained_llm": "Qwen/Qwen3-1.7B",
+            }
+        )
+        assert cfg.is_hybrid is False
+        assert cfg.text_config.layer_types == ["attention"] * 4
+
+    def test_custom_pretrained_llm(self):
+        """Config should accept different LLM backbones."""
+        cfg = NeMoSpeechLMConfig(
+            **{
+                **_DEFAULT_CONFIG_KWARGS,
+                "pretrained_llm": "Qwen/Qwen3-1.7B",
+            }
+        )
+        assert cfg.pretrained_llm == "Qwen/Qwen3-1.7B"
+        assert cfg.text_config is not None
+        assert cfg.llm_architectures == ["Qwen3ForCausalLM"]
+
+    def test_audio_locator_tag_default_accepted(self):
+        cfg = NeMoSpeechLMConfig(**_DEFAULT_CONFIG_KWARGS)
+        assert cfg.audio_locator_tag == "<|audio|>"
+
+    def test_audio_locator_tag_custom_rejected(self):
+        """Plugin only supports ``<|audio|>``; mismatched checkpoints fail at load time."""
+        with pytest.raises(ValueError, match="audio_locator_tag"):
+            NeMoSpeechLMConfig(
+                **{
+                    **_DEFAULT_CONFIG_KWARGS,
+                    "audio_locator_tag": "<|custom_audio|>",
+                }
+            )
+
+    @pytest.mark.parametrize(
+        "field",
+        [
+            "pretrained_llm",
+            "pretrained_asr",
+            "audio_locator_tag",
+            "prompt_format",
+            "pretrained_weights",
+        ],
+    )
+    def test_required_exported_fields(self, field):
+        kwargs = dict(_DEFAULT_CONFIG_KWARGS)
+        kwargs.pop(field)
+        with pytest.raises(ValueError, match=field):
+            NeMoSpeechLMConfig(**kwargs)
+
+    def test_unknown_attr_raises(self):
+        cfg = NeMoSpeechLMConfig(**_DEFAULT_CONFIG_KWARGS)
+        with pytest.raises(AttributeError):
+            _ = cfg.nonexistent_attribute_xyz
+
+    def test_encoder_chunk_size_seconds_default_none(self):
+        """Legacy checkpoints without a chunk size keep the single-pass encoder path."""
+        cfg = NeMoSpeechLMConfig(**_DEFAULT_CONFIG_KWARGS)
+        assert cfg.encoder_chunk_size_seconds is None
+
+    def test_encoder_chunk_size_seconds_round_trips(self):
+        """Chunk size set in config.json (e.g. SALMAutomodel default 30 s) survives load."""
+        cfg = NeMoSpeechLMConfig(
+            **{
+                **_DEFAULT_CONFIG_KWARGS,
+                "encoder_chunk_size_seconds": 30.0,
+            }
+        )
+        assert cfg.encoder_chunk_size_seconds == 30.0
+
+    def test_encoder_chunk_size_seconds_default_init_inert(self):
+        """No-arg default init must still expose ``encoder_chunk_size_seconds=None``."""
+        cfg = NeMoSpeechLMConfig()
+        assert cfg.encoder_chunk_size_seconds is None
+
+
+@pytest.mark.skipif(not (_HAS_CONFIG and _HAS_VLLM), reason="NeMoSpeechLMConfig or vLLM not available")
+class TestBackendSelection:
+    """Tests for ``backends.make_backend`` dispatch on hybrid/transformer configs."""
+
+    @pytest.fixture(autouse=True)
+    def mock_backbone_config(self, monkeypatch):
+        def from_pretrained(model_name: str, trust_remote_code: bool = True):
+            if "Nemotron" in model_name:
+                return SimpleNamespace(
+                    architectures=["NemotronHybridForCausalLM"],
+                    hidden_size=2048,
+                    vocab_size=131072,
+                    num_hidden_layers=4,
+                    num_key_value_heads=2,
+                    layer_norm_epsilon=1e-5,
+                )
+            return SimpleNamespace(
+                architectures=["Qwen3ForCausalLM"],
+                hidden_size=2048,
+                vocab_size=151936,
+                num_hidden_layers=4,
+                rms_norm_eps=1e-6,
+            )
+
+        monkeypatch.setattr(_config_module.AutoConfig, "from_pretrained", from_pretrained)
+
+    def test_hybrid_config_picks_hybrid_backend(self):
+        from nemo.collections.speechlm2.vllm.salm.backends import HybridBackend, make_backend
+
+        cfg = NeMoSpeechLMConfig(**_DEFAULT_CONFIG_KWARGS)
+        backend = make_backend(cfg)
+        assert isinstance(backend, HybridBackend)
+        assert backend.architectures() == ["NemotronHForCausalLM"]
+
+    def test_transformer_config_picks_transformer_backend(self):
+        from nemo.collections.speechlm2.vllm.salm.backends import TransformerBackend, make_backend
+
+        cfg = NeMoSpeechLMConfig(
+            **{
+                **_DEFAULT_CONFIG_KWARGS,
+                "pretrained_llm": "Qwen/Qwen3-1.7B",
+            }
+        )
+        backend = make_backend(cfg)
+        assert isinstance(backend, TransformerBackend)
+        assert backend.architectures() == ["Qwen3ForCausalLM"]
+
+
+@pytest.mark.skipif(not _HAS_VLLM, reason="vLLM not installed")
+class TestSpecialTokens:
+    """Tests for special token handling."""
+
+    def test_adds_missing_token(self):
+        from unittest.mock import MagicMock
+
+        from nemo.collections.speechlm2.vllm.salm.audio import _ensure_special_tokens
+
+        tokenizer = MagicMock()
+        tokenizer.get_vocab.return_value = {}
+        _ensure_special_tokens(tokenizer)
+        tokenizer.add_special_tokens.assert_called_once()
+
+    def test_skips_existing_token(self):
+        from unittest.mock import MagicMock
+
+        from nemo.collections.speechlm2.vllm.salm.audio import _ensure_special_tokens
+
+        tokenizer = MagicMock()
+        tokenizer.get_vocab.return_value = {"<|audio|>": 99}
+        _ensure_special_tokens(tokenizer)
+        tokenizer.add_special_tokens.assert_not_called()
+
+    def test_placeholder_str(self):
+        from nemo.collections.speechlm2.vllm.salm.model import NeMoSpeechLMForConditionalGeneration
+
+        assert NeMoSpeechLMForConditionalGeneration.get_placeholder_str("audio", 0) == "<|audio|>"
+        assert NeMoSpeechLMForConditionalGeneration.get_placeholder_str("image", 0) is None
+
+
+@pytest.mark.skipif(not _HAS_VLLM, reason="vLLM not installed")
+class TestAudioProcessing:
+    """Tests for audio encoding with a tiny perception module."""
+
+    def test_data_parser_normalizes_audio(self, monkeypatch):
+        from nemo.collections.speechlm2.vllm.salm.audio import NeMoSpeechLMProcessingInfo
+
+        info = object.__new__(NeMoSpeechLMProcessingInfo)
+        monkeypatch.setattr(info, "_get_expected_hidden_size", lambda: 2048)
+
+        parser = info.get_data_parser()
+
+        assert parser.audio_resampler.target_sr == 16000
+        assert parser.target_channels == 1
+
+    def test_processing_info_has_no_audio_duration_limit(self):
+        from nemo.collections.speechlm2.vllm.salm.audio import NeMoSpeechLMProcessingInfo
+
+        info = object.__new__(NeMoSpeechLMProcessingInfo)
+
+        assert not hasattr(info, "get_max_audio_len")
+        assert not hasattr(info, "get_max_audio_tokens")
+
+    def test_dummy_inputs_use_profiling_audio_length(self):
+        from nemo.collections.speechlm2.vllm.salm.audio import (
+            NeMoSpeechLMDummyInputsBuilder,
+            NeMoSpeechLMProcessingInfo,
+        )
+
+        info = object.__new__(NeMoSpeechLMProcessingInfo)
+        builder = object.__new__(NeMoSpeechLMDummyInputsBuilder)
+        builder.info = info
+
+        result = builder.get_dummy_mm_data(seq_len=0, mm_counts={"audio": 1}, mm_options={})
+
+        assert result["audio"][0].shape[-1] == 40 * 16000
+
+    def test_dummy_inputs_use_requested_audio_length(self, monkeypatch):
+        from nemo.collections.speechlm2.vllm.salm.audio import NeMoSpeechLMDummyInputsBuilder
+
+        builder = object.__new__(NeMoSpeechLMDummyInputsBuilder)
+        builder.info = SimpleNamespace(_get_encoder_chunk_size_seconds=lambda: None)
+        monkeypatch.setattr(
+            builder,
+            "_get_dummy_audios",
+            lambda length, num_audios: [SimpleNamespace(length=length) for _ in range(num_audios)],
+        )
+
+        result = builder.get_dummy_mm_data(
+            seq_len=0,
+            mm_counts={"audio": 1},
+            mm_options={"audio": SimpleNamespace(length=12345)},
+        )
+
+        assert result["audio"][0].length == 12345
+
+    def test_dummy_inputs_cap_requested_audio_length_to_text_budget(self, monkeypatch):
+        from nemo.collections.speechlm2.vllm.salm.audio import (
+            _DUMMY_AUDIO_TEXT_TOKEN_RESERVE,
+            NeMoSpeechLMDummyInputsBuilder,
+            NeMoSpeechLMProcessingInfo,
+        )
+
+        target_audio_tokens = 4
+        max_audio_len = NeMoSpeechLMProcessingInfo._samples_for_audio_tokens(target_audio_tokens)
+        builder = object.__new__(NeMoSpeechLMDummyInputsBuilder)
+        builder.info = SimpleNamespace(_get_encoder_chunk_size_seconds=lambda: None)
+        monkeypatch.setattr(
+            builder,
+            "_get_dummy_audios",
+            lambda length, num_audios: [SimpleNamespace(length=length) for _ in range(num_audios)],
+        )
+
+        result = builder.get_dummy_mm_data(
+            seq_len=_DUMMY_AUDIO_TEXT_TOKEN_RESERVE + target_audio_tokens,
+            mm_counts={"audio": 1},
+            mm_options={"audio": SimpleNamespace(length=max_audio_len + 16000)},
+        )
+
+        assert result["audio"][0].length == max_audio_len
+
+    def test_dummy_inputs_large_seq_len_uses_max_audio_cap(self, monkeypatch):
+        from nemo.collections.speechlm2.vllm.salm.audio import (
+            _DUMMY_AUDIO_MAX_DURATION_S,
+            _SAMPLING_RATE,
+            NeMoSpeechLMDummyInputsBuilder,
+        )
+
+        max_audio_len = int(_DUMMY_AUDIO_MAX_DURATION_S * _SAMPLING_RATE)
+        builder = object.__new__(NeMoSpeechLMDummyInputsBuilder)
+        builder.info = SimpleNamespace(_get_encoder_chunk_size_seconds=lambda: None)
+        monkeypatch.setattr(
+            builder,
+            "_get_dummy_audios",
+            lambda length, num_audios: [SimpleNamespace(length=length) for _ in range(num_audios)],
+        )
+
+        result = builder.get_dummy_mm_data(
+            seq_len=10_000_000,
+            mm_counts={"audio": 1},
+            mm_options={"audio": SimpleNamespace(length=max_audio_len + 16000)},
+        )
+
+        assert result["audio"][0].length == max_audio_len
+
+    def test_call_hf_processor_requires_matching_placeholder_count(self):
+        from nemo.collections.speechlm2.vllm.salm.audio import NeMoSpeechLMMultiModalProcessor
+
+        processor = object.__new__(NeMoSpeechLMMultiModalProcessor)
+        processor.info = SimpleNamespace(
+            get_tokenizer=_FakeTokenizer,
+            _estimate_audio_tokens=lambda samples, chunk_size_seconds=None: 2,
+            _get_encoder_chunk_size_seconds=lambda: None,
+        )
+
+        with pytest.raises(ValueError, match="placeholders"):
+            processor._call_hf_processor(
+                prompt="Transcribe this audio",
+                mm_data={"audios": [[0.0] * 16000]},
+                mm_kwargs={},
+                tok_kwargs={},
+            )
+
+    def test_call_hf_processor_emits_true_audio_lengths(self):
+        import torch
+
+        from nemo.collections.speechlm2.vllm.salm.audio import NeMoSpeechLMMultiModalProcessor
+
+        processor = object.__new__(NeMoSpeechLMMultiModalProcessor)
+        processor.info = SimpleNamespace(
+            get_tokenizer=_FakeTokenizer,
+            _estimate_audio_tokens=lambda samples, chunk_size_seconds=None: 2,
+            _get_encoder_chunk_size_seconds=lambda: None,
+        )
+
+        result = processor._call_hf_processor(
+            prompt="Transcribe: <|audio|>",
+            mm_data={"audios": [[0.0] * 12345]},
+            mm_kwargs={},
+            tok_kwargs={},
+        )
+
+        assert len(result["audio_signal"]) == 1
+        assert result["audio_signal"][0].shape[-1] == 12345
+        assert torch.equal(result["audio_signal_length"], torch.tensor([12345]))
+
+    def test_perception_forward(self):
+        """A small NeMo perception module should encode dummy audio to embeddings."""
+        import torch
+
+        if not torch.cuda.is_available():
+            pytest.skip("CUDA required")
+        from nemo.collections.speechlm2.vllm.salm.audio import _load_nemo_perception
+
+        perception_cfg = {
+            "output_dim": 256,
+            "encoder": {
+                "_target_": "nemo.collections.asr.modules.ConformerEncoder",
+                "feat_in": 128,
+                "feat_out": -1,
+                "n_layers": 2,
+                "d_model": 256,
+                "subsampling": "dw_striding",
+                "subsampling_factor": 8,
+                "subsampling_conv_channels": 64,
+                "ff_expansion_factor": 4,
+                "self_attention_model": "rel_pos",
+                "n_heads": 4,
+                "conv_kernel_size": 9,
+                "conv_norm_type": "batch_norm",
+                "dropout": 0.0,
+                "dropout_pre_encoder": 0.0,
+                "dropout_emb": 0.0,
+                "dropout_att": 0.0,
+            },
+            "modality_adapter": {
+                "_target_": "nemo.collections.speechlm2.modules.perception.IdentityConnector",
+                "d_model": 256,
+            },
+            "preprocessor": {
+                "_target_": "nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor",
+                "sample_rate": 16000,
+                "normalize": "per_feature",
+                "window_size": 0.025,
+                "window_stride": 0.01,
+                "window": "hann",
+                "features": 128,
+                "n_fft": 512,
+                "log": True,
+                "frame_splicing": 1,
+                "dither": 0.0,
+                "pad_to": 0,
+                "pad_value": 0.0,
+            },
+        }
+
+        perception = _load_nemo_perception(perception_cfg)
+        perception = perception.to("cuda", dtype=torch.float32)
+
+        dummy_audio = torch.randn(1, 16000, device="cuda")
+        audio_len = torch.tensor([16000], device="cuda")
+
+        with torch.no_grad():
+            embeds, embed_lens = perception(input_signal=dummy_audio, input_signal_length=audio_len)
+
+        assert embeds.ndim == 3
+        assert embeds.shape[0] == 1
+        assert embeds.shape[2] == 256
+        assert embed_lens[0] > 0
+
+
+@pytest.mark.skipif(not _HAS_VLLM, reason="vLLM not installed")
+class TestPluginRegistration:
+    """Tests for plugin registration with vLLM."""
+
+    def test_register_config(self, monkeypatch):
+        """register() should add nemo_speechlm to vLLM's config registry."""
+        from transformers import AutoConfig
+
+        from nemo.collections.speechlm2.vllm.salm import register
+
+        monkeypatch.setattr(
+            AutoConfig, "from_pretrained", lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError())
+        )
+
+        register()
+
+        from vllm.transformers_utils.config import _CONFIG_REGISTRY
+
+        assert "nemo_speechlm" in _CONFIG_REGISTRY
+
+    def test_register_model(self, monkeypatch):
+        """register() should make NeMoSpeechLMForConditionalGeneration importable.
+
+        The plugin now registers a single architecture name; the obsolete
+        ``NeMoSpeechLMHybridForConditionalGeneration`` no longer appears.
+        """
+        from transformers import AutoConfig
+
+        from nemo.collections.speechlm2.vllm.salm import register
+
+        monkeypatch.setattr(
+            AutoConfig, "from_pretrained", lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError())
+        )
+
+        register()
+
+        from vllm.model_executor.models.registry import ModelRegistry
+
+        from nemo.collections.speechlm2.vllm.salm.model import NeMoSpeechLMForConditionalGeneration
+
+        assert "NeMoSpeechLMForConditionalGeneration" in ModelRegistry.get_supported_archs()
+        assert NeMoSpeechLMForConditionalGeneration is not None
+
+    def test_register_does_not_patch_fast_tokenizer(self, monkeypatch):
+        from transformers import AutoConfig, PreTrainedTokenizerFast
+
+        from nemo.collections.speechlm2.vllm.salm import register
+
+        monkeypatch.setattr(
+            AutoConfig, "from_pretrained", lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError())
+        )
+
+        assert "_orig_batch_encode_plus" not in PreTrainedTokenizerFast.__dict__
+        register()
+        assert "_orig_batch_encode_plus" not in PreTrainedTokenizerFast.__dict__
+
+    def test_register_does_not_load_backbone_config(self, monkeypatch):
+        from unittest.mock import Mock
+
+        from transformers import AutoConfig
+
+        from nemo.collections.speechlm2.vllm.salm import register
+
+        from_pretrained = Mock(side_effect=AssertionError("register() must not load remote backbone configs"))
+        monkeypatch.setattr(AutoConfig, "from_pretrained", from_pretrained)
+
+        register()
+
+        from_pretrained.assert_not_called()
+
+
+class _FakeTokenizer:
+    def __init__(self):
+        self.added_special_tokens = None
+
+    def get_vocab(self):
+        return {}
+
+    def add_special_tokens(self, tokens):
+        self.added_special_tokens = tokens
+
+    def encode(self, prompt, add_special_tokens=True):
+        return list(range(max(1, len(prompt.split()))))