diff --git a/examples/speechlm2/conf/streaming_stt_automodel.yaml b/examples/speechlm2/conf/streaming_stt_automodel.yaml
new file mode 100644
index 000000000000..4ae64f4fdae3
--- /dev/null
+++ b/examples/speechlm2/conf/streaming_stt_automodel.yaml
@@ -0,0 +1,211 @@
+# Streaming STT SpeechLM trained with the NeMo Automodel LLM backend.
+#
+# Same recipe as ``streaming_stt_lora_multichunk.yaml``; the deltas are:
+#   * model.use_nemo_automodel: true      -> StreamingSTTModelAutomodel
+#   * model.lora uses Automodel PeftConfig keys (dim/alpha/dropout/target_modules)
+#     instead of HuggingFace PEFT keys (r/lora_alpha/lora_dropout)
+#   * trainer.strategy: AutomodelParallelStrategy (FSDP2/TP/PP/CP/EP) instead of DDP
+model:
+  pretrained_llm: Qwen/Qwen3-1.7B
+  pretrained_asr: nvidia/nemotron-speech-streaming-en-0.6b
+  load_llm_weights: true
+  load_asr_weights: true
+
+  # Use the NeMo Automodel backend (nemo_automodel package required).
+  # The LLM is loaded *and sharded* inside configure_model(), so each rank only
+  # materializes its own shard.
+  use_nemo_automodel: true
+
+  sample_rate: 16000
+  blank_token: "<blank>"
+  frame_length_in_secs: 0.08
+  # --- Multi chunk-size training ---
+  # A list of positive frame counts enables multi chunk-size training: one value
+  # is drawn at random per batch. Inference defaults to the longest value (here
+  # 14); override per call with generate(chunk_size_override=N).
+  chunk_size: [2, 6, 14]
+  # Encoder attention look-ahead. The LEFT context (70) is fixed; the RIGHT
+  # context (look-ahead) is set automatically each batch to (chunk_size - 1).
+  att_context_size: [70, 13]
+  audio_pad_to: 16
+
+  freeze_speech_encoder: false
+  freeze_modality_adapter: false
+  freeze_modality_proj: false
+  freeze_llm_model: true
+  freeze_llm_head: false
+  freeze_embed_tokens: false
+
+  freeze_params: []
+  prevent_freeze_params: []  # Use to make specific submodules trainable; overrides freeze_params
+
+  # Fine-tune from a previous training checkpoint (model weights only — optimizer,
+  # scheduler, and step counter start fresh). Supports DCP directories (from
+  # FSDP2/TP training), HuggingFace directories (model.safetensors), and
+  # single-file .ckpt checkpoints.
+  init_from_checkpoint: null
+
+  # LoRA on the LLM, applied by Automodel's PEFT implementation (NOT HuggingFace
+  # PEFT — the key names differ; passing `r`/`lora_alpha`/`lora_dropout` here
+  # raises an explicit error). LoRA params stay trainable even when the LLM is
+  # frozen. Comment the block out to disable.
+  lora:
+    dim: 128                 # HF PEFT `r`
+    alpha: 256               # HF PEFT `lora_alpha`
+    dropout: 0.01            # HF PEFT `lora_dropout`
+    target_modules: ["q_proj", "k_proj", "v_proj", "o_proj"]
+    # match_all_linear: false
+    # exclude_modules: []
+    # use_dora: false
+
+  # Automodel backend dispatch. ONLY valid for Automodel-native backbones (the
+  # MoE/Nemotron implementations); plain HuggingFace architectures such as
+  # Qwen3 are constructed by transformers and raise
+  # `TypeError: __init__() got an unexpected keyword argument 'backend'`.
+  # Defaults auto-select TransformerEngine/DeepEP when available.
+  # automodel_backend:
+  #   attn: sdpa                # "te" | "sdpa" | "flex"
+  #   linear: torch             # "torch" | "te"
+  #   rms_norm: torch_fp32      # "torch" | "torch_fp32" | "te"
+  # sdpa_method: ["flash_attention"]
+
+  # MoE knobs — only meaningful with an MoE ``pretrained_llm``; harmless no-ops
+  # for dense backbones such as Qwen3-1.7B.
+  aux_loss_coeff: 0.0
+  train_gate: false
+  moe_metrics:
+    enabled: false
+    mode: brief
+    detailed_every_steps: null
+    top_k_experts: 5
+
+  perception:
+    spec_augment:
+      _target_: nemo.collections.asr.modules.SpectrogramAugmentation
+      freq_masks: 2 # set to zero to disable it
+      # you may use lower time_masks for smaller models to have a faster convergence
+      time_masks: 10 # set to zero to disable it
+      freq_width: 27
+      time_width: 0.05
+
+    modality_adapter:
+      _target_: nemo.collections.speechlm2.modules.perception.IdentityConnector
+
+  optimizer:
+    _target_: torch.optim.AdamW
+    lr: 5e-4
+    betas: [0.9, 0.98]
+    weight_decay: 1e-3
+    foreach: true
+
+  lr_scheduler:
+    _target_: nemo.core.optim.lr_scheduler.CosineAnnealing
+    warmup_steps: 1000
+    min_lr: 1e-6
+    max_steps: ${trainer.max_steps}
+
+data:
+  # --- StreamingSTTDataset config ---
+  dataset:
+    sample_rate: ${model.sample_rate}
+    frame_length_in_secs: ${model.frame_length_in_secs}
+    # Inherits the multi chunk-size list from model.chunk_size; the dataset draws
+    # one value per batch and records it in StreamingSTTBatch.chunk_size so the
+    # model can match the encoder look-ahead.
+    chunk_size: ${model.chunk_size}
+    num_delay_frames: 0
+    audio_tag: "<audio>"
+    blank_token: ${model.blank_token}
+    system_role: system
+    system_prompt: "Transcribe the audio into text."
+
+  val_dataset_overrides:
+    chunk_size: 6
+
+  # --- Lhotse dataloader config ---
+  train_ds:
+    sample_rate: ${data.dataset.sample_rate}
+    input_cfg: null
+    manifest_filepath: null
+    seed: 42
+    shuffle: true
+    shard_seed: "randomized"
+    num_workers: 4
+    batch_size: 4
+
+  validation_ds:
+    datasets:
+      val_set_0:
+        input_cfg: null
+        manifest_filepath: null
+    sample_rate: ${data.dataset.sample_rate}
+    batch_size: 4
+    num_workers: 4
+    seed: 42
+    shard_seed: "randomized"
+
+trainer:
+  devices: -1
+  accelerator: gpu
+  num_nodes: 1
+  # bf16-true keeps params in bf16; the model derives its dtype from this value.
+  precision: bf16-true
+  logger: false
+  enable_checkpointing: false
+  use_distributed_sampler: false
+  max_steps: 100000
+  limit_train_batches: 5000
+  val_check_interval: ${trainer.limit_train_batches}
+  limit_val_batches: 10
+  log_every_n_steps: 10
+  num_sanity_val_steps: 1
+  gradient_clip_val: 1.0
+  accumulate_grad_batches: 1
+  strategy:
+    # AutomodelParallelStrategy delegates device mesh creation to nemo_automodel.
+    # The model's configure_model() receives the device_mesh and passes it to
+    # Automodel's from_pretrained for memory-efficient loading (each GPU only
+    # loads its own shard).
+    # This model supports FSDP2/HSDP (dp_size, dp_replicate_size) and EP.
+    # tp_size > 1 and cp_size > 1 are REJECTED at fit start: the interleaved
+    # audio/text sequence must stay on one rank, and the streaming loss is not
+    # DTensor-safe against a vocab-sharded lm_head.
+    _target_: nemo.collections.speechlm2.parts.parallel.AutomodelParallelStrategy
+    dp_size: null         # null = inferred from world_size / other dims
+    dp_replicate_size: 1  # HSDP replication group size (>1 enables hybrid sharding)
+    tp_size: 1            # must stay 1 (see above)
+    pp_size: 1            # must stay 1 (no pipeline schedule in this recipe)
+    cp_size: 1            # must stay 1 (see above)
+    ep_size: 1            # Expert parallel size (MoE backbones only)
+    activation_checkpointing_llm: false
+    activation_checkpointing_perception: false
+
+
+exp_manager:
+  exp_dir: null
+  explicit_log_dir: null
+  name: streaming_stt_automodel
+  create_tensorboard_logger: false
+  create_checkpoint_callback: true
+  use_datetime_version: true
+  max_time_per_run: "00:03:50:00"
+
+  resume_from_checkpoint: null
+  resume_if_exists: true
+  resume_ignore_no_checkpoint: true
+
+  create_wandb_logger: false
+  wandb_logger_kwargs:
+    name: ${exp_manager.name}
+    project: null
+    resume: true
+
+  checkpoint_callback_params:
+    filename: "{step}"
+    monitor: val_acc
+    mode: max
+    every_n_train_steps: null
+    every_n_epochs: 1
+    save_top_k: 1
+    always_save_nemo: false
+    save_nemo_on_train_end: false
diff --git a/examples/speechlm2/streaming_stt_train.py b/examples/speechlm2/streaming_stt_train.py
index aacea7c94bbb..8c58cdfad7c1 100644
--- a/examples/speechlm2/streaming_stt_train.py
+++ b/examples/speechlm2/streaming_stt_train.py
@@ -63,8 +63,18 @@ def train(cfg):
         forced_aligner = None
         defer_get_batch = False
 
+    # Set model.use_nemo_automodel=true to train with the NeMo Automodel LLM
+    # backend (FSDP2/TP/EP via AutomodelParallelStrategy) instead of the
+    # HuggingFace backend with DDP. The two classes share the same recipe;
+    # see conf/streaming_stt_automodel.yaml for the config deltas.
+    model_cls = StreamingSTTModel
+    if cfg.model.get("use_nemo_automodel", False):
+        from nemo.collections.speechlm2 import StreamingSTTModelAutomodel
+
+        model_cls = StreamingSTTModelAutomodel
+
     with trainer.init_module():
-        model = StreamingSTTModel(
+        model = model_cls(
             OmegaConf.to_container(cfg.model, resolve=True),
             forced_aligner=forced_aligner,
             data_cfg=dataset_cfg,
diff --git a/nemo/collections/speechlm2/__init__.py b/nemo/collections/speechlm2/__init__.py
index ed05ec96b0e4..1e64c49dec97 100644
--- a/nemo/collections/speechlm2/__init__.py
+++ b/nemo/collections/speechlm2/__init__.py
@@ -22,6 +22,7 @@
     SALMAutomodel,
     SALMWithAsrDecoder,
     StreamingSTTModel,
+    StreamingSTTModelAutomodel,
 )
 
 __all__ = [
@@ -40,4 +41,5 @@
     'SALMWithAsrDecoder',
     'NemotronVoiceChat',
     'StreamingSTTModel',
+    'StreamingSTTModelAutomodel',
 ]
diff --git a/nemo/collections/speechlm2/data/streaming_stt_dataset.py b/nemo/collections/speechlm2/data/streaming_stt_dataset.py
index 30be692f4419..77e22bf52230 100644
--- a/nemo/collections/speechlm2/data/streaming_stt_dataset.py
+++ b/nemo/collections/speechlm2/data/streaming_stt_dataset.py
@@ -16,6 +16,7 @@
 import math
 import random
 import re
+import warnings
 from dataclasses import dataclass
 from typing import Iterable, List, Optional, Union
 
@@ -541,6 +542,59 @@ def get_llm_messages_for_batch(
     return batch_messages
 
 
+def resolve_pad_id(tokenizer: AutoTokenizer) -> int:
+    """Padding token ID for a tokenizer that may not define one.
+
+    Falls back to ``<unk>`` and finally to id 0. Some LLM tokenizers ship without
+    a pad token (e.g. NVIDIA-Nemotron-3-Nano, where ``pad_id`` is ``None`` and
+    ``unk_id`` is 0) — passing that ``None`` into ``pad_sequence`` raises
+    ``TypeError: argument 'padding_value' must be float, not NoneType``.
+
+    The dataset and :attr:`StreamingSTTModel.text_pad_id` MUST agree on this
+    value: the model derives its attention mask as ``input_tokens != pad_id``, so
+    a mismatch would either unmask padding or mask real tokens. Both call here.
+    """
+    pad_id = tokenizer.pad_id
+    if pad_id is None:
+        pad_id = getattr(tokenizer, "unk_id", None)
+    if pad_id is None:
+        warnings.warn(
+            "The text tokenizer has no <pad> or <unk> token; using id 0 for padding "
+            "(this may lead to silent bugs).",
+            stacklevel=2,
+        )
+        pad_id = 0
+    return int(pad_id)
+
+
+def apply_chat_template_ids(hf_tok, messages: List[dict], **kwargs) -> list[int]:
+    """``apply_chat_template(..., tokenize=True)`` normalized to a flat list of token IDs.
+
+    transformers 4.x returns a plain ``list[int]`` from a tokenizing call, while
+    transformers 5.x always returns a ``BatchEncoding`` (``return_dict`` became
+    the default). Callers here only ever want the IDs, so unwrap the mapping —
+    and a batched ``[[ids]]`` layout too, in case a future version wraps single
+    conversations.
+
+    Args:
+        hf_tok: A HuggingFace tokenizer (``tokenizer.tokenizer``).
+        messages: ``[{"role": ..., "content": ...}, ...]``.
+        kwargs: Forwarded to ``apply_chat_template`` (e.g. ``add_generation_prompt``,
+            ``enable_thinking``). Do not pass ``tokenize`` or ``return_dict``.
+
+    Returns:
+        The conversation's token IDs as a flat ``list[int]``.
+    """
+    out = hf_tok.apply_chat_template(messages, tokenize=True, **kwargs)
+    if hasattr(out, "keys"):  # BatchEncoding / dict — transformers >= 5
+        out = out["input_ids"]
+    if hasattr(out, "tolist"):  # torch tensor (only when return_tensors was requested)
+        out = out.tolist()
+    if len(out) > 0 and isinstance(out[0], (list, tuple)):  # batched [[ids]]
+        out = out[0]
+    return list(out)
+
+
 def parse_chat_template_ids(hf_tok, last_turn: bool = False) -> tuple[list[int], list[int], list[int]]:
     """Discover turn-structure token IDs from a HuggingFace chat template.
 
@@ -679,13 +733,13 @@ def _tokenize_compact_with_assistant_mask(
     # --- System section: keep Qwen3-style wrapping ---
     system_msgs = [m for m in messages if m["role"] == "system"]
     if system_msgs:
-        system_ids = hf_tok.apply_chat_template(
+        system_ids = apply_chat_template_ids(
+            hf_tok,
             system_msgs,
-            tokenize=True,
             add_generation_prompt=False,
             enable_thinking=False,
         )
-        input_ids.extend(list(system_ids))
+        input_ids.extend(system_ids)
         assistant_mask.extend([0] * len(system_ids))
 
     # --- Per-turn compact encoding ---
@@ -777,14 +831,7 @@ def _tokenize_with_assistant_mask(
     assistant_mask = [0] * len(input_ids)
 
     msgs_sentinel = [{**m, "content": _SENTINEL_CHAR} if m["role"] == "assistant" else m for m in messages]
-    ids_sentinel_result = hf_tok.apply_chat_template(
-        msgs_sentinel,
-        tokenize=True,
-        enable_thinking=False,
-    )
-    ids_sentinel = list(
-        ids_sentinel_result["input_ids"] if hasattr(ids_sentinel_result, "keys") else ids_sentinel_result
-    )
+    ids_sentinel = apply_chat_template_ids(hf_tok, msgs_sentinel, enable_thinking=False)
 
     eos_id = getattr(hf_tok, 'eos_token_id', None)
     i, j = 0, 0  # pointers into input_ids and ids_sentinel
@@ -1122,13 +1169,14 @@ def get_batch_data(
             all_input_ids.append(torch.tensor(input_ids, dtype=torch.long))
             all_target_ids.append(torch.tensor(target_ids, dtype=torch.long))
 
+        pad_id = resolve_pad_id(self.tokenizer)
         if chunk_size >= 0:  # fixed chunking or dynamic chunking: right-pad
-            input_tokens = right_collate_vectors(all_input_ids, padding_value=self.tokenizer.pad_id)
+            input_tokens = right_collate_vectors(all_input_ids, padding_value=pad_id)
             target_tokens = right_collate_vectors(all_target_ids, padding_value=IGNORE_INDEX)
             input_token_lens = torch.tensor([len(ids) for ids in all_input_ids], dtype=torch.long)
             target_token_lens = torch.tensor([len(ids) for ids in all_target_ids], dtype=torch.long)
         else:  # offline mode: left-pad
-            input_tokens = left_collate_vectors(all_input_ids, padding_value=self.tokenizer.pad_id)
+            input_tokens = left_collate_vectors(all_input_ids, padding_value=pad_id)
             target_tokens = left_collate_vectors(all_target_ids, padding_value=IGNORE_INDEX)
             # length is the same size as input_tokens.shape[1] since they're left-padded
             input_token_lens = torch.tensor(
diff --git a/nemo/collections/speechlm2/models/__init__.py b/nemo/collections/speechlm2/models/__init__.py
index cb7b7849efe4..c30e1faab028 100644
--- a/nemo/collections/speechlm2/models/__init__.py
+++ b/nemo/collections/speechlm2/models/__init__.py
@@ -20,6 +20,7 @@
 from .salm_asr_decoder import SALMWithAsrDecoder
 from .salm_automodel import SALMAutomodel
 from .streaming_stt_model import StreamingState, StreamingSTTModel
+from .streaming_stt_model_automodel import StreamingSTTModelAutomodel
 
 __all__ = [
     'DuplexS2SModel',
@@ -32,4 +33,5 @@
     'NemotronVoiceChat',
     'StreamingState',
     'StreamingSTTModel',
+    'StreamingSTTModelAutomodel',
 ]
diff --git a/nemo/collections/speechlm2/models/streaming_stt_model.py b/nemo/collections/speechlm2/models/streaming_stt_model.py
index cb47fa3996ac..a6b5f87a05ed 100644
--- a/nemo/collections/speechlm2/models/streaming_stt_model.py
+++ b/nemo/collections/speechlm2/models/streaming_stt_model.py
@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import math
-import warnings
 from collections import defaultdict
 from copy import deepcopy
 from dataclasses import dataclass, field
@@ -39,9 +38,11 @@
     IGNORE_INDEX,
     StreamingSTTBatch,
     StreamingSTTDataset,
+    apply_chat_template_ids,
     build_compact_turn_markers,
     decode_with_blank,
     parse_chat_template_ids,
+    resolve_pad_id,
 )
 from nemo.collections.speechlm2.parts.alignments import ForcedAligner
 from nemo.collections.speechlm2.parts.hf_hub import HFHubMixin
@@ -368,10 +369,80 @@ def __init__(
         self.save_hyperparameters()
         self.cfg = DictConfig(cfg)
         self.core_cfg: StreamingSTTModelConfig = to_dataclass(StreamingSTTModelConfig, cfg)
+        self._normalize_chunk_size()
 
-        # Normalize chunk_size: a list (e.g. [2, 6, 13]) enables multi chunk-size
-        # training (one drawn per batch) and configurable inference (longest by
-        # default). A scalar keeps the original single fixed/dynamic/offline mode.
+        # --- LLM ---
+        self.tokenizer = AutoTokenizer(self.core_cfg.pretrained_llm, use_fast=True)
+        self.llm = load_pretrained_hf(
+            self.core_cfg.pretrained_llm,
+            pretrained_weights=self.core_cfg.load_llm_weights,
+        )
+
+        self._register_special_tokens()
+
+        # Separate embedding layer to avoid FSDP/TP conflicts (same pattern as SALM)
+        self.embed_tokens = self.llm.model.embed_tokens
+        del self.llm.model.embed_tokens
+
+        # --- Speech encoder (perception module) ---
+        self.perception = setup_perception(
+            cfg=self.cfg,
+            output_dim=self.llm.config.hidden_size,
+            pretrained_asr=self.core_cfg.pretrained_asr,
+            pretrained_weights=self.core_cfg.load_asr_weights,
+            audio_pad_to=self.core_cfg.audio_pad_to,
+            att_context_size=self.core_cfg.att_context_size,
+        )
+
+        # --- Aux chunk-boundary classifier (only built when enabled) ---
+        # Only valid in dynamic-chunking mode (chunk_size == 0). When disabled,
+        # no module is built and existing runs / checkpoints see no behavior
+        # change. The boundary decision falls back to the LM head's blank /
+        # user_footer_first signal.
+        if self.core_cfg.use_chunk_classifier:
+            assert self.core_cfg.chunk_size == 0, (
+                "use_chunk_classifier=True requires dynamic chunking "
+                f"(chunk_size=0), got chunk_size={self.core_cfg.chunk_size}"
+            )
+            self._build_chunk_classifier()
+            # Aux training/eval reads self._user_footer_first_id (the BCE positive
+            # label). It's normally set lazily by _ensure_inference_cache, but
+            # training runs before any inference call — so prime the cache now.
+            self._ensure_inference_cache()
+
+        self._apply_freeze_config()
+
+        # --- LoRA ---
+        if "lora" in self.cfg:
+            # Install LoRA after freezing the LLM body to avoid freezing the LoRA weights
+            maybe_install_lora(self)
+            # huggingface PEFT library freezes the whole LLM, so we need to unfreeze the lm_head if needed
+            if (lm_head := self._lm_head_module) is not None:
+                if self.core_cfg.freeze_llm_head:
+                    freeze_module(lm_head)
+                else:
+                    unfreeze_module(lm_head)
+
+        self._setup_forced_aligner(forced_aligner, data_cfg, val_data_cfg, dataset_cls)
+
+        logging.info("\n" + str(ModelSummary(self, max_depth=2)))
+
+    # ------------------------------------------------------------------
+    # __init__ building blocks
+    # ------------------------------------------------------------------
+    # Kept as separate methods (rather than inlined in ``__init__``) so that
+    # subclasses which build the LLM lazily — e.g.
+    # :class:`~nemo.collections.speechlm2.models.streaming_stt_model_automodel.StreamingSTTModelAutomodel`,
+    # which defers LLM construction to ``configure_model()`` — can reuse them
+    # in a different order.
+
+    def _normalize_chunk_size(self) -> None:
+        """Resolve ``core_cfg.chunk_size`` into ``_chunk_size_candidates`` / ``_chunk_size_repr``.
+
+        A list (e.g. ``[2, 6, 13]``) enables multi chunk-size training (one drawn
+        per batch) and configurable inference (longest by default). A scalar keeps
+        the original single fixed/dynamic/offline mode.
+        """
         cs = self.core_cfg.chunk_size
         if isinstance(cs, (list, tuple, ListConfig)):
             self._chunk_size_candidates = [int(x) for x in cs]
@@ -388,13 +459,17 @@ def __init__(
             self._chunk_size_candidates = None
             self._chunk_size_repr = int(cs)
 
-        # --- LLM ---
-        self.tokenizer = AutoTokenizer(self.core_cfg.pretrained_llm, use_fast=True)
-        self.llm = load_pretrained_hf(
-            self.core_cfg.pretrained_llm,
-            pretrained_weights=self.core_cfg.load_llm_weights,
-        )
+    def _resize_llm_embeddings(self) -> None:
+        """Grow the LLM's embedding table to cover newly added special tokens."""
+        self.llm.resize_token_embeddings(len(self.tokenizer.tokenizer))
+
+    def _register_special_tokens(self) -> None:
+        """Register ``blank`` / ``end_of_audio`` / ``write`` tokens on the tokenizer.
 
+        Each newly added token grows the LLM embedding table via
+        :meth:`_resize_llm_embeddings` (overridden by the Automodel subclass,
+        where the LLM does not exist yet at this point).
+        """
         # Ensure <blank> token is in the vocabulary.
         # Unescape Python escape sequences (e.g. "\\n" → "\n") because Hydra/OmegaConf
         # loads YAML strings literally without interpreting backslash escapes.
@@ -406,7 +481,7 @@ def __init__(
             logging.info("blank_token is empty: blank mechanism disabled")
         elif not token_in_vocab(self.blank_token, self.tokenizer):
             self.tokenizer.add_special_tokens({"additional_special_tokens": [self.blank_token]})
-            self.llm.resize_token_embeddings(len(self.tokenizer.tokenizer))
+            self._resize_llm_embeddings()
             logging.info(f"Added blank token `{self.blank_token}` to tokenizer: {self.blank_token_id}")
         else:
             logging.info(f"Blank token `{str(self.blank_token)}` already in tokenizer: {self.blank_token_id}")
@@ -418,7 +493,7 @@ def __init__(
             eoa = self.core_cfg.end_of_audio_token
             if not token_in_vocab(eoa, self.tokenizer):
                 self.tokenizer.add_special_tokens({"additional_special_tokens": [eoa]})
-                self.llm.resize_token_embeddings(len(self.tokenizer.tokenizer))
+                self._resize_llm_embeddings()
                 eoa_id = self.tokenizer.tokenizer.convert_tokens_to_ids(eoa)
                 logging.info(f"Added end_of_audio_token `{eoa}` to tokenizer: {eoa_id}")
             else:
@@ -433,7 +508,7 @@ def __init__(
             wt = self.core_cfg.write_token
             if not token_in_vocab(wt, self.tokenizer):
                 self.tokenizer.add_special_tokens({"additional_special_tokens": [wt]})
-                self.llm.resize_token_embeddings(len(self.tokenizer.tokenizer))
+                self._resize_llm_embeddings()
                 wt_id = self.tokenizer.tokenizer.convert_tokens_to_ids(wt)
                 logging.info(f"Added write_token `{wt}` to tokenizer: {wt_id}")
             else:
@@ -455,48 +530,8 @@ def __init__(
                     f"{self.core_cfg.write_token!r}."
                 )
 
-        # Separate embedding layer to avoid FSDP/TP conflicts (same pattern as SALM)
-        self.embed_tokens = self.llm.model.embed_tokens
-        del self.llm.model.embed_tokens
-
-        # --- Speech encoder (perception module) ---
-        self.perception = setup_perception(
-            cfg=self.cfg,
-            output_dim=self.llm.config.hidden_size,
-            pretrained_asr=self.core_cfg.pretrained_asr,
-            pretrained_weights=self.core_cfg.load_asr_weights,
-            audio_pad_to=self.core_cfg.audio_pad_to,
-            att_context_size=self.core_cfg.att_context_size,
-        )
-
-        # --- Aux chunk-boundary classifier (only built when enabled) ---
-        # Only valid in dynamic-chunking mode (chunk_size == 0). When disabled,
-        # no module is built and existing runs / checkpoints see no behavior
-        # change. The boundary decision falls back to the LM head's blank /
-        # user_footer_first signal.
-        if self.core_cfg.use_chunk_classifier:
-            assert self.core_cfg.chunk_size == 0, (
-                "use_chunk_classifier=True requires dynamic chunking "
-                f"(chunk_size=0), got chunk_size={self.core_cfg.chunk_size}"
-            )
-            self._build_chunk_classifier()
-            # Aux training/eval reads self._user_footer_first_id (the BCE positive
-            # label). It's normally set lazily by _ensure_inference_cache, but
-            # training runs before any inference call — so prime the cache now.
-            self._ensure_inference_cache()
-
-        self._apply_freeze_config()
-
-        # --- LoRA ---
-        if "lora" in self.cfg:
-            # Install LoRA after freezing the LLM body to avoid freezing the LoRA weights
-            maybe_install_lora(self)
-            # huggingface PEFT library freezes the whole LLM, so we need to unfreeze the lm_head if needed
-            if self.core_cfg.freeze_llm_head:
-                freeze_module(self.llm.lm_head)
-            else:
-                unfreeze_module(self.llm.lm_head)
-
+    def _setup_forced_aligner(self, forced_aligner, data_cfg, val_data_cfg, dataset_cls) -> None:
+        """Attach the optional online forced aligner and its dataset(s)."""
         if forced_aligner is not None:
             assert data_cfg is not None, "Dataset config is required for online forced alignment"
             assert dataset_cls is not None, "Dataset class is required for online forced alignment"
@@ -514,8 +549,6 @@ def __init__(
             self.dataset = None
             self.val_dataset = None
 
-        logging.info("\n" + str(ModelSummary(self, max_depth=2)))
-
     def _build_chunk_classifier(self) -> None:
         """Construct the aux backbone + linear head.
 
@@ -579,10 +612,12 @@ def _apply_freeze_config(self) -> None:
             unfreeze_module(self.llm.model)
 
         # lm_head is inside self.llm, so re-apply after the LLM-wide freeze
-        if self.core_cfg.freeze_llm_head:
-            freeze_module(self.llm.lm_head)
-        else:
-            unfreeze_module(self.llm.lm_head)
+        lm_head = self._lm_head_module
+        if lm_head is not None:
+            if self.core_cfg.freeze_llm_head:
+                freeze_module(lm_head)
+            else:
+                unfreeze_module(lm_head)
 
         # embed_tokens is a separate top-level module (moved out of llm)
         if self.core_cfg.freeze_embed_tokens:
@@ -610,16 +645,9 @@ def text_vocab_size(self):
 
     @property
     def text_pad_id(self) -> int:
-        pad_id = self.tokenizer.pad_id
-        if pad_id is None:
-            pad_id = self.tokenizer.unk_id
-        if pad_id is None:
-            warnings.warn(
-                "The text tokenizer has no <pad> or <unk> token; using id 0 for "
-                "padding (this may lead to silent bugs)."
-            )
-            pad_id = 0
-        return pad_id
+        # Shared with the dataset's collation so the padding value the batch was
+        # built with always matches the one the attention mask is derived from.
+        return resolve_pad_id(self.tokenizer)
 
     @property
     def text_eos_id(self) -> int:
@@ -654,6 +682,48 @@ def blank_token_id(self) -> int:
     def has_blank(self) -> bool:
         return self.blank_token != ""
 
+    # ------------------------------------------------------------------
+    # LLM-backend indirection hooks
+    # ------------------------------------------------------------------
+    # Every training / inference code path below goes through these four hooks
+    # instead of touching ``self.llm`` / ``self.embed_tokens`` directly, so a
+    # subclass can swap the LLM backend without duplicating those paths. See
+    # :class:`~nemo.collections.speechlm2.models.streaming_stt_model_automodel.StreamingSTTModelAutomodel`
+    # for the NeMo Automodel implementations (sharded DTensor embeddings, LLM
+    # outputs that aren't HuggingFace ``ModelOutput`` objects, and an
+    # ``embed_tokens`` that stays inside the LLM).
+
+    def _embed_tokens(self, input_ids: Tensor) -> Tensor:
+        """Embed token IDs with the text embedding table."""
+        return self.embed_tokens(input_ids)
+
+    @property
+    def _lm_head_module(self) -> Optional[nn.Module]:
+        """The LLM's output projection, or ``None`` when the backbone has no ``lm_head``."""
+        return getattr(self.llm, "lm_head", None)
+
+    @property
+    def _embed_ref_tensor(self) -> Tensor:
+        """A tensor carrying the embedding table's dtype/device, for ``Tensor.type_as``."""
+        return self.embed_tokens.weight
+
+    def _llm_forward(self, **kwargs):
+        """Run the LLM forward pass.
+
+        Returns an object supporting both attribute (``out.logits``) and
+        mapping (``out["logits"]``) access, i.e. a HuggingFace ``ModelOutput``.
+        """
+        return self.llm(**kwargs)
+
+    def _move_embedding_ctx(self):
+        """Context manager that makes ``embed_tokens`` reachable from inside the LLM.
+
+        Needed by HuggingFace generation utilities that look the embedding table
+        up on the LLM; here ``embed_tokens`` was moved to the top level in
+        ``__init__``, so it is temporarily put back.
+        """
+        return move_embedding(self)
+
     # ------------------------------------------------------------------
     # Cached auto-detection: content-score mode and annotation availability
     # ------------------------------------------------------------------
@@ -737,7 +807,7 @@ def _build_input_embeds(
         # --- text embeddings ---
         # Zero-out audio positions so embed_tokens gets valid indices.
         text_tokens = input_tokens.where(~audio_mask, torch.zeros_like(input_tokens))
-        text_embeds = self.embed_tokens(text_tokens)  # (B, L, H)
+        text_embeds = self._embed_tokens(text_tokens)  # (B, L, H)
 
         # --- audio embeddings ---
         audio_embs, _audio_emb_lens = self.perception(
@@ -772,7 +842,7 @@ def forward(
         ``hidden_states`` (B, L, H) — the LLM's last-layer hidden state, used
         as input to the aux chunk-boundary classifier.
         """
-        out = self.llm(
+        out = self._llm_forward(
             inputs_embeds=input_embeds,
             attention_mask=attention_mask,
             past_key_values=cache,
@@ -1584,7 +1654,7 @@ def _autoregressive_decode(
                 for b in range(B):
                     state.seq_lens[b] += 1
 
-            token_emb = self.embed_tokens(tokens_to_feed.unsqueeze(1))  # (B, 1, H)
+            token_emb = self._embed_tokens(tokens_to_feed.unsqueeze(1))  # (B, 1, H)
 
             state.attention_mask = torch.cat(
                 [
@@ -1593,7 +1663,7 @@ def _autoregressive_decode(
                 ],
                 dim=1,
             )
-            out = self.llm(
+            out = self._llm_forward(
                 inputs_embeds=token_emb,
                 past_key_values=cache,
                 attention_mask=state.attention_mask,
@@ -1703,9 +1773,9 @@ def get_init_streaming_state(
         # Tokenize each prompt
         all_sys_ids = []
         for prompt in prompts:
-            ids = hf_tok.apply_chat_template(
+            ids = apply_chat_template_ids(
+                hf_tok,
                 [{"role": "system", "content": prompt}],
-                tokenize=True,
                 add_generation_prompt=False,
                 enable_thinking=False,
             )
@@ -1722,11 +1792,11 @@ def get_init_streaming_state(
 
         if not needs_padding:
             # Fast path: all same length, no padding needed
-            sys_embs = self.embed_tokens(
+            sys_embs = self._embed_tokens(
                 torch.tensor(all_sys_ids[0], device=device, dtype=torch.long).unsqueeze(0)
             ).expand(batch_size, -1, -1)
             attention_mask = torch.ones(batch_size, sys_lens[0], dtype=torch.long, device=device)
-            out = self.llm(
+            out = self._llm_forward(
                 inputs_embeds=sys_embs,
                 attention_mask=attention_mask,
                 use_cache=True,
@@ -1741,7 +1811,7 @@ def get_init_streaming_state(
             sys_embs = torch.zeros(batch_size, max_sys_len, H, device=device, dtype=dtype)
             attention_mask = torch.zeros(batch_size, max_sys_len, dtype=torch.long, device=device)
             for b in range(batch_size):
-                embs = self.embed_tokens(
+                embs = self._embed_tokens(
                     torch.tensor(all_sys_ids[b], device=device, dtype=torch.long).unsqueeze(0)
                 ).squeeze(
                     0
@@ -1749,7 +1819,7 @@ def get_init_streaming_state(
                 offset = max_sys_len - sys_lens[b]
                 sys_embs[b, offset:] = embs
                 attention_mask[b, offset:] = 1
-            out = self.llm(
+            out = self._llm_forward(
                 inputs_embeds=sys_embs,
                 attention_mask=attention_mask,
                 use_cache=True,
@@ -1813,7 +1883,7 @@ def _chunked_streaming_step(
         B = state.batch_size
 
         if _audio_embs is not None:
-            audio_chunk_embs = _audio_embs.type_as(self.embed_tokens.weight)
+            audio_chunk_embs = _audio_embs.type_as(self._embed_ref_tensor)
         else:
             # 0. Update audio feature buffer — B frames, one per stream
             if audio_chunk_lens is None:
@@ -1828,7 +1898,7 @@ def _chunked_streaming_step(
             ]
             features, right_paddings = state.audio_feature_buffer.update(frames)
             # Stack B feature buffers → (B, D, fbl)
-            processed_signal = torch.stack(features).type_as(self.embed_tokens.weight)
+            processed_signal = torch.stack(features).type_as(self._embed_ref_tensor)
             processed_signal_length = torch.tensor(
                 [processed_signal.shape[-1] - int(rp) for rp in right_paddings],
                 device=device,
@@ -1867,7 +1937,7 @@ def _chunked_streaming_step(
         audio_mask = turn_ids_t == AUDIO_TOKEN_IDX  # (B, L)
 
         text_tokens = turn_ids_t.where(~audio_mask, torch.zeros_like(turn_ids_t))
-        input_embeds = self.embed_tokens(text_tokens)  # (B, L, H)
+        input_embeds = self._embed_tokens(text_tokens)  # (B, L, H)
 
         # Replace audio placeholder positions with actual audio embeddings
         input_embeds[audio_mask] = audio_chunk_embs.reshape(-1, audio_chunk_embs.shape[-1])
@@ -1878,7 +1948,7 @@ def _chunked_streaming_step(
             [state.attention_mask, torch.ones(B, input_len, dtype=state.attention_mask.dtype, device=device)],
             dim=1,
         )
-        out = self.llm(
+        out = self._llm_forward(
             inputs_embeds=input_embeds,
             past_key_values=state.cache,
             attention_mask=state.attention_mask,
@@ -1903,14 +1973,14 @@ def _chunked_streaming_step(
         any_needs_footer = any(not fc for fc in footer_consumed)
         if any_needs_footer and self._asst_footer_ids:
             flen = len(self._asst_footer_ids)
-            asst_footer_embs = self.embed_tokens(
+            asst_footer_embs = self._embed_tokens(
                 torch.tensor(self._asst_footer_ids, device=device).unsqueeze(0).expand(B, -1)
             )
             state.attention_mask = torch.cat(
                 [state.attention_mask, torch.ones(B, flen, dtype=state.attention_mask.dtype, device=device)],
                 dim=1,
             )
-            out = self.llm(
+            out = self._llm_forward(
                 inputs_embeds=asst_footer_embs,
                 past_key_values=state.cache,
                 attention_mask=state.attention_mask,
@@ -2001,22 +2071,22 @@ def _generate_offline(
         prompts = [system_prompt] * B if isinstance(system_prompt, str) else system_prompt
         all_sys_embs = []
         for prompt in prompts:
-            sys_ids = hf_tok.apply_chat_template(
+            sys_ids = apply_chat_template_ids(
+                hf_tok,
                 [{"role": "system", "content": prompt}],
-                tokenize=True,
                 add_generation_prompt=False,
                 enable_thinking=False,
             )
-            embs = self.embed_tokens(torch.tensor(sys_ids, device=device, dtype=torch.long).unsqueeze(0)).squeeze(
+            embs = self._embed_tokens(torch.tensor(sys_ids, device=device, dtype=torch.long).unsqueeze(0)).squeeze(
                 0
             )  # (L_sys_b, H)
             all_sys_embs.append(embs)
 
         # 2. Embed turn template components (shared across batch)
-        user_header_embs = self.embed_tokens(
+        user_header_embs = self._embed_tokens(
             torch.tensor(self._user_header_ids, device=device, dtype=torch.long).unsqueeze(0)
         )  # (1, L_uh, H)
-        uf_ah_embs = self.embed_tokens(
+        uf_ah_embs = self._embed_tokens(
             torch.tensor(self._user_footer_and_asst_header_ids, device=device, dtype=torch.long).unsqueeze(0)
         )  # (1, L_uf, H)
 
@@ -2026,7 +2096,7 @@ def _generate_offline(
             input_signal=audios,
             input_signal_length=audio_lens_t,
         )  # (B, T_enc_max, H), (B,)
-        batch_audio_embs = batch_audio_embs.type_as(self.embed_tokens.weight)
+        batch_audio_embs = batch_audio_embs.type_as(self._embed_ref_tensor)
         all_audio_embs = [batch_audio_embs[b, : int(batch_emb_lens[b].item())] for b in range(B)]
 
         # 4. Build per-sample input sequences:
@@ -2052,7 +2122,7 @@ def _generate_offline(
             attention_mask[b, offset:] = 1
 
         # 6. LLM prefill (single forward pass)
-        out = self.llm(
+        out = self._llm_forward(
             inputs_embeds=input_embeds,
             attention_mask=attention_mask,
             use_cache=True,
@@ -2236,7 +2306,7 @@ def _generate_dynamic_streaming(
         # Use the blank token embedding when blank is enabled (a real token the
         # model knows). Otherwise fall back to the text pad id.
         pad_token_id = self.blank_token_id if self.has_blank else self.text_pad_id
-        pad_emb = self.embed_tokens(torch.tensor([pad_token_id], device=device)).squeeze(0)  # (H,)
+        pad_emb = self._embed_tokens(torch.tensor([pad_token_id], device=device)).squeeze(0)  # (H,)
 
         for _step in range(max_steps):
             # --- Refill audio embedding buffers for LISTENING streams ---
@@ -2266,7 +2336,7 @@ def _generate_dynamic_streaming(
 
                 # Feature buffer selectively updates only the submitted stream_ids
                 features, right_paddings = state.audio_feature_buffer.update(frames)
-                processed_signal = torch.stack(features).type_as(self.embed_tokens.weight)  # (S, D, T)
+                processed_signal = torch.stack(features).type_as(self._embed_ref_tensor)  # (S, D, T)
                 processed_signal_length = torch.tensor(
                     [processed_signal.shape[-1] - int(rp) for rp in right_paddings],
                     device=device,
@@ -2319,22 +2389,22 @@ def _generate_dynamic_streaming(
                         embs_list.append(pad_emb)
                 elif stream_state[b] == FOOTER:
                     tid = uf_ah_ids[template_pos[b]]
-                    embs_list.append(self.embed_tokens(torch.tensor([tid], device=device)).squeeze(0))  # (H,)
+                    embs_list.append(self._embed_tokens(torch.tensor([tid], device=device)).squeeze(0))  # (H,)
                 elif stream_state[b] == GENERATING:
                     embs_list.append(
-                        self.embed_tokens(torch.tensor([last_gen_token[b]], device=device)).squeeze(0)  # (H,)
+                        self._embed_tokens(torch.tensor([last_gen_token[b]], device=device)).squeeze(0)  # (H,)
                     )
                 elif stream_state[b] == BLANK_FEED:
                     # Only reached when has_blank is True (guarded at transition sites).
                     embs_list.append(
-                        self.embed_tokens(torch.tensor([self.blank_token_id], device=device)).squeeze(0)  # (H,)
+                        self._embed_tokens(torch.tensor([self.blank_token_id], device=device)).squeeze(0)  # (H,)
                     )
                 elif stream_state[b] == ASST_FOOTER:
                     tid = af_ids[template_pos[b]]
-                    embs_list.append(self.embed_tokens(torch.tensor([tid], device=device)).squeeze(0))  # (H,)
+                    embs_list.append(self._embed_tokens(torch.tensor([tid], device=device)).squeeze(0))  # (H,)
                 elif stream_state[b] == HEADER:
                     tid = uh_ids[template_pos[b]]
-                    embs_list.append(self.embed_tokens(torch.tensor([tid], device=device)).squeeze(0))  # (H,)
+                    embs_list.append(self._embed_tokens(torch.tensor([tid], device=device)).squeeze(0))  # (H,)
                 else:  # DONE
                     embs_list.append(pad_emb)
 
@@ -2355,7 +2425,7 @@ def _generate_dynamic_streaming(
                     dim=1,
                 )
                 llm_kwargs["attention_mask"] = state.attention_mask
-            out = self.llm(**llm_kwargs)
+            out = self._llm_forward(**llm_kwargs)
             state.cache = out.past_key_values
             for b in range(B):
                 state.seq_lens[b] += 1
@@ -3032,7 +3102,7 @@ def generate(
         chunk_size = self._resolve_inference_chunk_size(chunk_size_override)
         self._set_encoder_att_context(chunk_size, recompute_streaming=True)
 
-        with move_embedding(self):
+        with self._move_embedding_ctx():
             B = audios.shape[0]
             n_samples_list = [int(audio_lens[b].item()) for b in range(B)]
 
diff --git a/nemo/collections/speechlm2/models/streaming_stt_model_automodel.py b/nemo/collections/speechlm2/models/streaming_stt_model_automodel.py
new file mode 100644
index 000000000000..64e70a0f2b3d
--- /dev/null
+++ b/nemo/collections/speechlm2/models/streaming_stt_model_automodel.py
@@ -0,0 +1,845 @@
+# 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.
+"""Streaming STT SpeechLM with a NeMo Automodel LLM backend.
+
+:class:`StreamingSTTModelAutomodel` is to
+:class:`~nemo.collections.speechlm2.models.streaming_stt_model.StreamingSTTModel`
+what :class:`~nemo.collections.speechlm2.models.salm_automodel.SALMAutomodel` is
+to :class:`~nemo.collections.speechlm2.models.salm.SALM`: the training and
+inference recipe is unchanged, only the LLM backend and the parallelism story
+differ.
+
+The differences relative to the HuggingFace-backed parent class:
+
+* **Deferred construction.** ``__init__`` only builds the tokenizer (and the
+  optional forced aligner); the LLM and the perception module are created in
+  :meth:`configure_model`, which the Lightning strategy calls *after* the device
+  mesh exists. This lets Automodel shard the LLM while loading it, so each rank
+  only ever materializes its own shard.
+* **Parallelism.** FSDP2 / HSDP (and EP for MoE backbones) via
+  :class:`~nemo.collections.speechlm2.parts.parallel.AutomodelParallelStrategy`
+  instead of DDP. The perception module and the optional aux chunk classifier
+  are sharded with ``fully_shard`` on the same mesh Automodel uses for the LLM.
+  Tensor and context parallelism are rejected at ``on_fit_start`` — see
+  :meth:`_validate_parallelism_compatibility`.
+* **Embeddings stay inside the LLM.** The parent moves ``embed_tokens`` to the
+  top level to dodge FSDP/TP conflicts; Automodel needs it in place for its own
+  parallelization plan and state-dict adapters, so here it is accessed through
+  the ``embed_tokens`` property and embedded via ``F.embedding`` on the
+  all-gathered weight (see :meth:`_embed_tokens`).
+* **Vocabulary growth.** ``resize_token_embeddings`` cannot run on a sharded
+  embedding table, so added special tokens (``<blank>`` / ``write`` /
+  ``end_of_audio``) reuse the spare rows most checkpoints already have
+  (``config.vocab_size > len(tokenizer)``); see :meth:`_sync_llm_vocab_size`.
+* **LoRA** is applied by Automodel's own PEFT implementation (config keys
+  ``dim`` / ``alpha`` / ``dropout`` / ``target_modules``), not HuggingFace PEFT.
+
+Everything else — the streaming data recipe, the losses, the fixed / dynamic /
+offline generation paths — is inherited unchanged from ``StreamingSTTModel``.
+"""
+
+from contextlib import nullcontext
+from typing import Any, Optional
+
+import torch
+import torch.nn.functional as F
+from lightning.pytorch.utilities.model_summary import ModelSummary
+from omegaconf import DictConfig, OmegaConf
+from torch import Tensor
+from torch.distributed.fsdp import fully_shard
+from torch.distributed.tensor import DTensor, Partial
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers.utils import ModelOutput
+
+from nemo.collections.common.tokenizers import AutoTokenizer
+from nemo.collections.speechlm2.data.streaming_stt_dataset import StreamingSTTDataset
+from nemo.collections.speechlm2.models.streaming_stt_model import StreamingSTTModel, StreamingSTTModelConfig
+from nemo.collections.speechlm2.parts.automodel_lora import (
+    LORA_PARAM_PATTERN,
+    ensure_lora_trainable,
+    make_peft_config,
+    maybe_install_lora,
+)
+from nemo.collections.speechlm2.parts.pretrained import (
+    load_pretrained_automodel_llm,
+    maybe_load_pretrained_models,
+    setup_perception,
+)
+from nemo.collections.speechlm2.parts.utils import freeze_module, to_dataclass, unfreeze_module
+from nemo.utils import logging
+
+# Fields of ``transformers.modeling_outputs.CausalLMOutputWithPast``; used to
+# normalize LLM outputs that come back as a plain dict.
+_CAUSAL_LM_OUTPUT_FIELDS = ("loss", "logits", "past_key_values", "hidden_states", "attentions")
+
+# HuggingFace-PEFT LoRA keys. Automodel's ``PeftConfig`` uses different names
+# (``dim`` / ``alpha`` / ``dropout``), so a config copy-pasted from the HF-backed
+# recipe would fail deep inside Automodel with an unhelpful ``TypeError``.
+_HF_PEFT_LORA_KEYS = ("r", "lora_alpha", "lora_dropout", "task_type", "peft_type", "modules_to_save")
+
+
+class StreamingSTTModelAutomodel(StreamingSTTModel):
+    """Streaming STT SpeechLM whose LLM is loaded and parallelized by NeMo Automodel.
+
+    Args:
+        cfg: Model configuration as a plain Python dict (required for
+            hyperparameter serialization in PTL checkpoints). Accepts every key
+            of :class:`~nemo.collections.speechlm2.models.streaming_stt_model.StreamingSTTModelConfig`
+            plus the Automodel-specific ones: ``lora`` (Automodel ``PeftConfig``
+            format), ``automodel_backend``, ``sdpa_method``, ``compile``,
+            ``aux_loss_coeff``, ``train_gate``, ``moe_metrics``,
+            ``trust_remote_code``, ``init_from_checkpoint`` and
+            ``init_configure_model``.
+        forced_aligner: Optional forced aligner for online alignment.
+        data_cfg: Training dataset config (required with ``forced_aligner``).
+        val_data_cfg: Optional validation dataset config; falls back to ``data_cfg``.
+        dataset_cls: Dataset class used when online forced alignment is enabled.
+    """
+
+    def __init__(
+        self,
+        cfg: dict,
+        forced_aligner=None,
+        data_cfg: Optional[DictConfig] = None,
+        val_data_cfg: Optional[DictConfig] = None,
+        dataset_cls=StreamingSTTDataset,
+    ) -> None:
+        assert isinstance(cfg, dict), (
+            "You must pass the config to StreamingSTTModelAutomodel as a Python dict to support hyperparameter "
+            f"serialization in PTL checkpoints (we got: '{type(cfg)=}')."
+        )
+        # NOTE: StreamingSTTModel.__init__ is deliberately bypassed — it builds
+        # the LLM and the perception module eagerly, whereas the Automodel path
+        # must defer both to configure_model() so they can be created directly
+        # on the device mesh. The pieces of the parent __init__ that do not
+        # depend on the LLM are reused through the extracted helpers below.
+        # ``super(StreamingSTTModel, self)`` (rather than a direct
+        # ``LightningModule.__init__(self)`` call) is deliberate: referencing
+        # ``super`` gives this frame the ``__class__`` cell that Lightning's
+        # ``save_hyperparameters()`` needs to discover the init args — without
+        # it ``self.hparams`` comes out empty and checkpoints lose ``cfg``.
+        super(StreamingSTTModel, self).__init__()
+        self.save_hyperparameters()
+        self.cfg = DictConfig(cfg)
+        self.core_cfg: StreamingSTTModelConfig = to_dataclass(StreamingSTTModelConfig, cfg)
+        self._normalize_chunk_size()
+
+        self.tokenizer = AutoTokenizer(self.core_cfg.pretrained_llm, use_fast=True)
+        self.llm = None  # populated by configure_model()
+        self.perception = None  # populated by configure_model()
+
+        self._use_fsdp = False
+        self._configured_with_mesh = False
+        self._device_mesh = None
+        self._current_batch_idx = 0
+
+        self._register_special_tokens()
+        self._validate_lora_cfg()
+        self._setup_forced_aligner(forced_aligner, data_cfg, val_data_cfg, dataset_cls)
+
+        if self.cfg.get("init_configure_model", False):
+            self.configure_model()
+
+    # ------------------------------------------------------------------
+    # Config validation
+    # ------------------------------------------------------------------
+
+    def _validate_lora_cfg(self) -> None:
+        """Fail fast on a HuggingFace-PEFT-style ``lora`` block.
+
+        The parent class installs LoRA through HuggingFace PEFT
+        (``r`` / ``lora_alpha`` / ``lora_dropout``); Automodel has its own
+        implementation with different key names. Silently forwarding the HF keys
+        raises a ``TypeError`` several frames deep inside Automodel, so translate
+        it into an actionable message here.
+        """
+        if "lora" not in self.cfg or not self.cfg.lora:
+            return
+        offending = [k for k in _HF_PEFT_LORA_KEYS if k in self.cfg.lora]
+        if offending:
+            raise ValueError(
+                f"model.lora contains HuggingFace-PEFT keys {offending}, but "
+                f"{type(self).__name__} installs LoRA with NeMo Automodel's PEFT implementation. "
+                "Use the Automodel key names instead, e.g.:\n"
+                "  lora:\n"
+                "    dim: 128           # was `r`\n"
+                "    alpha: 256         # was `lora_alpha`\n"
+                "    dropout: 0.01      # was `lora_dropout`\n"
+                "    target_modules: [\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\"]"
+            )
+
+    def _unfreeze_lora_params(self) -> int:
+        """Re-enable gradients on the LoRA adapters after the LLM-body freeze.
+
+        On the ``device_mesh`` path Automodel installs LoRA *inside*
+        ``from_pretrained``, i.e. before :meth:`_apply_freeze_config` runs, so
+        ``freeze_module(self.llm.model)`` also freezes ``lora_A`` / ``lora_B``.
+        ``ensure_lora_trainable`` only appends the ``prevent_freeze_params``
+        regex, and ``optim_setup.freeze_and_subset`` skips parameters that
+        already have ``requires_grad=False`` *before* consulting that regex — so
+        without this the adapters would silently never train (the only hint being
+        an "UNMATCHED freeze-preventing pattern" warning).
+
+        Returns the number of adapter parameters re-enabled.
+        """
+        import re
+
+        pattern = re.compile(LORA_PARAM_PATTERN)
+        n = 0
+        for name, param in self.llm.named_parameters():
+            if pattern.match(name) and not param.requires_grad:
+                param.requires_grad_(True)
+                n += 1
+        if n:
+            logging.info(f"Re-enabled gradients on {n} LoRA parameters frozen by the LLM-body freeze.")
+        return n
+
+    # ------------------------------------------------------------------
+    # LLM-backend indirection hooks (see StreamingSTTModel for the contract)
+    # ------------------------------------------------------------------
+
+    @property
+    def device(self) -> torch.device:
+        """Infer the device from the LLM's parameters.
+
+        ``LightningModule.device`` is set by the Trainer and defaults to CPU
+        during standalone inference (no Trainer). Query the actual parameter
+        storage instead so ``.to(self.device)`` works for both regular and
+        ``DTensor`` (FSDP2 / TP) parameters.
+        """
+        if self.llm is not None:
+            p = next(self.llm.parameters(), None)
+            if p is not None:
+                return p._local_tensor.device if isinstance(p, DTensor) else p.device
+        return super().device
+
+    @property
+    def embed_tokens(self):
+        """The LLM's embedding layer (kept inside the LLM, unlike the parent class)."""
+        if self.llm is None:
+            return None
+        return self.llm.model.embed_tokens
+
+    def _embed_tokens(self, input_ids: Tensor) -> Tensor:
+        """Embed token IDs using the LLM's embedding table.
+
+        Uses ``F.embedding`` instead of calling the ``nn.Embedding`` module to
+        avoid triggering FSDP2 pre-forward hooks (which lazily initialize the
+        child before the root LLM module, causing a ``RuntimeError``).
+
+        When the weight is a sharded ``DTensor`` (FSDP2) it is ``full_tensor()``-ed
+        first to all-gather the complete embedding table — the same operation
+        FSDP2 performs inside the LLM's forward pass.
+
+        ``grad_placements`` matters whenever the embedding table is trainable
+        (``freeze_embed_tokens: false``, which the streaming recipe needs so the
+        new ``<blank>`` token can be learned). The lookup happens *outside* the
+        LLM's FSDP2 forward, so FSDP2 never reduce-scatters this parameter's
+        gradient. Without declaring the gradient ``Partial``, ``full_tensor()``
+        defaults to a ``Replicate`` backward and the redistribute back to
+        ``Shard(0)`` is a pure local chunk — each rank would keep only its own
+        micro-batch contribution. ``Partial(reduce_op="avg")`` makes that
+        redistribute an averaging reduce-scatter, matching what FSDP2 does for
+        every other parameter.
+        """
+        weight = self.embed_tokens.weight
+        if isinstance(weight, DTensor):
+            if weight.requires_grad:
+                grad_placements = [Partial(reduce_op="avg")] * weight.device_mesh.ndim
+                weight = weight.full_tensor(grad_placements=grad_placements)
+            else:
+                weight = weight.full_tensor()
+        return F.embedding(input_ids, weight)
+
+    @property
+    def _embed_ref_tensor(self) -> Tensor:
+        """Dtype/device reference for ``Tensor.type_as``.
+
+        ``type_as`` on a ``DTensor`` would try to convert to the DTensor *type*;
+        the local shard carries the same dtype and device, so use it instead.
+        """
+        weight = self.embed_tokens.weight
+        return weight.to_local() if isinstance(weight, DTensor) else weight
+
+    def _llm_forward(self, **kwargs):
+        """Run the LLM and normalize the output into a HuggingFace ``ModelOutput``.
+
+        Automodel's own model implementations (e.g. the MoE backbones) do not
+        always honor ``return_dict=True`` — some return the logits tensor
+        directly. Wrap those so the inherited training and streaming-inference
+        code can keep using ``out.logits`` / ``out.past_key_values``.
+        """
+        out = self.llm(**kwargs)
+
+        if isinstance(out, ModelOutput):
+            normalized = out
+        elif isinstance(out, torch.Tensor):
+            normalized = CausalLMOutputWithPast(logits=out)
+        elif isinstance(out, dict):
+            normalized = CausalLMOutputWithPast(**{k: v for k, v in out.items() if k in _CAUSAL_LM_OUTPUT_FIELDS})
+        elif isinstance(out, (tuple, list)) and out:
+            normalized = CausalLMOutputWithPast(logits=out[0])
+        else:
+            raise RuntimeError(f"Unsupported LLM output type: {type(out)!r}; expected a ModelOutput, dict or Tensor.")
+
+        if kwargs.get("use_cache", False) and normalized.past_key_values is None:
+            raise RuntimeError(
+                f"The Automodel LLM {type(self.llm).__name__!r} did not return a KV cache "
+                "(`past_key_values`) with `use_cache=True`. Streaming inference requires "
+                "incremental decoding support from the backbone; use the HuggingFace-backed "
+                "StreamingSTTModel for inference with this LLM."
+            )
+        if kwargs.get("output_hidden_states", False) and normalized.hidden_states is None:
+            raise RuntimeError(
+                f"The Automodel LLM {type(self.llm).__name__!r} did not return hidden states with "
+                "`output_hidden_states=True`, which the aux chunk-boundary classifier "
+                "(`use_chunk_classifier=true`) needs."
+            )
+        return normalized
+
+    def _move_embedding_ctx(self):
+        """No-op: ``embed_tokens`` already lives inside the LLM here."""
+        return nullcontext()
+
+    # ------------------------------------------------------------------
+    # Vocabulary sizing
+    # ------------------------------------------------------------------
+
+    def _resize_llm_embeddings(self) -> None:
+        """Defer the embedding resize — the LLM does not exist during ``__init__``."""
+        if self.llm is None:
+            return
+        self._sync_llm_vocab_size()
+
+    def _sync_llm_vocab_size(self) -> None:
+        """Ensure the LLM embedding table covers every tokenizer ID.
+
+        Unlike the parent class this never *shrinks* the table: most checkpoints
+        ship with spare rows (``config.vocab_size > len(tokenizer)``, e.g. 151936
+        vs 151669 for Qwen3), which the newly added special tokens can occupy
+        without touching the parameter shapes. That matters because a sharded
+        (``DTensor``) embedding cannot be resized in place at all — growing one
+        raises with an actionable message instead of failing obscurely later.
+        """
+        target = len(self.tokenizer.tokenizer)
+        embed = self.embed_tokens
+        current = int(embed.weight.shape[0])
+        if current >= target:
+            if current > target:
+                logging.info(
+                    f"LLM embedding table has {current} rows for a tokenizer of {target} tokens "
+                    f"({current - target} spare rows) — added special tokens reuse the spare rows, no resize needed."
+                )
+            return
+
+        if isinstance(embed.weight, DTensor):
+            raise RuntimeError(
+                f"The tokenizer has {target} tokens but the sharded LLM embedding table only has {current} rows, "
+                "and a DTensor-sharded embedding cannot be resized. Extend the checkpoint's vocabulary offline "
+                "(so that config.vocab_size >= len(tokenizer)), or reduce the number of added special tokens "
+                "(`blank_token` / `write_token` / `end_of_audio_token`)."
+            )
+        self.llm.resize_token_embeddings(target)
+        logging.info(f"Resized the LLM embedding table from {current} to {target} rows.")
+
+    # ------------------------------------------------------------------
+    # Aux chunk-boundary classifier
+    # ------------------------------------------------------------------
+
+    def _build_chunk_classifier(self) -> None:
+        """Build the aux backbone + head, skipping the warm-start when the LLM is sharded.
+
+        The parent's warm start copies the last K LLM layers into the aux
+        backbone. Under FSDP2 those weights are ``DTensor`` shards, which cannot
+        be loaded into a plain (unsharded) module, so fall back to random init.
+        """
+        llm_param = next(self.llm.parameters(), None)
+        if isinstance(llm_param, DTensor) and self.core_cfg.chunk_classifier_init_from_llm:
+            logging.warning(
+                "chunk_classifier_init_from_llm=True is not supported when the LLM is sharded "
+                "(DTensor parameters); falling back to random init for the aux chunk classifier."
+            )
+            init_from_llm = self.core_cfg.chunk_classifier_init_from_llm
+            self.core_cfg.chunk_classifier_init_from_llm = False
+            try:
+                super()._build_chunk_classifier()
+            finally:
+                self.core_cfg.chunk_classifier_init_from_llm = init_from_llm
+            return
+        super()._build_chunk_classifier()
+
+    # ------------------------------------------------------------------
+    # Model construction (called by the Lightning strategy with the device mesh)
+    # ------------------------------------------------------------------
+
+    def configure_model(
+        self,
+        device_mesh=None,
+        distributed_config=None,
+        moe_config=None,
+        moe_mesh=None,
+        activation_checkpointing_llm: Optional[bool] = None,
+        activation_checkpointing_perception: Optional[bool] = None,
+    ) -> None:
+        """Build (and parallelize) the LLM, the perception module and the aux head.
+
+        Called by Lightning after ``setup_environment()``, so ``device_mesh`` is
+        available and Automodel can shard the LLM while loading its weights.
+        Safe to call twice — the second call is a no-op.
+        """
+        resolved_mesh = device_mesh if device_mesh is not None else self.device_mesh
+        if self.llm is not None:
+            # Already built. Refuse the case where the first build had no mesh
+            # (``init_configure_model=true``, or ``from_pretrained`` — see
+            # parts/hf_hub.py) but one exists now: returning silently would leave
+            # every rank with a full unsharded replica and, since
+            # ModelParallelStrategy installs no DDP wrapper, no gradient
+            # synchronization at all.
+            if resolved_mesh is not None and not self._configured_with_mesh:
+                raise RuntimeError(
+                    "configure_model() already ran without a device mesh (model.init_configure_model=true "
+                    "or a from_pretrained() restore), but a device mesh is available now. The existing model "
+                    "would train unsharded and WITHOUT gradient synchronization. Remove `init_configure_model` "
+                    "from the model config for distributed training — the strategy calls configure_model() "
+                    "with the device mesh itself."
+                )
+            return
+
+        # Use the provided device_mesh, or fall back to the LightningModule property.
+        device_mesh = resolved_mesh
+        if device_mesh is not None:
+            self._device_mesh = device_mesh
+
+        dtype = self._resolve_dtype()
+
+        # Fall back to trainer.strategy for configs (Lightning training path).
+        if distributed_config is None and self._trainer is not None:
+            distributed_config = getattr(self._trainer.strategy, "distributed_config", None)
+        if moe_mesh is None and self._trainer is not None:
+            moe_mesh = getattr(self._trainer.strategy, "moe_mesh", None)
+        if moe_config is None and self._trainer is not None:
+            moe_config = getattr(self._trainer.strategy, "moe_config", None)
+        if activation_checkpointing_llm is None and self._trainer is not None:
+            activation_checkpointing_llm = getattr(self._trainer.strategy, "activation_checkpointing_llm", None)
+        if activation_checkpointing_llm is None:
+            activation_checkpointing_llm = False
+        if activation_checkpointing_perception is None and self._trainer is not None:
+            activation_checkpointing_perception = getattr(
+                self._trainer.strategy, "activation_checkpointing_perception", None
+            )
+        if activation_checkpointing_perception is None:
+            activation_checkpointing_perception = False
+
+        automodel_kwargs = self._build_automodel_kwargs(
+            device_mesh=device_mesh,
+            distributed_config=distributed_config,
+            moe_config=moe_config,
+            moe_mesh=moe_mesh,
+            activation_checkpointing_llm=activation_checkpointing_llm,
+        )
+
+        # When LoRA is configured and we have a device_mesh, pass peft_config
+        # through Automodel so LoRA is applied before FSDP2 sharding (which
+        # handles meta-device init correctly).
+        peft_config = make_peft_config(self.cfg.lora) if "lora" in self.cfg else None
+        if peft_config is not None and device_mesh is not None:
+            automodel_kwargs["peft_config"] = peft_config
+
+        # --- LLM ---
+        self.llm = load_pretrained_automodel_llm(
+            self.core_cfg.pretrained_llm,
+            pretrained_weights=self.core_cfg.load_llm_weights,
+            dtype=dtype,
+            trust_remote_code=self.cfg.get("trust_remote_code", False),
+            **automodel_kwargs,
+        )
+        # Special tokens were registered on the tokenizer in __init__, when the
+        # LLM did not exist yet — reconcile the vocabulary now.
+        self._sync_llm_vocab_size()
+
+        # --- Speech encoder (perception module) ---
+        self.perception = setup_perception(
+            cfg=self.cfg,
+            output_dim=self.llm.config.hidden_size,
+            pretrained_asr=self.core_cfg.pretrained_asr,
+            pretrained_weights=self.core_cfg.load_asr_weights,
+            audio_pad_to=self.core_cfg.audio_pad_to,
+            att_context_size=self.core_cfg.att_context_size,
+        )
+
+        # --- Aux chunk-boundary classifier (only built when enabled) ---
+        if self.core_cfg.use_chunk_classifier:
+            assert self.core_cfg.chunk_size == 0, (
+                "use_chunk_classifier=True requires dynamic chunking "
+                f"(chunk_size=0), got chunk_size={self.core_cfg.chunk_size}"
+            )
+            self._build_chunk_classifier()
+            # Aux training/eval reads self._user_footer_first_id (the BCE positive
+            # label). It's normally set lazily by _ensure_inference_cache, but
+            # training runs before any inference call — so prime the cache now.
+            self._ensure_inference_cache()
+
+        # Activation checkpointing on the perception encoder layers. Must run
+        # BEFORE FSDP2 wrapping (same as the LLM path inside Automodel) so that
+        # checkpoint_wrapper sees the pristine layer objects.
+        self.perception.set_activation_checkpointing(activation_checkpointing_perception)
+
+        # Module-level freezing FIRST — everything that marks parameters trainable
+        # (LoRA adapters, MoE gates) must run *after* it, because
+        # ``freeze_module(self.llm.model)`` recurses over the whole LLM body and
+        # ``optim_setup.freeze_and_subset`` skips ``requires_grad=False`` params
+        # before it ever consults ``prevent_freeze_params``.
+        self._apply_freeze_config()
+
+        # --- LoRA ---
+        # With a device_mesh, LoRA was already applied inside Automodel's
+        # from_pretrained (before sharding); otherwise apply it now.
+        if peft_config is not None:
+            if device_mesh is None:
+                maybe_install_lora(self)
+            else:
+                # Still need the prevent_freeze_params pattern for configure_optimizers.
+                ensure_lora_trainable(self)
+            # In the device_mesh path the adapters existed before the freeze above
+            # and were frozen along with the LLM body; re-enable them.
+            self._unfreeze_lora_params()
+            # Automodel's LoRA does not freeze the base model, but keep the
+            # explicit lm_head switch honored (parity with the parent class).
+            if (lm_head := self._lm_head_module) is not None:
+                if self.core_cfg.freeze_llm_head:
+                    freeze_module(lm_head)
+                else:
+                    unfreeze_module(lm_head)
+
+        # MoE options last: ``train_gate`` unfreezes router weights that live
+        # inside ``llm.model`` and would be re-frozen by ``_apply_freeze_config``.
+        self.setup_moe_options()
+
+        if device_mesh is None:
+            maybe_load_pretrained_models(self)
+            logging.info("\n" + str(ModelSummary(self, max_depth=2)))
+            return
+
+        # Cast the perception module to the training dtype BEFORE FSDP2 wrapping.
+        # The LLM is already in the target dtype (loaded with torch_dtype=dtype)
+        # and FSDP2 requires a uniform parameter dtype.
+        if dtype != torch.float32:
+            self.perception.to(dtype=dtype)
+
+        self._configured_with_mesh = True
+
+        # Use the same FSDP mesh Automodel uses for the LLM so that gradient
+        # clipping can torch.stack norms from all parameters.
+        dim_names = device_mesh.mesh_dim_names
+        if "dp_replicate" in dim_names and "dp_shard_cp" in dim_names:
+            fsdp_mesh = device_mesh["dp_replicate", "dp_shard_cp"]
+        elif "dp_shard_cp" in dim_names:
+            fsdp_mesh = device_mesh["dp_shard_cp"]
+        else:
+            fsdp_mesh = device_mesh["dp"]
+
+        if fsdp_mesh.size() > 1:
+            self._use_fsdp = True
+            self.perception = fully_shard(self.perception, mesh=fsdp_mesh)
+            # Every parameter must belong to some FSDP2 group, otherwise its
+            # gradients are never reduced across ranks and the replicas diverge.
+            if self.core_cfg.use_chunk_classifier:
+                self.chunk_classifier_backbone = fully_shard(self.chunk_classifier_backbone, mesh=fsdp_mesh)
+                self.chunk_classifier_head = fully_shard(self.chunk_classifier_head, mesh=fsdp_mesh)
+
+        # Optionally initialize weights from a previous training checkpoint
+        # (fresh optimizer/scheduler). Must happen after FSDP wrapping so that
+        # DCP loading can fill DTensor parameters with the correct shards.
+        maybe_load_pretrained_models(self)
+        logging.info("\n" + str(ModelSummary(self, max_depth=2)))
+
+    def _resolve_dtype(self) -> torch.dtype:
+        """Training dtype: trainer precision first, then ``model.dtype`` from the config."""
+        if self._trainer is not None:
+            precision = str(self._trainer.precision)
+            if "bf16" in precision:
+                return torch.bfloat16
+            if "16" in precision:
+                return torch.float16
+            return torch.float32
+        cfg_dtype = self.cfg.get("torch_dtype", None) or getattr(self.core_cfg, "dtype", None)
+        if cfg_dtype is None:
+            return torch.float32
+        return getattr(torch, cfg_dtype) if isinstance(cfg_dtype, str) else cfg_dtype
+
+    def _build_automodel_kwargs(
+        self,
+        device_mesh,
+        distributed_config,
+        moe_config,
+        moe_mesh,
+        activation_checkpointing_llm: bool,
+    ) -> dict:
+        """Assemble the kwargs forwarded to Automodel's ``from_pretrained`` / ``from_config``."""
+        automodel_kwargs: dict[str, Any] = {}
+        if device_mesh is not None:
+            automodel_kwargs["device_mesh"] = device_mesh
+            # Automodel's instantiate_infrastructure unconditionally calls
+            # .to_dict() on these configs, so always provide defaults.
+            if distributed_config is None:
+                from nemo_automodel.components.distributed.config import FSDP2Config
+
+                distributed_config = FSDP2Config()
+            if moe_config is None:
+                from nemo_automodel.components.moe.config import MoEParallelizerConfig
+
+                moe_config = MoEParallelizerConfig()
+            # Route the single LLM AC flag to both paths: the EP/MoE parallelizer
+            # reads ``activation_checkpointing`` directly (MoEParallelizerConfig
+            # has no such field), while FSDP2's AC wrapping reads the field on
+            # FSDP2Config.
+            if activation_checkpointing_llm:
+                distributed_config.activation_checkpointing = True
+            automodel_kwargs["distributed_config"] = distributed_config
+            automodel_kwargs["moe_config"] = moe_config
+            automodel_kwargs["activation_checkpointing"] = activation_checkpointing_llm
+        if moe_mesh is not None:
+            automodel_kwargs["moe_mesh"] = moe_mesh
+
+        # torch.compile support.
+        compile_cfg = self.cfg.get("compile", None)
+        if compile_cfg is not None:
+            from nemo_automodel.components.utils.compile_utils import CompileConfig
+
+            automodel_kwargs["compile_config"] = CompileConfig(**dict(compile_cfg))
+
+        # Backend dispatch — lets YAML pick attn/linear/rms_norm/MoE backends
+        # (e.g. attn=sdpa to bypass TransformerEngine).
+        backend_cfg = self.cfg.get("automodel_backend", None)
+        if backend_cfg is not None:
+            from nemo_automodel.components.models.common import BackendConfig
+
+            automodel_kwargs["backend"] = BackendConfig(**OmegaConf.to_container(backend_cfg, resolve=True))
+
+        # Pin the SDPA kernel used by attn=sdpa (e.g. ["flash_attention"]).
+        sdpa_method = self.cfg.get("sdpa_method", None)
+        if sdpa_method is not None:
+            automodel_kwargs["sdpa_method"] = list(OmegaConf.to_container(sdpa_method, resolve=True))
+
+        return automodel_kwargs
+
+    # ------------------------------------------------------------------
+    # Training loop integration
+    # ------------------------------------------------------------------
+
+    def on_fit_start(self) -> None:
+        """Validate the parallelism configuration and configure the MoE aux-loss scaler."""
+        self._validate_parallelism_compatibility()
+        self._configure_moe_aux_loss_scaler()
+
+    def _validate_parallelism_compatibility(self) -> None:
+        """Reject parallelism dimensions the streaming recipe cannot support.
+
+        * **Context parallelism** splits the sequence across ranks, which is
+          incompatible with the interleaved audio/text layout and the causal
+          chunk-by-chunk supervision this model is trained with.
+        * **Tensor parallelism** makes the LLM's ``lm_head`` emit vocab-sharded
+          ``DTensor`` logits, which the inherited loss cannot consume: the
+          blank / non-blank breakdown indexes the per-token loss with a plain
+          boolean mask (``per_token_loss[is_blank]``), and mixing a
+          ``torch.Tensor`` mask with a ``DTensor`` raises at the first step.
+          Supporting TP also requires padding the interleaved sequence to a
+          multiple of ``tp_size`` (see ``SALMAutomodel.prepare_inputs``).
+
+        Only FSDP2 (``dp_size`` / ``dp_replicate_size``) and expert parallelism
+        (``ep_size``, MoE backbones) are supported.
+        """
+        device_mesh = getattr(self, "_device_mesh", None)
+        if device_mesh is None:
+            return
+        names = device_mesh.mesh_dim_names or ()
+        if "cp" in names and device_mesh["cp"].size() > 1:
+            raise ValueError(
+                f"{type(self).__name__} does not support context parallelism "
+                f"(got cp_size={device_mesh['cp'].size()}). The interleaved audio/text sequence must stay "
+                "on a single rank; use FSDP2 (dp_size) instead."
+            )
+        if "tp" in names and device_mesh["tp"].size() > 1:
+            raise ValueError(
+                f"{type(self).__name__} does not support tensor parallelism "
+                f"(got tp_size={device_mesh['tp'].size()}). TP shards the lm_head output, and the streaming "
+                "loss (blank / non-blank breakdown) indexes the per-token loss with a plain boolean mask, "
+                "which is not DTensor-safe. Use FSDP2 (dp_size) instead."
+            )
+
+    def training_step(self, batch, batch_idx: int):
+        # Recorded for _setup_moe_fsdp_sync(), which runs during backward().
+        self._current_batch_idx = batch_idx
+        ans = super().training_step(batch, batch_idx)
+        self.maybe_log_moe_metrics(batch_idx)
+        return ans
+
+    def backward(self, *args, **kwargs):
+        self._setup_moe_fsdp_sync()
+        super().backward(*args, **kwargs)
+
+    def _setup_moe_fsdp_sync(self):
+        """Configure MoE FSDP gradient sync for gradient accumulation.
+
+        When ``accumulate_grad_batches > 1``, disables gradient all-reduce and
+        resharding on intermediate backward passes and re-enables them on the
+        final backward before ``optimizer.step()``. No-op when the LLM lacks
+        Automodel's ``MoEFSDPSyncMixin`` or gradient accumulation is inactive.
+        """
+        if not self._use_fsdp or not hasattr(self.llm, 'prepare_for_grad_accumulation'):
+            return
+        acc = self.trainer.accumulate_grad_batches if self._trainer else 1
+        if acc <= 1:
+            return
+        batch_idx = getattr(self, '_current_batch_idx', 0)
+        is_final = (batch_idx + 1) % acc == 0 or (batch_idx + 1) == self.trainer.num_training_batches
+        if is_final:
+            self.llm.prepare_for_final_backward()
+        else:
+            self.llm.prepare_for_grad_accumulation()
+
+    def configure_gradient_clipping(self, optimizer, gradient_clip_val, gradient_clip_algorithm=None):
+        """Mesh-aware gradient clipping.
+
+        When Automodel parallelizes the LLM, some parameters end up as DTensors
+        on the ``(dp_replicate, dp_shard_cp)`` mesh while others may be on the
+        flattened ``dp`` mesh. PyTorch's ``clip_grad_norm_`` requires all norms
+        to share the same mesh for ``torch.stack``, so delegate to Automodel's
+        implementation which groups parameters by ``(mesh_id, placements)``.
+        """
+        if not self._use_fsdp or gradient_clip_val is None or gradient_clip_val <= 0:
+            return super().configure_gradient_clipping(optimizer, gradient_clip_val, gradient_clip_algorithm)
+        from nemo_automodel.components.training.utils import _clip_grad_norm_impl
+
+        params = [p for group in optimizer.param_groups for p in group["params"] if p.grad is not None]
+        if params:
+            _clip_grad_norm_impl(params, max_norm=gradient_clip_val)
+
+    # ------------------------------------------------------------------
+    # MoE helpers (no-ops for dense LLMs)
+    # ------------------------------------------------------------------
+
+    def setup_moe_options(self):
+        """Apply MoE config overrides and enable load-balance tracking.
+
+        Must be called after ``self.llm`` is created. Safe no-op when the LLM
+        has no ``Gate`` modules (dense backbone) or when Automodel's MoE
+        components are unavailable.
+        """
+        aux_loss_coeff = self.cfg.get("aux_loss_coeff", 0.0)
+        moe_metrics_cfg = self.cfg.get("moe_metrics", None)
+        moe_requested = (
+            aux_loss_coeff > 0
+            or self.cfg.get("train_gate", False)
+            or (moe_metrics_cfg is not None and moe_metrics_cfg.get("enabled", False))
+        )
+        try:
+            from nemo_automodel.components.moe.layers import Gate
+        except ImportError:
+            if moe_requested:
+                logging.warning(
+                    "MoE options (aux_loss_coeff / train_gate / moe_metrics) were requested but "
+                    "nemo_automodel's MoE components are not importable — ignoring them."
+                )
+            return
+
+        if aux_loss_coeff > 0:
+            for module in self.llm.modules():
+                if isinstance(module, Gate):
+                    module.aux_loss_coeff = aux_loss_coeff
+
+        if self.cfg.get("train_gate", False):
+            for module in self.llm.modules():
+                if isinstance(module, Gate):
+                    module.train_gate = True
+                    module.weight.requires_grad_(True)
+                    if module.bias is not None:
+                        module.bias.requires_grad_(True)
+
+        if moe_metrics_cfg is not None and moe_metrics_cfg.get("enabled", False):
+            from nemo_automodel.components.moe.load_balance_metrics import enable_load_balance_tracking
+
+            enable_load_balance_tracking(self.llm)
+
+    def maybe_log_moe_metrics(self, step: int):
+        """Collect and log MoE load-balance metrics.
+
+        All ranks must call this method (the all-reduce inside
+        ``collect_expert_loads`` is collective).
+        """
+        moe_metrics_cfg = self.cfg.get("moe_metrics", None)
+        if moe_metrics_cfg is None or not moe_metrics_cfg.get("enabled", False):
+            return
+
+        from nemo_automodel.components.moe.load_balance_metrics import (
+            collect_expert_loads,
+            compute_brief_metrics,
+            compute_detailed_metrics,
+        )
+
+        layer_loads = collect_expert_loads(self.llm, dp_group=self._get_moe_dp_group())
+        if not layer_loads:
+            return
+
+        top_k = moe_metrics_cfg.get("top_k_experts", 5)
+        if moe_metrics_cfg.get("mode", "brief") == "detailed":
+            detailed_every = moe_metrics_cfg.get("detailed_every_steps", None)
+            if detailed_every is not None and step % detailed_every != 0:
+                metrics = compute_brief_metrics(layer_loads, top_k=top_k)
+            else:
+                metrics = compute_detailed_metrics(layer_loads, top_k=top_k)
+        else:
+            metrics = compute_brief_metrics(layer_loads, top_k=top_k)
+
+        self.log_dict(metrics, on_step=True)
+
+    def _get_moe_dp_group(self):
+        """Return the DP process group for MoE metrics all-reduce.
+
+        Prefers the ``dp_cp`` submesh (includes context parallelism) and falls
+        back to ``dp``. Both are flattened submeshes registered in
+        ``device_mesh._flatten_mapping``, so they are resolved via
+        ``get_flat_mesh``. Returns ``None`` when no device mesh is available.
+        """
+        device_mesh = getattr(self, "_device_mesh", None)
+        if device_mesh is None:
+            return None
+        try:
+            from nemo_automodel.components.distributed.mesh_utils import get_flat_mesh
+        except ImportError:
+            return None
+
+        try:
+            if "cp" in device_mesh.mesh_dim_names and device_mesh["cp"].size() > 1:
+                return get_flat_mesh(device_mesh, "dp_cp").get_group()
+            return get_flat_mesh(device_mesh, "dp").get_group()
+        except KeyError:
+            return None
+
+    def _configure_moe_aux_loss_scaler(self) -> None:
+        """Cancel FSDP's gradient averaging on MoE aux-loss grads.
+
+        ``MoEAuxLossAutoScaler`` multiplies aux-loss-derived gradients by
+        ``main_loss_backward_scale`` during backward; FSDP's all-reduce then
+        divides every gradient by ``dp_group_size``. Setting the scaler to
+        ``dp_group_size`` cancels that division out. No-op when
+        ``nemo_automodel`` is unavailable.
+        """
+        try:
+            from nemo_automodel.components.moe.megatron.moe_utils import MoEAuxLossAutoScaler
+        except ImportError:
+            return
+        dp_group = self._get_moe_dp_group()
+        dp_size = dp_group.size() if dp_group is not None else 1
+        MoEAuxLossAutoScaler.main_loss_backward_scale = torch.tensor(float(dp_size))
diff --git a/tests/collections/speechlm2/test_streaming_stt_automodel.py b/tests/collections/speechlm2/test_streaming_stt_automodel.py
new file mode 100644
index 000000000000..465ba77d5c90
--- /dev/null
+++ b/tests/collections/speechlm2/test_streaming_stt_automodel.py
@@ -0,0 +1,510 @@
+# 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.
+"""Tests for :class:`StreamingSTTModelAutomodel`.
+
+``nemo_automodel`` is an optional dependency, so these tests stub the single
+entry point that needs it (``load_pretrained_automodel_llm``) with a small
+HuggingFace LLM. Everything else — deferred construction, the vocabulary
+reconciliation, the freeze policy, the embedding indirection and a real
+training step — is exercised for real.
+"""
+
+import os
+from types import SimpleNamespace
+
+import pytest
+import torch
+from omegaconf import OmegaConf
+from transformers import AutoConfig, AutoModelForCausalLM
+
+from nemo.collections.speechlm2.data.streaming_stt_dataset import AUDIO_TOKEN_IDX, IGNORE_INDEX, StreamingSTTBatch
+from nemo.collections.speechlm2.models import StreamingSTTModelAutomodel
+
+BLANK_TOKEN = "<blank>"
+CHUNK_SIZE = 2
+NUM_AUDIO_SLOTS = 8
+
+
+def resolve_pretrained_llm() -> str:
+    if os.path.exists("/home/TestData/speechlm/pretrained_models/Qwen--Qwen3-1.7B"):
+        return "/home/TestData/speechlm/pretrained_models/Qwen--Qwen3-1.7B"
+    return "Qwen/Qwen3-1.7B"
+
+
+def make_cfg(**overrides) -> dict:
+    """Model config with a tiny, offline-constructible perception module."""
+    cfg = {
+        "pretrained_llm": resolve_pretrained_llm(),
+        "pretrained_asr": "unused-because-load_asr_weights-is-false",
+        "load_llm_weights": False,
+        "load_asr_weights": False,
+        "use_nemo_automodel": True,
+        "blank_token": BLANK_TOKEN,
+        "chunk_size": CHUNK_SIZE,
+        "att_context_size": [70, 1],
+        "audio_pad_to": 0,
+        "sample_rate": 16000,
+        "frame_length_in_secs": 0.08,
+        "dtype": "float32",
+        "freeze_speech_encoder": False,
+        "freeze_modality_adapter": False,
+        "freeze_modality_proj": False,
+        "freeze_llm_model": True,
+        "freeze_llm_head": False,
+        "freeze_embed_tokens": False,
+        "freeze_params": [],
+        "prevent_freeze_params": [],
+        "perception": {
+            "target": "nemo.collections.speechlm2.modules.perception.AudioPerceptionModule",
+            "output_dim": 128,
+            "encoder": {
+                "_target_": "nemo.collections.asr.modules.ConformerEncoder",
+                "att_context_size": [70, 1],
+                "causal_downsampling": True,
+                "conv_context_size": "causal",
+                "conv_kernel_size": 9,
+                "d_model": 64,
+                "feat_in": 80,
+                "feat_out": -1,
+                "ff_expansion_factor": 4,
+                "n_heads": 4,
+                "n_layers": 2,
+                "self_attention_model": "rel_pos",
+                "subsampling": "dw_striding",
+                "subsampling_conv_channels": 64,
+                "subsampling_factor": 8,
+            },
+            "modality_adapter": {
+                "_target_": "nemo.collections.speechlm2.modules.perception.IdentityConnector",
+                "d_model": 64,
+            },
+            "preprocessor": {
+                "_target_": "nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor",
+                "features": 80,
+                "normalize": "per_feature",
+                "sample_rate": 16000,
+                "window_size": 0.025,
+                "window_stride": 0.01,
+            },
+        },
+        "optimizer": {"_target_": "torch.optim.AdamW", "lr": 1e-4},
+    }
+    cfg.update(overrides)
+    return cfg
+
+
+def tiny_llm_factory(monkeypatch):
+    """Patch ``load_pretrained_automodel_llm`` with a small real HF LLM.
+
+    The architecture and (crucially) ``vocab_size`` come from the real Qwen3
+    config, so the spare-embedding-rows logic is exercised faithfully; only the
+    layer sizes are shrunk.
+    """
+    import nemo.collections.speechlm2.models.streaming_stt_model_automodel as mod
+
+    calls = {}
+
+    def _fake_loader(model_path_or_name, pretrained_weights=True, dtype=torch.float32, **kwargs):
+        calls["kwargs"] = kwargs
+        calls["dtype"] = dtype
+        calls["pretrained_weights"] = pretrained_weights
+        config = AutoConfig.from_pretrained(model_path_or_name)
+        config.num_hidden_layers = 2
+        config.hidden_size = 128
+        config.intermediate_size = 256
+        config.num_attention_heads = 4
+        config.num_key_value_heads = 2
+        config.head_dim = 32
+        config.tie_word_embeddings = False
+        return AutoModelForCausalLM.from_config(config, dtype=dtype)
+
+    monkeypatch.setattr(mod, "load_pretrained_automodel_llm", _fake_loader)
+    return calls
+
+
+@pytest.fixture
+def model(monkeypatch):
+    calls = tiny_llm_factory(monkeypatch)
+    m = StreamingSTTModelAutomodel(make_cfg())
+    m.configure_model()
+    m._automodel_calls = calls
+    return m
+
+
+def make_batch(model, batch_size: int = 2, n_text_tokens: int = 6) -> StreamingSTTBatch:
+    """Synthetic batch: [audio slots ..., text tokens ...] with targets on text."""
+    torch.manual_seed(0)
+    total_len = NUM_AUDIO_SLOTS + n_text_tokens
+    input_tokens = torch.full((batch_size, total_len), model.text_pad_id, dtype=torch.long)
+    input_tokens[:, :NUM_AUDIO_SLOTS] = AUDIO_TOKEN_IDX
+    text_ids = torch.randint(low=10, high=1000, size=(batch_size, n_text_tokens))
+    input_tokens[:, NUM_AUDIO_SLOTS:] = text_ids
+
+    target_tokens = torch.full((batch_size, total_len), IGNORE_INDEX, dtype=torch.long)
+    target_tokens[:, NUM_AUDIO_SLOTS:] = text_ids
+
+    n_samples = 16000  # 1 second
+    audios = torch.randn(batch_size, n_samples) * 0.1
+    return StreamingSTTBatch(
+        audios=audios,
+        audio_lens=torch.full((batch_size,), n_samples, dtype=torch.long),
+        input_tokens=input_tokens,
+        input_token_lens=torch.full((batch_size,), total_len, dtype=torch.long),
+        target_tokens=target_tokens,
+        target_token_lens=torch.full((batch_size,), total_len, dtype=torch.long),
+        text=["dummy"] * batch_size,
+        chunk_size=CHUNK_SIZE,
+    )
+
+
+# ===========================================================================
+# Deferred construction
+# ===========================================================================
+
+
+def test_init_defers_llm_and_perception(monkeypatch):
+    tiny_llm_factory(monkeypatch)
+    m = StreamingSTTModelAutomodel(make_cfg())
+    assert m.llm is None
+    assert m.perception is None
+    # The tokenizer is ready immediately (the DataModule needs it before fit).
+    assert m.tokenizer is not None
+    assert m.blank_token == BLANK_TOKEN
+    assert m.blank_token_id > 0
+    # ...and no embedding table exists yet.
+    assert m.embed_tokens is None
+
+
+def test_configure_model_builds_modules(model):
+    assert model.llm is not None
+    assert model.perception is not None
+    # embed_tokens stays *inside* the LLM (unlike StreamingSTTModel).
+    assert model.embed_tokens is model.llm.model.embed_tokens
+    assert "embed_tokens" not in dict(model.named_children())
+    assert any(k.startswith("llm.model.embed_tokens") for k in model.state_dict())
+
+
+def test_configure_model_is_idempotent(model):
+    llm_before = model.llm
+    perception_before = model.perception
+    model.configure_model()
+    assert model.llm is llm_before
+    assert model.perception is perception_before
+
+
+def test_perception_output_dim_matches_llm(model):
+    assert model.perception.proj.out_features == model.llm.config.hidden_size
+
+
+def test_dtype_resolution_from_config(model):
+    assert model._automodel_calls["dtype"] == torch.float32
+    assert model._automodel_calls["pretrained_weights"] is False
+
+
+# ===========================================================================
+# Vocabulary reconciliation
+# ===========================================================================
+
+
+def test_added_special_token_fits_in_spare_rows(model):
+    """Qwen3 ships 151936 embedding rows for a 151669-token vocab, so <blank>
+    lands in a spare row and no resize is needed."""
+    assert len(model.tokenizer.tokenizer) > 0
+    assert model.blank_token_id < model.embed_tokens.weight.shape[0]
+    assert model.text_vocab_size >= len(model.tokenizer.tokenizer)
+
+
+def test_sync_vocab_grows_table_when_too_small(model):
+    """When the checkpoint has no spare rows the table is grown, not shrunk."""
+    target = len(model.tokenizer.tokenizer)
+    model.llm.resize_token_embeddings(target - 5)
+    assert model.embed_tokens.weight.shape[0] == target - 5
+    model._sync_llm_vocab_size()
+    assert model.embed_tokens.weight.shape[0] >= target
+
+
+def test_sync_vocab_keeps_spare_rows(model):
+    """A larger-than-needed table is left alone (the parent class would shrink it)."""
+    rows_before = model.embed_tokens.weight.shape[0]
+    assert rows_before > len(model.tokenizer.tokenizer)
+    model._sync_llm_vocab_size()
+    assert model.embed_tokens.weight.shape[0] == rows_before
+
+
+# ===========================================================================
+# Embedding indirection
+# ===========================================================================
+
+
+def test_embed_tokens_matches_module_call(model):
+    ids = torch.tensor([[1, 2, 3], [4, 5, 6]])
+    torch.testing.assert_close(model._embed_tokens(ids), model.embed_tokens(ids))
+
+
+def test_embed_ref_tensor_dtype(model):
+    assert model._embed_ref_tensor.dtype == model.embed_tokens.weight.dtype
+
+
+def test_move_embedding_ctx_is_a_noop(model):
+    """The parent moves embed_tokens in/out of the LLM for generation; here it
+    must stay put, otherwise the LLM loses its embedding table."""
+    with model._move_embedding_ctx():
+        assert model.llm.model.embed_tokens is not None
+    assert model.llm.model.embed_tokens is not None
+
+
+@pytest.fixture
+def fake_dist():
+    """Minimal env for the fake distributed backend used with LocalTensorMode."""
+    os.environ.setdefault("MASTER_ADDR", "localhost")
+    os.environ.setdefault("MASTER_PORT", "0")
+    yield
+    if torch.distributed.is_initialized():
+        torch.distributed.destroy_process_group()
+
+
+@pytest.mark.parametrize("requires_grad", [True, False])
+def test_embed_tokens_gradient_is_reduced_across_dp_ranks(fake_dist, requires_grad):
+    """A trainable embedding table is looked up OUTSIDE the LLM's FSDP2 forward,
+    so FSDP2 never reduce-scatters its gradient. ``_embed_tokens`` must declare
+    the gradient ``Partial`` so the redistribute back to ``Shard(0)`` performs
+    the DP reduction — otherwise each rank keeps only its own micro-batch
+    contribution and the ``<blank>`` embedding trains on 1/dp_size of the data.
+
+    Rank 0 weighs its lookup by 1.0 and rank 1 by 3.0, so a correctly averaged
+    gradient is 2.0 while an unreduced one is 1.0.
+    """
+    local_tensor_mode = pytest.importorskip(
+        "torch.distributed._local_tensor", reason="LocalTensorMode requires PyTorch >= 2.10"
+    )
+    from torch.distributed.device_mesh import init_device_mesh
+    from torch.distributed.tensor import Shard, distribute_tensor
+
+    world_size = 2
+    torch.distributed.init_process_group(backend="fake", rank=0, world_size=world_size)
+    with local_tensor_mode.LocalTensorMode(world_size):
+        mesh = init_device_mesh("cpu", (world_size,), mesh_dim_names=("dp",))
+        weight = distribute_tensor(torch.zeros(8, 4), mesh, [Shard(0)]).requires_grad_(requires_grad)
+        per_rank_scale = distribute_tensor(torch.tensor([1.0, 3.0]), mesh, [Shard(0)]).to_local()
+        stub = SimpleNamespace(embed_tokens=SimpleNamespace(weight=weight))
+
+        embedded = StreamingSTTModelAutomodel._embed_tokens(stub, torch.zeros(1, dtype=torch.long))
+        assert embedded.shape == (1, 4)
+
+        if not requires_grad:
+            assert not embedded.requires_grad
+            return
+        (embedded * per_rank_scale).sum().backward()
+        row0_grad = weight.grad.full_tensor()[0]
+        torch.testing.assert_close(row0_grad, torch.full((4,), 2.0))
+
+
+# ===========================================================================
+# LLM output normalization
+# ===========================================================================
+
+
+class _TensorOutLLM(torch.nn.Module):
+    """Stand-in for an Automodel backbone that returns bare logits."""
+
+    def __init__(self, vocab=7, hidden=4):
+        super().__init__()
+        self.proj = torch.nn.Linear(hidden, vocab)
+
+    def forward(self, inputs_embeds=None, **kwargs):
+        return self.proj(inputs_embeds)
+
+
+def test_llm_forward_wraps_bare_tensor(model):
+    model.llm = _TensorOutLLM()
+    out = model._llm_forward(inputs_embeds=torch.randn(2, 3, 4))
+    assert out.logits.shape == (2, 3, 7)
+    assert out.past_key_values is None
+
+
+def test_llm_forward_raises_without_cache(model):
+    model.llm = _TensorOutLLM()
+    with pytest.raises(RuntimeError, match="did not return a KV cache"):
+        model._llm_forward(inputs_embeds=torch.randn(2, 3, 4), use_cache=True)
+
+
+def test_llm_forward_raises_without_hidden_states(model):
+    model.llm = _TensorOutLLM()
+    with pytest.raises(RuntimeError, match="did not return hidden states"):
+        model._llm_forward(inputs_embeds=torch.randn(2, 3, 4), output_hidden_states=True)
+
+
+def test_llm_forward_passes_through_model_output(model):
+    out = model._llm_forward(inputs_embeds=torch.randn(2, 3, model.llm.config.hidden_size))
+    assert out.logits.shape[:2] == (2, 3)
+    # Both mapping and attribute access work (the inherited code paths use both).
+    torch.testing.assert_close(out["logits"], out.logits)
+
+
+# ===========================================================================
+# Freeze policy
+# ===========================================================================
+
+
+def test_freeze_llm_body_keeps_head_and_embeddings_trainable(model):
+    assert all(not p.requires_grad for p in model.llm.model.layers.parameters())
+    assert all(p.requires_grad for p in model.llm.lm_head.parameters())
+    # embed_tokens lives inside llm.model, so the per-module switch must win.
+    assert model.embed_tokens.weight.requires_grad
+    assert all(p.requires_grad for p in model.perception.encoder.parameters())
+
+
+def test_freeze_embed_tokens(monkeypatch):
+    tiny_llm_factory(monkeypatch)
+    m = StreamingSTTModelAutomodel(make_cfg(freeze_embed_tokens=True))
+    m.configure_model()
+    assert not m.embed_tokens.weight.requires_grad
+
+
+# ===========================================================================
+# Config validation
+# ===========================================================================
+
+
+def test_hf_peft_lora_keys_rejected(monkeypatch):
+    tiny_llm_factory(monkeypatch)
+    cfg = make_cfg(lora={"task_type": "CAUSAL_LM", "r": 128, "lora_alpha": 256, "lora_dropout": 0.01})
+    with pytest.raises(ValueError, match="HuggingFace-PEFT keys"):
+        StreamingSTTModelAutomodel(cfg)
+
+
+class _FakeMesh:
+    """Minimal stand-in for a DeviceMesh with per-dimension sizes."""
+
+    def __init__(self, **sizes):
+        self.mesh_dim_names = tuple(sizes)
+        self._sizes = sizes
+
+    def __getitem__(self, name):
+        size = self._sizes[name]
+        return SimpleNamespace(size=lambda: size)
+
+
+def test_context_parallelism_rejected(model):
+    model._device_mesh = _FakeMesh(dp=1, cp=2, tp=1)
+    with pytest.raises(ValueError, match="context parallelism"):
+        model._validate_parallelism_compatibility()
+
+
+def test_tensor_parallelism_rejected(model):
+    """TP shards lm_head's output; the inherited blank/non-blank loss breakdown
+    indexes the per-token loss with a plain bool mask, which is not DTensor-safe."""
+    model._device_mesh = _FakeMesh(dp=1, cp=1, tp=2)
+    with pytest.raises(ValueError, match="tensor parallelism"):
+        model._validate_parallelism_compatibility()
+
+
+def test_data_parallel_only_mesh_accepted(model):
+    model._device_mesh = _FakeMesh(dp=8, cp=1, tp=1)
+    model._validate_parallelism_compatibility()  # must not raise
+
+
+def test_configure_model_refuses_to_skip_sharding(model):
+    """A model already built WITHOUT a mesh must not silently no-op when the
+    strategy later calls configure_model() WITH one — that would train every
+    rank unsharded and unsynchronized."""
+    assert model._configured_with_mesh is False  # built by the fixture without a mesh
+    with pytest.raises(RuntimeError, match="already ran without a device mesh"):
+        model.configure_model(device_mesh=_FakeMesh(dp=2, cp=1, tp=1))
+
+
+# ===========================================================================
+# LoRA trainability (the freeze order that bit the sharded path)
+# ===========================================================================
+
+
+def _inject_fake_lora(model):
+    """Attach lora_A/lora_B params the way Automodel does inside from_pretrained."""
+    q_proj = model.llm.model.layers[0].self_attn.q_proj
+    hidden = model.llm.config.hidden_size
+    q_proj.lora_A = torch.nn.Linear(hidden, 4, bias=False)
+    q_proj.lora_B = torch.nn.Linear(4, hidden, bias=False)
+    return [q_proj.lora_A.weight, q_proj.lora_B.weight]
+
+
+def test_unfreeze_lora_params_restores_grads(model):
+    """`freeze_module(llm.model)` also freezes adapters installed by Automodel
+    before the freeze; `prevent_freeze_params` alone cannot bring them back
+    because freeze_and_subset skips requires_grad=False params first."""
+    lora_params = _inject_fake_lora(model)
+    model._apply_freeze_config()  # freeze_llm_model=True -> freezes the adapters too
+    assert all(not p.requires_grad for p in lora_params)
+
+    n = model._unfreeze_lora_params()
+    assert n == len(lora_params)
+    assert all(p.requires_grad for p in lora_params)
+
+
+def test_lora_params_reach_the_optimizer(model):
+    lora_params = _inject_fake_lora(model)
+    model._apply_freeze_config()
+    model.cfg.prevent_freeze_params = [r"^.+\.lora_.+$"]
+    model._unfreeze_lora_params()
+
+    optimizer = model.configure_optimizers()["optimizer"]
+    optimizer_params = {id(p) for group in optimizer.param_groups for p in group["params"]}
+    assert all(id(p) in optimizer_params for p in lora_params)
+
+
+# ===========================================================================
+# End-to-end training step
+# ===========================================================================
+
+
+def test_training_step_produces_gradients(model):
+    batch = make_batch(model)
+    out = model.training_step(batch, batch_idx=0)
+    loss = out["loss"]
+    assert torch.isfinite(loss)
+    loss.backward()
+    grads = [p.grad for p in model.perception.parameters() if p.requires_grad]
+    assert any(g is not None and torch.isfinite(g).all() and g.abs().sum() > 0 for g in grads)
+    assert model.llm.lm_head.weight.grad is not None
+    # The frozen LLM body must not accumulate gradients.
+    assert all(p.grad is None for p in model.llm.model.layers.parameters())
+
+
+def test_validation_step_runs(model):
+    model.on_validation_epoch_start()
+    model.validation_step({"val_set_0": make_batch(model)}, batch_idx=0)
+    assert model._partial_val_losses["val_set_0"]
+    assert torch.isfinite(model._partial_val_losses["val_set_0"][0])
+
+
+def test_configure_optimizers_only_sees_trainable_params(model):
+    ans = model.configure_optimizers()
+    optimizer = ans["optimizer"]
+    optimizer_params = {id(p) for group in optimizer.param_groups for p in group["params"]}
+    frozen = [p for p in model.llm.model.layers.parameters()]
+    assert optimizer_params, "optimizer received no parameters"
+    assert all(id(p) not in optimizer_params for p in frozen)
+    assert id(model.embed_tokens.weight) in optimizer_params
+
+
+def test_oomptimizer_schema(model):
+    schema = model.oomptimizer_schema
+    assert schema["cls"] is StreamingSTTBatch
+    names = {entry["name"] for entry in schema["inputs"]}
+    assert {"input_tokens", "target_tokens", "audios", "audio_lens"} <= names
+
+
+def test_hyperparameters_are_serializable(model):
+    """PTL checkpoints store cfg as a plain dict."""
+    assert isinstance(model.hparams["cfg"], dict)
+    OmegaConf.create(model.hparams["cfg"])
diff --git a/tests/collections/speechlm2/test_streaming_stt_dataset.py b/tests/collections/speechlm2/test_streaming_stt_dataset.py
index cd6e9b58586f..6ca66169818d 100644
--- a/tests/collections/speechlm2/test_streaming_stt_dataset.py
+++ b/tests/collections/speechlm2/test_streaming_stt_dataset.py
@@ -29,8 +29,10 @@
 """
 
 import math
+from types import SimpleNamespace
 
 import pytest
+import torch
 
 from nemo.collections.speechlm2.data.streaming_stt_dataset import (
     AUDIO_TOKEN_IDX,
@@ -38,11 +40,14 @@
     _replace_audio_chunks,
     _tokenize_compact_with_assistant_mask,
     _tokenize_with_assistant_mask,
+    apply_chat_template_ids,
     build_compact_turn_markers,
     compute_word_spans,
     decode_with_blank,
     get_llm_messages_for_batch,
     get_llm_messages_for_sample,
+    resolve_pad_id,
+    right_collate_vectors,
 )
 from nemo.collections.speechlm2.parts.alignments import WordAlignment
 
@@ -1720,6 +1725,90 @@ def real_tokenizer(request):
     return label, hf_tok
 
 
+class TestResolvePadId:
+    """Some LLM tokenizers ship without a pad token (e.g. NVIDIA-Nemotron-3-Nano:
+    ``pad_id`` None, ``unk_id`` 0). Feeding that None to ``pad_sequence`` raises,
+    and the dataset's padding value must match the model's ``text_pad_id``,
+    which is how the attention mask is derived."""
+
+    @staticmethod
+    def _tok(pad_id, unk_id):
+        return SimpleNamespace(pad_id=pad_id, unk_id=unk_id)
+
+    def test_uses_pad_id_when_present(self):
+        assert resolve_pad_id(self._tok(pad_id=7, unk_id=0)) == 7
+
+    def test_falls_back_to_unk(self):
+        assert resolve_pad_id(self._tok(pad_id=None, unk_id=0)) == 0
+
+    def test_falls_back_to_zero_with_warning(self):
+        with pytest.warns(UserWarning, match="no <pad> or <unk> token"):
+            assert resolve_pad_id(self._tok(pad_id=None, unk_id=None)) == 0
+
+    def test_tolerates_tokenizer_without_unk_attribute(self):
+        with pytest.warns(UserWarning):
+            assert resolve_pad_id(SimpleNamespace(pad_id=None)) == 0
+
+    def test_collation_pads_with_the_resolved_id(self):
+        """A padless tokenizer must not crash collation, and the pad value used
+        must be the resolved one (not None)."""
+        pad = resolve_pad_id(self._tok(pad_id=None, unk_id=0))
+        collated = right_collate_vectors(
+            [torch.tensor([5, 6, 7]), torch.tensor([8])],
+            padding_value=pad,
+        )
+        assert collated[1].tolist() == [8, pad, pad]
+
+
+class TestApplyChatTemplateIds:
+    """``apply_chat_template(tokenize=True)`` returns ``list[int]`` on transformers 4.x
+    but a ``BatchEncoding`` on 5.x. ``apply_chat_template_ids`` must flatten both to
+    the same ``list[int]`` — otherwise callers iterate the mapping's *keys* and blow up
+    later with "too many dimensions 'str'"."""
+
+    MESSAGES = [{"role": "system", "content": "hi"}]
+
+    class _FakeTok:
+        """Returns whichever shape a given transformers version would."""
+
+        IDS = [1, 2, 3]
+
+        def __init__(self, shape):
+            self.shape = shape
+            self.seen_kwargs = None
+
+        def apply_chat_template(self, messages, **kwargs):
+            self.seen_kwargs = kwargs
+            if self.shape == "list":  # transformers 4.x
+                return list(self.IDS)
+            if self.shape == "dict":  # transformers 5.x
+                return {"input_ids": list(self.IDS), "attention_mask": [1, 1, 1]}
+            if self.shape == "batched_dict":  # hypothetical future batching
+                return {"input_ids": [list(self.IDS)], "attention_mask": [[1, 1, 1]]}
+            if self.shape == "tensor_dict":
+                import torch
+
+                return {"input_ids": torch.tensor(self.IDS)}
+            raise AssertionError(self.shape)
+
+    @pytest.mark.parametrize("shape", ["list", "dict", "batched_dict", "tensor_dict"])
+    def test_all_return_shapes_flatten_to_ids(self, shape):
+        tok = self._FakeTok(shape)
+        assert apply_chat_template_ids(tok, self.MESSAGES) == [1, 2, 3]
+
+    def test_forwards_kwargs_and_forces_tokenize(self):
+        tok = self._FakeTok("dict")
+        apply_chat_template_ids(tok, self.MESSAGES, add_generation_prompt=False, enable_thinking=False)
+        assert tok.seen_kwargs == {"tokenize": True, "add_generation_prompt": False, "enable_thinking": False}
+
+    def test_matches_the_raw_call_on_this_transformers_version(self, real_tokenizer):
+        """Whatever the installed version returns, the helper yields the same IDs."""
+        _, hf_tok = real_tokenizer
+        raw = hf_tok.apply_chat_template(self.MESSAGES, tokenize=True, add_generation_prompt=False)
+        expected = list(raw["input_ids"]) if hasattr(raw, "keys") else list(raw)
+        assert apply_chat_template_ids(hf_tok, self.MESSAGES, add_generation_prompt=False) == expected
+
+
 class TestTokenizeWithAssistantMaskRealTokenizers:
     """Verify the diff-based fallback works with real HF tokenizers."""