chore: import upstream snapshot with attribution

docs/source/speechlm2/configs.rst

@@ -0,0 +1,544 @@
+Configuration Files
+===================
+
+The SpeechLM2 models use YAML configuration files to define model architecture, training parameters, and data settings.
+This page describes the configuration structure and important parameters for each model type in the collection.
+
+Configuration Structure
+-----------------------
+
+SpeechLM2 configuration files typically have the following high-level structure:
+
+.. code-block:: yaml
+
+    model:
+      # Model architecture settings
+      ...
+    
+    trainer:
+      # PyTorch Lightning trainer settings
+      ...
+    
+    exp_manager:
+      # Experiment logging settings
+      ...
+    
+    data:
+      # Dataset settings
+      ...
+
+SALM Configuration
+------------------
+
+The SALM (Speech-Augmented Language Model) configuration includes settings for the pretrained LLM, audio perception module, and training parameters.
+See the `SALM paper <https://arxiv.org/abs/2310.09424>`_ for more details.
+
+.. code-block:: yaml
+
+    model:
+      # Pretrained model paths
+      pretrained_llm: "TinyLlama/TinyLlama_v1.1"  # HF model path
+      pretrained_asr: "stt_en_fastconformer_hybrid_large_streaming_80ms"  # NeMo checkpoint name
+      pretrained_weights: True  # Whether to load weights or just architecture
+
+      # Fine-tune from a previous training checkpoint (weights only, fresh optimizer)
+      init_from_checkpoint: null  # path to .ckpt, DCP dir, or HF dir
+
+      # Special token settings
+      audio_locator_tag: "<audio>"  # Tag to replace with audio embeddings
+      
+      # Freezing parameters
+      freeze_params:
+        - "^llm\\.model\\.layers\\.[0-4]\\..+$"  # Regex patterns for parameters to freeze
+      prevent_freeze_params: []  # Override freeze_params for specific submodules
+      
+      # Optional LoRA settings for efficient fine-tuning
+      lora:
+        task_type: CAUSAL_LM
+        r: 8
+        lora_alpha: 32
+        lora_dropout: 0.1
+      
+      # Audio perception module configuration
+      perception:
+        target: nemo.collections.speechlm2.modules.perception.AudioPerceptionModule
+        
+        preprocessor:
+          normalize: 'NA'
+        
+        encoder:
+          self_attention_model: rel_pos
+          att_context_size: [-1, -1]
+          conv_context_size: regular
+          conv_norm_type: batch_norm
+        
+        modality_adapter:
+          _target_: nemo.collections.asr.modules.ConformerEncoder
+          feat_in: 1024
+          feat_out: -1
+          n_layers: 2
+          d_model: 1024
+          subsampling: dw_striding
+          subsampling_factor: 1
+          subsampling_conv_channels: 256
+          causal_downsampling: false
+          ff_expansion_factor: 4
+          self_attention_model: rel_pos
+          n_heads: 8
+          att_context_size: [-1, -1]
+          att_context_style: regular
+          xscaling: true
+          untie_biases: true
+          pos_emb_max_len: 5000
+          conv_kernel_size: 9
+          conv_norm_type: batch_norm
+          conv_context_size: null
+          dropout: 0
+          dropout_pre_encoder: 0
+          dropout_emb: 0.0
+
+SALMAutomodel Configuration
+----------------------------
+
+The SALMAutomodel configuration extends the SALM configuration with NeMo Automodel
+support. The key difference is ``use_nemo_automodel: true`` and the use of
+``AutomodelParallelStrategy`` instead of ``DDPStrategy``.
+
+The example below shows a configuration for training with NVIDIA Nemotron Nano V3
+MoE as the LLM backbone, with Expert Parallelism across 8 GPUs:
+
+.. code-block:: yaml
+
+    model:
+      use_nemo_automodel: true  # Selects SALMAutomodel in salm_train.py
+      pretrained_llm: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16
+      pretrained_asr: "nvidia/canary-1b-flash"
+      pretrained_weights: True
+      encoder_chunk_size_seconds: 30.0
+
+      freeze_params:
+        - "^llm\\..+$"
+        - "^perception\\.preprocessor\\..+$"
+        - "^perception\\.encoder\\..+$"
+      prevent_freeze_params: []
+
+      # LoRA uses Automodel-native format (not HF PEFT):
+      # lora:
+      #   dim: 128
+      #   alpha: 256
+      #   dropout: 0.01
+      #   target_modules: ["q_proj", "v_proj"]
+
+      perception:
+        target: nemo.collections.speechlm2.modules.perception.AudioPerceptionModule
+        output_dim: 2048
+        modality_adapter:
+          _target_: nemo.collections.speechlm2.modules.perception.IdentityConnector
+          d_model: 1024
+
+    trainer:
+      strategy:
+        _target_: nemo.collections.speechlm2.parts.parallel.AutomodelParallelStrategy
+        ep_size: 8  # Expert Parallelism across 8 GPUs for MoE
+        # tp_size: 1
+        # dp_size: null  # inferred
+
+NeMo Automodel applies MoE-specific optimizations automatically when an MoE model
+is detected:
+
+* **Grouped GEMM** — fuses expert computations into a single batched matrix multiply
+  for higher GPU throughput.
+* **DeepEP** (Deep Expert Parallelism) — efficient all-to-all expert routing across
+  GPUs, minimizing communication overhead for MoE layers.
+
+Note the differences from the SALM configuration:
+
+* ``model.use_nemo_automodel: true`` — selects ``SALMAutomodel`` in the training script.
+* ``model.pretrained_llm`` can point to MoE models like ``nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16``.
+* ``trainer.strategy._target_`` uses ``AutomodelParallelStrategy`` instead of ``ModelParallelStrategy``.
+* ``ep_size`` controls Expert Parallelism on the FSDP data-parallel axis — dense layers are sharded via FSDP2, while MoE layers use EP for expert routing on the same GPUs.
+* LoRA config uses ``dim``/``alpha`` keys (Automodel native) instead of ``r``/``lora_alpha`` (HF PEFT).
+* No ``embed_tokens`` freeze pattern — embeddings stay inside the LLM.
+* ``encoder_chunk_size_seconds`` controls long-audio chunking for the speech encoder.
+  Audio rows longer than this value are split on the time axis, encoded as a chunk
+  batch, and concatenated back into one embedding sequence before the LLM forward.
+  Set it to ``null`` to disable chunking.
+
+SALMAutomodel-Specific Options
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+The SALMAutomodel config exposes a few extra knobs that pass through to NeMo
+Automodel. All are optional — defaults preserve standard behavior.
+
+**MoE training:**
+
+.. code-block:: yaml
+
+    model:
+      # MoE auxiliary load-balancing loss coefficient. > 0 to enable.
+      # Gradients are injected during backward; reported CE loss is unchanged.
+      aux_loss_coeff: 0.0
+
+      # When true, unfreezes Gate.weight so the router can adapt to new data.
+      # Default false keeps pretrained routing frozen.
+      train_gate: false
+
+      # Per-step expert balance / utilization metrics.
+      moe_metrics:
+        enabled: true
+        mode: brief                  # "brief" or "detailed"
+        detailed_every_steps: null   # null = every step when mode=detailed
+        top_k_experts: 5             # top/bottom utilization experts to report
+
+When ``aux_loss_coeff > 0``, SALMAutomodel sets ``MoEAuxLossAutoScaler.main_loss_backward_scale``
+to the DP group size at ``on_fit_start`` so FSDP's gradient averaging cancels out and the
+net aux-loss gradient scale stays at 1.
+
+**torch.compile:**
+
+.. code-block:: yaml
+
+    model:
+      compile:
+        enabled: false
+        mode: default          # "default" | "reduce-overhead" | "max-autotune"
+        fullgraph: false
+        dynamic: true          # Recommended for variable-length audio
+        backend: null          # null = inductor
+        dynamo_cache_size_limit: 256
+
+**Backend dispatch (attention / linear / norm / MoE kernels):**
+
+.. code-block:: yaml
+
+    model:
+      automodel_backend:
+        attn: te                  # "te" | "sdpa" | "flex"
+        linear: te                # "torch" | "te"
+        rms_norm: torch_fp32      # "torch" | "torch_fp32" | "te"
+        rope_fusion: true
+        experts: torch_mm         # "torch" | "te" | "gmm" | "torch_mm"
+        dispatcher: deepep        # "torch" | "deepep" | "hybridep" | "uccl_ep"
+        dispatcher_num_sms: 20
+
+      # Pin SDPA kernel when automodel_backend.attn=sdpa.
+      # E.g. ["flash_attention"] forces FA2 and errors if unavailable.
+      sdpa_method: null
+
+Defaults come from Automodel's ``BackendConfig`` and auto-select TransformerEngine /
+DeepEP when available; override here to pin a specific backend (for example,
+``attn: sdpa`` to bypass TE).
+
+**Packed sequences (THD):**
+
+.. code-block:: yaml
+
+    model:
+      packed_sequences: true   # default false (right-padded BSHD path)
+      automodel_backend:
+        attn: te               # THD path dispatches TE varlen FlashAttention
+
+When ``packed_sequences`` is true, ``SALMAutomodel.prepare_inputs`` packs
+each minibatch into a single flat ``[T_total, H]`` sequence with a
+``cu_seqlens`` index instead of right-padding to ``[B, T_max, H]``.
+``SALMAutomodel`` then forwards the THD metadata (``qkv_format``,
+``cu_seqlens``, ``position_ids``, ``max_seqlen``) through ``forward()`` to
+the LLM. The TE attention preprocessor splits the singular ``max_seqlen``
+into the ``max_seqlen_q`` / ``max_seqlen_kv`` pair that
+``DotProductAttention`` requires for ``qkv_format="thd"``. The packing also
+rounds each utterance's flat length up to a multiple of ``2 * cp_size`` so
+the same THD batch satisfies TE's CP DualChunkSwap contract — see the
+"Context Parallelism (CP)" subsection in
+:doc:`training_and_scaling` for the recommended pairing with ``cp_size > 1``.
+
+Padding overhead drops from ``O(B * (T_max - T_avg))`` to
+``O(per-utt rounding to 2*cp_size)``. Throughput improvement scales with
+the variance of utterance lengths in your bucketing.
+
+DuplexS2SModel Configuration
+-----------------------------
+
+The DuplexS2SModel adds speech generation capabilities to the configuration:
+
+.. code-block:: yaml
+
+    model:
+      # Pretrained model paths
+      pretrained_llm: "TinyLlama/TinyLlama_v1.1"
+      pretrained_audio_codec: "path/to/audio_codec.nemo"
+      pretrained_asr: "stt_en_fastconformer_hybrid_large_streaming_80ms"
+      scoring_asr: "stt_en_fastconformer_transducer_large"  # used only in validation
+      
+      # Loss weights
+      audio_loss_weight: 4
+      text_loss_weight: 3
+      
+      # Perception module config (similar to SALM)
+      perception:
+        # ... (similar to SALM perception module)
+
+DuplexS2SSpeechDecoderModel Configuration
+-----------------------------------------
+
+The DuplexS2SSpeechDecoderModel is similar to DuplexS2SModel, but focuses on an additional speech generation transformer decoder:
+
+.. code-block:: yaml
+
+    model:
+      # Pretrained model paths
+      pretrained_llm: "TinyLlama/TinyLlama_v1.1"
+      pretrained_audio_codec: "path/to/audio_codec.nemo"
+      pretrained_asr: "stt_en_fastconformer_hybrid_large_streaming_80ms"
+
+      # Speech decoder settings
+      speech_decoder:
+        target: nemo.collections.speechlm2.modules.speech_generation.TransformerARSpeechDecoder
+        d_model: 1024
+        n_layers: 12
+        n_heads: 16
+        d_kv: 64
+        d_ff: 4096
+        max_seq_len: 2048
+        dropout: 0.1
+        layernorm_epsilon: 1e-5
+        activation_function: "gelu_new"
+        init_method_std: 0.02
+        use_cache: True
+
+      # ... other settings
+
+DuplexSTTModel Configuration
+--------------------------------------
+
+The DuplexSTTModel is a speech-to-text model that processes duplex audio conversations and generates agent text responses:
+
+.. code-block:: yaml
+
+    model:
+      # Pretrained model paths
+      pretrained_llm: "TinyLlama/TinyLlama_v1.1"
+      pretrained_asr: "stt_en_fastconformer_hybrid_large_streaming_80ms"
+        
+      # ... other settings
+
+Trainer Configuration
+---------------------
+
+The trainer section contains PyTorch Lightning Trainer settings:
+
+.. code-block:: yaml
+
+    trainer:
+      devices: 1
+      num_nodes: 1
+      accelerator: gpu
+      precision: bf16-true
+      logger: false
+      enable_checkpointing: false  # handled by exp_manager
+      replace_sampler_ddp: false   # handled by lhotse
+      max_epochs: null
+      max_steps: 100000
+      log_every_n_steps: 10
+      val_check_interval: 2000
+      accumulate_grad_batches: 1
+      gradient_clip_val: 1.0
+
+Experiment Manager Configuration
+--------------------------------
+
+The exp_manager section contains settings for experiment logging and model checkpointing:
+
+.. code-block:: yaml
+
+    exp_manager:
+      explicit_log_dir: path/to/output_dir
+      exp_dir: null
+      name: ${name}
+      create_wandb_logger: false  # set to true if you want to use wandb
+      wandb_logger_kwargs:
+        project: null
+        name: null
+      resume_if_exists: true
+      resume_ignore_no_checkpoint: true
+      create_checkpoint_callback: true
+      checkpoint_callback_params:
+        monitor: val_loss
+        filename: "{step}"  # checkpoint name will be step=<step>.ckpt
+        save_top_k: 1
+        mode: min
+      create_tensorboard_logger: false  # set to true if you want to use tensorboard
+      version: null
+
+Data Configuration
+------------------
+
+The data section defines dataset paths, preprocessing, and data loading parameters:
+
+.. code-block:: yaml
+
+    data:
+      train_ds:
+        sample_rate: ${data.target_sample_rate}
+        input_cfg:
+          - type: lhotse_shar
+            shar_path: /path/to/train_data
+        seed: 42
+        shard_seed: "randomized"
+        num_workers: 4
+        batch_size: 16
+        # Optional bucketing settings
+        # batch_duration: 100
+        # bucket_duration_bins: [8.94766,10.1551,11.64118,19.30376,42.85]
+        # use_bucketing: true
+        # num_buckets: 5
+        # bucket_buffer_size: 5000
+    
+      validation_ds:
+        datasets:
+          val_set_name:
+            shar_path: /path/to/validation_data
+        sample_rate: ${data.target_sample_rate}
+        batch_size: 1
+        seed: 42
+        shard_seed: "randomized"
+
+Depending on the model, there may be additional options available under ``data`` namespace that are passed to the corresponding Dataset class.
+For example, S2S models have:
+
+.. code-block:: yaml
+
+    data:
+      frame_length: 0.08
+      source_sample_rate: 16000
+      target_sample_rate: 22050
+      input_roles: ["user", "User"]
+      output_roles: ["agent", "Assistant"]
+
+      train_ds: ...
+
+Important Configuration Parameters
+-----------------------------------
+
+Model Parameters
+^^^^^^^^^^^^^^^^
+
+- **pretrained_llm**: Path to the pretrained HuggingFace LLM
+- **pretrained_asr**: Name of the pretrained NeMo ASR model used for perception
+- **encoder_chunk_size_seconds**: Speech-encoder chunk size in seconds for long audio inputs
+  (supported by both ``SALM`` and ``SALMAutomodel``). Leave as ``null`` to encode each audio
+  row directly
+- **pretrained_audio_codec**: Path to the pretrained audio codec model (for speech generation)
+- **init_from_checkpoint**: Path to a training checkpoint to initialize model weights from (see :ref:`fine-tuning-from-checkpoint` below)
+- **freeze_params**: Regex patterns of parameters to freeze during training
+- **audio_loss_weight/text_loss_weight**: Weighting of different loss components
+
+Perception Module
+^^^^^^^^^^^^^^^^^
+
+- **self_attention_model**: Type of attention mechanism ("rel_pos" or "abs_pos")
+- **att_context_size**: Context window size for attention ([left, right])
+- **conv_context_size**: Context type for convolutions ("causal" or "regular")
+- **n_layers**: Number of encoder layers
+- **d_model**: Model dimension size
+
+Data Parameters
+^^^^^^^^^^^^^^^
+
+- **frame_length**: Frame duration in seconds
+- **source_sample_rate/target_sample_rate**: Sample rates for input/output audio
+- **input_roles/output_roles**: Speaker roles for input and output
+- **batch_size**: Number of samples per batch
+- **use_bucketing**: Whether to use length-based bucketing for efficient batching
+
+Example Configuration Files
+---------------------------
+
+Example configurations for all model types can be found in the example directory:
+
+- SALM: `examples/speechlm2/conf/salm.yaml`
+- SALMAutomodel: `examples/speechlm2/conf/salm_automodel.yaml`
+- DuplexS2SModel: `examples/speechlm2/conf/s2s_duplex.yaml`
+- DuplexS2SSpeechDecoderModel: `examples/speechlm2/conf/s2s_duplex_speech_decoder.yaml`
+- DuplexSTTModel: `examples/speechlm2/conf/duplex_stt.yaml`
+
+Using Configuration Files
+-------------------------
+
+You can use these configurations with the training scripts by specifying the config path:
+
+.. code-block:: bash
+
+    # Train SALM model
+    python examples/speechlm2/salm_train.py \
+      --config-path=conf \
+      --config-name=salm
+
+    # Train SALMAutomodel
+    python examples/speechlm2/salm_train.py \
+      --config-name=salm_automodel
+
+You can also override configuration values from the command line:
+
+.. code-block:: bash
+
+    python examples/speechlm2/salm_train.py \
+      --config-path=conf \
+      --config-name=salm \
+      model.pretrained_llm="different/llm/path" \
+      trainer.max_steps=1000 \
+      data.train_ds.batch_size=8
+
+.. _fine-tuning-from-checkpoint:
+
+Fine-Tuning from a Previous Checkpoint
+---------------------------------------
+
+To start a new training run initialized from a previous checkpoint — with a fresh
+optimizer, LR scheduler, and step counter — set ``model.init_from_checkpoint``:
+
+.. code-block:: yaml
+
+    model:
+      init_from_checkpoint: /path/to/checkpoints/step=6375.ckpt
+
+Or pass it as a Hydra override:
+
+.. code-block:: bash
+
+    python examples/speechlm2/salm_train.py \
+      --config-name=salm_automodel \
+      ++model.init_from_checkpoint=/path/to/checkpoints/step=6375.ckpt
+
+This differs from ``exp_manager.resume_from_checkpoint`` which restores the
+**full** training state (optimizer, scheduler, step counter) to continue an
+interrupted run. ``init_from_checkpoint`` only loads model weights, giving you a
+clean starting point for fine-tuning on different data or with different
+hyperparameters.
+
+Supported Checkpoint Formats
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Three checkpoint formats are supported:
+
+* **Distributed checkpoints (DCP)**: Directories with a ``.metadata`` file, produced
+  by ``ModelParallelStrategy`` / ``AutomodelParallelStrategy``. This is the default
+  format when training with FSDP2 or TP. DCP loading handles automatic resharding
+  when the parallelism configuration differs between the source and target runs.
+
+* **HuggingFace model directories**: Directories containing ``model.safetensors``,
+  such as the output of ``to_hf.py``.
+
+* **Single-file checkpoints**: Standard ``.ckpt`` or ``.pt`` files with a
+  ``state_dict`` key.
+
+The model architecture is still defined by ``pretrained_llm`` and ``pretrained_asr``
+(needed for config and tokenizer initialization), but all weights are overridden by
+the checkpoint.
+
+This feature works with both ``SALM`` and ``SALMAutomodel``.
+
+.. note::
+   ``init_from_checkpoint`` requires the source and target models to use the
+   same model class (e.g., both ``SALMAutomodel``). Cross-model loading
+   (e.g., ``SALM`` checkpoint into ``SALMAutomodel``) will encounter state dict
+   key mismatches because the two classes structure the embedding layer differently.

docs/source/speechlm2/datasets.rst

@@ -0,0 +1,453 @@
+Datasets
+========
+
+The speechlm2 collection supports datasets that contain both audio and text data for training models that can understand speech and generate appropriate responses.
+This section describes the dataset format, preparation, and usage with the speechlm2 models.
+
+Dataset Format
+--------------
+
+Duplex S2S models use the Lhotse framework for audio data management. The primary datasets used are:
+
+1. **DuplexS2SDataset**: For general duplex speech-to-speech models
+2. **SALMDataset**: Specifically for the Speech-Augmented Language Model (SALM), which processes speech+text and outputs text.
+3. **DuplexSTTDataset**: For the DuplexSTTModel, which processes conversational audio and generates text responses.
+4. **DuplexEARTTSDataset**: Dataset for Duplex EARTTS model, extending DuplexS2SDataset with additional output fields for TTS, including audio prompting. It optionally prepends an audio prompt (speaker reference) to target_audio, which is used to initialize speaker conditioning in the EARTTS model. The dataset provides audio_prompt, audio_prompt_lens, non_prompt_mask, aligned_attention_mask, and aligned_position_ids, and supports custom speaker reference audio through the context_audio field, while preserving full compatibility with the original data format.
+
+DuplexS2S Dataset Structure
+^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+A typical dataset for speechlm2 models consists of:
+
+1. **Audio files**: Contains source audio (input speech) and possibly target audio (output speech)
+2. **Text transcriptions**: Associated text for both input and output speech
+3. **Role identifiers**: To distinguish between speakers (e.g., "user" vs "agent")
+
+The dataset organization is built around the concept of conversation turns, with each turn containing audio and text from either a user or an agent/assistant.
+
+The datasets are primarily managed using Lhotse's CutSet format, which provides efficient handling of audio data and annotations. A typical Lhotse manifest includes:
+
+- Audio recording information (path, duration, sample rate)
+- Supervision information (transcripts, speaker roles, timing)
+- Optional additional annotations
+
+Example of a Lhotse cut:
+
+.. code-block:: python
+
+    {
+        "id": "conversation_1",
+        "start": 0,
+        "duration": 10.7,
+        "channel": 0,
+        "supervisions": [
+            {
+                "id": "conversation_1_turn_0",
+                "text": "Can you help me with this problem?",
+                "start": 0,
+                "duration": 5.2,
+                "speaker": "user"
+            },
+            {
+                "id": "conversation_1_turn_1",
+                "text": "I can help you with that.",
+                "start": 5.2,
+                "duration": 3.1,
+                "speaker": "assistant"
+            }
+        ],
+        "recording": {
+            "id": "conversation_1_user",
+            "path": "/path/to/audio/conversation_1_user.wav",
+            "sampling_rate": 16000,
+            "num_samples": 171200,
+            "duration": 10.7
+        },
+        "custom": {
+            "target_audio": {
+                "id": "conversation_1_assistant",
+                "path": "/path/to/audio/conversation_1_assistant.wav",
+                "sampling_rate": 22050,
+                "num_samples": 235935,
+                "duration": 10.7
+            }
+        }
+    }
+
+The DuplexS2SDataset performs several key operations when processing data:
+
+1. **Turn Identification**: Each cut contains a list of `supervisions` with objects of type `lhotse.SupervisionSegment` that represent conversation turns with corresponding text and speaker information.
+
+2. **Speaker Role Separation**: The text of each supervision is tokenized and identified as the model's output (when `supervision.speaker` is in `output_roles`, e.g., "agent" or "Assistant") or the model's input (when in `input_roles`, e.g., "user" or "User").
+
+3. **Token Sequence Generation**:
+   - `target_tokens` and `source_tokens` arrays are created with a length equal to `lhotse.utils.compute_num_frames(cut.duration, frame_length, cut.sampling_rate)`
+   - The `frame_length` parameter (typically 80ms) determines the temporal resolution of token assignments
+   - Each token is assigned to a position based on its corresponding audio segment's timing
+
+4. **Token Offset Calculation**:
+   - The starting position for each turn's tokens is determined using `lhotse.utils.compute_num_frames(supervision.start, frame_length)`
+   - This ensures tokens are aligned with their corresponding audio segments
+
+5. **Length Validation**:
+   - If token sequences are too long compared to the audio duration, warnings are emitted
+   - Tokens that extend beyond the audio length are truncated
+
+This process ensures that the model can correctly align audio input with corresponding text, and learn to generate appropriate responses based on the conversation context.
+
+DuplexS2SDataset
+****************
+
+This dataset class is designed for models that handle both speech understanding and speech generation. It processes audio inputs and prepares them for the model along with corresponding text.
+
+.. code-block:: python
+
+    from nemo.collections.speechlm2.data import DuplexS2SDataset
+    
+    dataset = DuplexS2SDataset(
+        tokenizer=model.tokenizer,                   # Text tokenizer
+        frame_length=0.08,                          # Frame length in seconds
+        source_sample_rate=16000,                   # Input audio sample rate
+        target_sample_rate=22050,                   # Output audio sample rate
+        input_roles=["user", "User"],               # Roles considered as input
+        output_roles=["agent", "Assistant"]         # Roles considered as output
+    )
+
+SALMDataset Structure
+^^^^^^^^^^^^^^^^^^^^^
+
+Data used for SALM can be either regular speech-to-text data (in any NeMo or Lhotse format), or a dataset of multi-turn conversions.
+For the most part, please refer to :doc:`the ASR datasets documentation <../asr/datasets>` for details on data formats and multimodal dataloading.
+
+When using speech-to-text data, you'll need read it with a special ``lhotse_as_conversation`` data reader
+that creates a two-turn, query+response, multi-modal conversation data types out of regular Lhotse cuts.
+This approach makes SALM training more flexible, allowing straightforward combination of single-turn and multi-turn data.
+
+Each audio turn is represented as a single token, defined in ``audio_locator_tag`` property, and automatically added to the model's tokenizer inside model code.
+This token is replaced during the training/generation pass with its corresponding audio segment representation.
+
+Example YAML configuration using existing ASR datasets with ``lhotse_as_conversation``:
+
+.. code-block:: yaml
+
+    data:
+      train_ds:
+        prompt_format: "llama3"  # Choose based on your model
+        token_equivalent_duration: 0.08
+        input_cfg:
+          # Example 1: Using standard ASR Lhotse manifests (JSONL)
+          - type: lhotse_as_conversation
+            cuts_path: /path/to/librispeech_train_clean_100.jsonl.gz
+            audio_locator_tag: "<|audioplaceholder|>"
+            tags:
+              context: "Transcribe the following audio:"
+              # Optional system prompt can be uncommented
+              # system_prompt: "You are a helpful assistant that transcribes audio accurately."
+          
+          # Example 2: Using tarred NeMo manifests
+          - type: lhotse_as_conversation
+            manifest_filepath: /path/to/tedlium_train_manifest.jsonl.gz
+            tarred_audio_filepaths: /path/to/tedlium_shards/shard-{000000..000009}.tar
+            audio_locator_tag: "<|audioplaceholder|>"
+            tags:
+              context: "Write down what is said in this recording:"
+              
+          # Example 3: Using Lhotse SHAR format
+          - type: lhotse_as_conversation
+            shar_path: /path/to/fisher_shar/
+            audio_locator_tag: "<|audioplaceholder|>"
+            tags:
+              context: "Listen to this clip and write a transcript:"
+    
+      # ... other settings
+
+Alternatively, one can provide an existing YAML file with their dataset composition and wrap 
+it in a ``lhotse_as_conversation`` reader as follows:
+
+.. code-block:: yaml
+
+    data:
+      train_ds:
+        input_cfg:
+          - type: lhotse_as_conversation
+            input_cfg: /path/to/dataset_config.yaml
+            audio_locator_tag: "<|audioplaceholder|>"
+            tags:
+              context: "Transcribe the following audio:"
+              # Optional system prompt can be uncommented
+              # system_prompt: "You are a helpful assistant that transcribes audio accurately."
+
+
+The ``lhotse_as_conversation`` reader automatically creates a two-turn conversation from each ASR example:
+1. Optionally, if ``system_prompt`` tag is provided, it's added as a special system turn for LLM models that support system prompts.
+2. A user turn containing the audio and a text context (from the ``context`` tag)
+3. An assistant turn containing the transcription (from the cut's supervision text)
+
+If a ``context`` tag is provided in the configuration, it's added as a text turn before the audio.
+
+SALMDataset
+***********
+
+This dataset class is specialized for the SALM model, which focuses on understanding speech input and generating text output.
+
+.. code-block:: python
+
+    from nemo.collections.speechlm2.data import SALMDataset
+
+    dataset = SALMDataset(
+        tokenizer=model.tokenizer,                   # Text tokenizer
+    )
+
+AIStore GetBatch (multimodal conversations) (experimental)
+**********************************************************
+
+`AIStore GetBatch <https://docs.nvidia.com/aistore/get_batch>`_ is a server-side
+batched object-fetch API; see the `paper <https://arxiv.org/html/2602.22434v1>`_
+for the design and motivation.
+
+For tarred multimodal conversation manifests (``NeMoMultimodalConversationJsonlAdapter``
+and ``NeMoMultimodalConversationShareGPTJsonlAdapter``), set the environment variable
+``USE_AIS_GET_BATCH=true`` to enable AIStore GetBatch loading:
+
+.. code-block:: bash
+
+    USE_AIS_GET_BATCH=true python examples/speechlm2/salm_train.py \
+      --config-name=salm_automodel ...
+
+When enabled:
+
+* The adapters skip opening tar files and instead build URL-backed cuts whose
+  ``AudioSource`` points at the per-shard audio location (the JSONL ``value``
+  field is trusted to match the tar layout).
+* ``SALMDataset`` constructs its loader as
+  ``AudioSamples(use_batch_loader=True, fault_tolerant=True, mono_downmix=True)``,
+  which issues a single batched fetch per minibatch instead of per-cut reads.
+* ``collate_conversation_audio_fault_tolerant`` delegates loading and collation
+  to ``AudioSamples`` and drops every conversation whose cuts didn't survive
+  the fetch — preserving the legacy fault-tolerant semantics.
+
+Leave the env var unset to keep the original tar-iterating loader.
+
+DuplexSTTDataset
+****************
+
+This dataset class is specialized for the DuplexSTTModel, which processes duplex conversational audio and generates text responses. Unlike DuplexS2SDataset which outputs speech, this dataset prepares data for speech-to-text conversion in duplex conversations.
+
+.. code-block:: python
+
+    from nemo.collections.speechlm2.data import DuplexSTTDataset
+
+    dataset = DuplexSTTDataset(
+        tokenizer=model.tokenizer,                   # Text tokenizer
+        frame_length=0.08,                           # Frame length for audio processing
+        source_sample_rate=16000,                    # Audio sample rate
+        input_roles=["User"],                        # Roles to use as input
+        output_roles=["Assistant"],                  # Roles to generate text for
+    )
+
+DataModule
+----------
+
+The DataModule class in the speechlm2 collection manages dataset loading, preparation, and batching for PyTorch Lightning training:
+
+.. code-block:: python
+
+    from nemo.collections.speechlm2.data import DataModule
+    
+    datamodule = DataModule(
+        cfg_data,                  # Configuration dictionary for data
+        tokenizer=model.tokenizer, # Text tokenizer
+        dataset=dataset            # Instance of DuplexS2SDataset, DuplexSTTDataset, or SALMDataset
+    )
+
+The DataModule takes care of:
+1. Setting up proper data parallel ranks for dataloaders
+2. Instantiating the dataloaders with configuration from YAML
+3. Managing multiple datasets for validation/testing
+
+Bucketing for Efficient Training
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+The DataModule supports bucketing for more efficient training. Bucketing groups samples of similar lengths together, which reduces padding and improves training efficiency. The key bucketing parameters are:
+
+1. **batch_duration**: Target cumulative duration (in seconds) of samples in a batch
+2. **bucket_duration_bins**: List of duration thresholds for bucketing
+3. **use_bucketing**: Flag to enable/disable bucketing
+4. **num_buckets**: Number of buckets to create
+5. **bucket_buffer_size**: Number of samples to load into memory for bucket assignment
+
+Example bucketing configuration:
+
+.. code-block:: yaml
+
+    train_ds:
+      # ... other settings
+      batch_duration: 100  # Target 100 seconds per batch
+      bucket_duration_bins: [8.94766, 10.1551, 11.64118, 19.30376, 42.85]  # Duration thresholds
+      use_bucketing: true  # Enable bucketing
+      num_buckets: 5  # Create 5 buckets
+      bucket_buffer_size: 5000  # Buffer size for bucket assignment
+
+When bucketing is enabled:
+
+1. Samples are grouped into buckets based on their duration
+2. Each batch contains samples from the same bucket
+3. The actual batch size can vary to maintain a consistent total duration
+4. The target batch_duration ensures efficient GPU memory usage
+
+Bucketing helps to:
+- Reduce padding and increase effective batch size
+- Improve training efficiency and convergence
+- Manage memory usage with variable-length inputs
+
+Data Configuration
+------------------
+
+A typical data configuration in YAML includes:
+
+.. code-block:: yaml
+
+    data:
+
+      train_ds:
+        sample_rate: ${data.target_sample_rate}
+        input_cfg:
+          - type: lhotse_shar
+            shar_path: /path/to/train_data
+        seed: 42
+        shard_seed: "randomized"
+        num_workers: 4
+        # Optional bucketing settings
+        batch_duration: 100
+        bucket_duration_bins: [8.94766, 10.1551, 11.64118, 19.30376, 42.85]
+        use_bucketing: true
+        num_buckets: 5
+        bucket_buffer_size: 5000
+        # batch_size: 4  # alternative to bucketing
+    
+      validation_ds:
+        datasets:
+          val_set_name_0:
+            shar_path: /path/to/validation_data_0
+          val_set_name_1:
+            shar_path: /path/to/validation_data_1
+        sample_rate: ${data.target_sample_rate}
+        batch_size: 4
+        seed: 42
+        shard_seed: "randomized"
+
+Note that the actual dataset paths and blend are defined by the YAML config, not Python code. This makes it easy to change the dataset composition without modifying the code.
+To learn more about the YAML data config, see :ref:`the Extended multi-dataset configuration format <asr-dataset-config-format>` section in the ASR documentation.
+
+Preparing S2S Datasets
+----------------------
+
+Creating Lhotse Manifests
+^^^^^^^^^^^^^^^^^^^^^^^^^
+
+To prepare your own dataset, you'll need to create Lhotse manifests from your audio files and transcripts:
+
+.. code-block:: python
+
+    from lhotse import CutSet, Recording, SupervisionSegment
+    
+    # Create a recording for user and assistant
+    recording_user = Recording(
+        id="conversation_1_user",
+        path="/path/to/audio/conversation_1_user.wav",
+        sampling_rate=16000,
+        num_samples=171200,
+        duration=10.7
+    )
+    recording_assistant = Recording(
+        id="conversation_1_assistant",
+        path="/path/to/audio/conversation_1_assistant.wav",
+        sampling_rate=22050,
+        num_samples=235935,
+        duration=10.7
+    )
+    
+    # Create supervision for this recording
+    supervisions = [
+        SupervisionSegment(
+            id="conversation_1_turn_0",
+            recording_id="conversation_1",
+            start=0,
+            duration=5.2,
+            text="Can you help me with this problem?",
+            speaker="user"
+        ),
+        SupervisionSegment(
+            id="conversation_1_turn_1",
+            recording_id="conversation_1",
+            start=5.5,
+            duration=3.1,
+            text="I can help you with that.",
+            speaker="assistant"
+        ),
+    ]
+    
+    # Create a CutSet
+    # The assistant's response is located in target_audio field which makes it easy to replace
+    # when using multiple models or speakers for synthetic data generation.
+    cut = recording.to_cut()
+    cut.supervisions = supervisions
+    cut.target_audio = recording_assistant
+    cutset = CutSet([cut])
+    
+    # Save to disk
+    cutset.to_file("path/to/manifest.jsonl.gz")
+
+Converting to SHAR Format
+^^^^^^^^^^^^^^^^^^^^^^^^^
+
+For efficient training, it's recommended to convert your Lhotse manifests to SHAR (SHarded ARchive) format:
+
+.. code-block:: python
+
+    from lhotse import CutSet
+    from lhotse.shar import SharWriter
+    
+    cutset = CutSet.from_file("path/to/manifest.jsonl.gz")
+    cutset.to_shar("path/to/train_shar", fields={"recording": "flac", "target_audio": "flac"}, shard_size=100)
+    
+
+Training with Prepared Datasets
+-------------------------------
+
+Once your datasets are prepared, you can use them to train a model:
+
+.. code-block:: python
+
+    # Load configuration
+    config_path = "path/to/config.yaml"
+    cfg = OmegaConf.load(config_path)
+    
+    # The training data paths are available in the config file:
+    # cfg.data.train_ds.input_cfg[0].shar_path = "path/to/train_shar"
+    
+    # Create dataset and datamodule
+    dataset = DuplexS2SDataset(
+        tokenizer=model.tokenizer,
+        frame_length=cfg.data.frame_length,
+        source_sample_rate=cfg.data.source_sample_rate,
+        target_sample_rate=cfg.data.target_sample_rate,
+        input_roles=cfg.data.input_roles,
+        output_roles=cfg.data.output_roles,
+    )
+    datamodule = DataModule(cfg.data, tokenizer=model.tokenizer, dataset=dataset)
+    
+    # Train the model
+    trainer.fit(model, datamodule)
+
+Example S2S Datasets
+--------------------
+
+While there are no publicly available datasets specifically formatted for Duplex S2S models yet, you can adapt conversation datasets with audio recordings such as:
+
+1. Fisher Corpus
+2. Switchboard Corpus
+3. CallHome
+4. Synthetic conversation datasets generated using TTS
+
+You would need to format these datasets as Lhotse manifests with appropriate speaker role annotations to use them with the speechlm2 S2S models. 

docs/source/speechlm2/intro.rst

@@ -0,0 +1,399 @@
+SpeechLM2
+================================
+
+.. note::
+   The SpeechLM2 collection is still in active development and the code is likely to keep changing.
+
+.. note::
+   Install your chosen compatible PyTorch stack first, then install SpeechLM2 with
+   ``uv pip install 'nemo-toolkit[speechlm2]'`` (or, from a source checkout, ``uv pip install -e '.[speechlm2]'``)
+   to get all required dependencies including NeMo Automodel. See :ref:`installation` for details.
+
+SpeechLM2 refers to a collection that augments pre-trained Large Language Models (LLMs) with speech understanding and generation capabilities.
+
+This collection is designed to be compact, efficient, and to support easy swapping of different LLMs backed by HuggingFace AutoModel or NeMo Automodel. 
+It has a first-class support for using dynamic batch sizes via Lhotse and various model parallelism techniques (e.g., FSDP2, Tensor Parallel, Sequence Parallel) via PyTorch DTensor API.
+
+We currently support six main model types:
+
+* **SALM** (Speech-Augmented Language Model) - a simple but effective approach to augmenting pre-trained LLMs with speech understanding capabilities. Available in two variants:
+
+  * ``SALM`` — uses HuggingFace Transformers for the LLM backbone with optional HF PEFT LoRA.
+  * ``SALMAutomodel`` — uses `NeMo Automodel <https://github.com/NVIDIA-NeMo/Automodel>`_ for the LLM backbone with native LoRA, advanced parallelism (FSDP2, TP, SP, EP via ``AutomodelParallelStrategy``), and MoE optimizations (Grouped GEMM, DeepEP) for efficient training with models like `NVIDIA Nemotron Nano V3 <https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16>`_.
+* **DuplexS2SModel** - a full-duplex speech-to-speech model with an ASR encoder, directly predicting discrete audio codes.
+* **DuplexS2SSpeechDecoderModel** - a variant of DuplexS2SModel with a separate transformer decoder for speech generation.
+* **DuplexEARTTS** - a ready-to-use duplex text-to-speech model that supports user interruption via a special text interruption token.
+* **DuplexSTTModel** - a decoder model to generate agent text in duplex, in response to both user speech and text inputs.
+* **NemotronVoiceChat** - an *inference-only* pipeline that seamlessly merges `DuplexSTTModel` and `DuplexEARTTS` to deliver an end-to-end, full-duplex conversational agent with high-fidelity speech generation.
+
+Using Pretrained Models
+-----------------------
+
+After :ref:`installing NeMo<installation>`, you can load and use a pretrained speechlm2 model as follows:
+
+.. code-block:: python
+
+    import nemo.collections.speechlm2 as slm
+    
+    # Load a pretrained SALM model
+    model = slm.models.SALM.from_pretrained("model_name_or_path")
+
+    # Set model to evaluation mode
+    model = model.eval()
+
+Inference with Pretrained Models
+--------------------------------
+
+SALM
+****
+
+You can run inference using the loaded pretrained SALM model:
+
+.. code-block:: python
+
+    import torch
+    import soundfile as sf
+    from nemo.collections.audio.parts.utils.transforms import resample
+    import nemo.collections.speechlm2 as slm
+
+    model = slm.models.SALM.from_pretrained("path/to/pretrained_checkpoint").eval()
+    
+    # Load audio file
+    audio_path = "path/to/audio.wav"
+    audio_signal, sample_rate = sf.read(audio_path)
+    audio_signal = torch.tensor(audio_signal).unsqueeze(0)
+    
+    # Resample if needed
+    if sample_rate != 16000:  # Most models expect 16kHz audio
+        audio_signal = resample(audio_signal, sample_rate, 16000)
+        sample_rate = 16000
+    
+    # Prepare audio for model
+    audio_signal = audio_signal.to(model.device)
+    audio_len = torch.tensor([audio_signal.shape[1]], device=model.device)
+    
+    # Create a prompt for SALM model inference
+    # The audio_locator_tag is a special token that will be replaced with audio embeddings
+    prompt = [{"role": "user", "content": f"{model.audio_locator_tag}"}]
+    
+    # Generate response
+    with torch.inference_mode():
+        output = model.generate(
+            prompts=[prompt],
+            audios=audio_signal,
+            audio_lens=audio_len,
+            generation_config=None  # You can customize generation parameters here
+        )
+    
+    # Process the output tokens
+    response = model.tokenizer.ids_to_text(output[0])
+    print(f"Model response: {response}")
+
+SALMAutomodel
+*************
+
+``SALMAutomodel`` is the NeMo Automodel variant of SALM. It enables efficient training of
+Speech LLMs with MoE architectures like `NVIDIA Nemotron Nano V3 <https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16>`_
+using MoE-specific optimizations (Grouped GEMM, DeepEP). It uses deferred initialization
+(``configure_model()``) and supports distributed training and inference via
+``AutomodelParallelStrategy``.
+
+.. code-block:: python
+
+    import torch
+    import nemo.collections.speechlm2 as slm
+    from nemo.collections.speechlm2.parts.parallel import setup_distributed
+
+    # Initialize distributed and create an Automodel-compatible device mesh with EP=2.
+    # setup_distributed delegates mesh creation to nemo_automodel, which builds
+    # the full (pp, dp_replicate, dp_shard, cp, tp) mesh with MoE submeshes.
+    strategy = setup_distributed(ep_size=2)
+
+    # Load a pretrained SALMAutomodel with the Automodel device mesh
+    model = slm.models.SALMAutomodel.from_pretrained(
+        "path/to/checkpoint",
+        device_mesh=strategy.device_mesh,
+        distributed_config=strategy.distributed_config,
+        moe_config=strategy.moe_config,
+        moe_mesh=strategy.moe_mesh,
+    ).eval()
+
+    # Inference is identical to SALM
+    with torch.inference_mode():
+        output = model.generate(
+            prompts=[prompt],
+            audios=audio_signal,
+            audio_lens=audio_len,
+        )
+
+DuplexS2SModel
+**************
+
+You can run inference using the loaded pretrained DuplexS2SModel:
+
+.. code-block:: python
+
+    import torch
+    import soundfile as sf
+    from nemo.collections.audio.parts.utils.transforms import resample
+    import nemo.collections.speechlm2 as slm
+
+    model = slm.models.DuplexS2SModel.from_pretrained("path/to/pretrained_checkpoint").eval()
+    
+    # Load audio file
+    audio_path = "path/to/audio.wav"
+    audio_signal, sample_rate = sf.read(audio_path)
+    audio_signal = torch.tensor(audio_signal).unsqueeze(0)
+    
+    # Resample if needed
+    if sample_rate != 16000:  # Most models expect 16kHz audio
+        audio_signal = resample(audio_signal, sample_rate, 16000)
+        sample_rate = 16000
+    
+    # Prepare audio for model
+    audio_signal = audio_signal.to(model.device)
+    audio_len = torch.tensor([audio_signal.shape[1]], device=model.device)
+    
+    # Run offline inference
+    results = model.offline_inference(
+        input_signal=audio_signal,
+        input_signal_lens=audio_len
+    )
+
+    # Decode text and audio tokens
+    transcription = results["text"][0]
+    audio = results["audio"][0]
+
+DuplexSTTModel
+**************
+
+You can run inference using the loaded pretrained DuplexSTTModel:
+
+.. code-block:: python
+
+    import torch
+    import soundfile as sf
+    from nemo.collections.audio.parts.utils.transforms import resample
+    import nemo.collections.speechlm2 as slm
+
+    model = slm.models.DuplexSTTModel.from_pretrained("path/to/pretrained_checkpoint").eval()
+
+    # Load audio file
+    audio_path = "path/to/audio.wav"
+    audio_signal, sample_rate = sf.read(audio_path)
+    audio_signal = torch.tensor(audio_signal).unsqueeze(0)
+
+    # Resample if needed
+    if sample_rate != 16000:
+        audio_signal = resample(audio_signal, sample_rate, 16000)
+        sample_rate = 16000
+
+    # Prepare audio for model
+    audio_signal = audio_signal.to(model.device)
+    audio_len = torch.tensor([audio_signal.shape[1]], device=model.device)
+
+    # Run offline inference - generates text output only
+    results = model.offline_inference(
+        input_signal=audio_signal,
+        input_signal_lens=audio_len
+    )
+
+    # Decode text tokens
+    transcription = results["text"][0]
+    print(f"Transcription: {transcription}")
+
+DuplexEARTTS
+************
+
+Because `DuplexEARTTS` relies on precise token padding and EOS placement to handle potential user interruptions, inference and evaluation are handled via the `duplex_eartts_eval.py` script following the MagpieTTS dataset format recipe. 
+
+The evaluation script processes a `JSONL` file where each line is a dictionary containing the text, the reference audio for the speaker, and the desired output audio filename. 
+
+**JSONL Format Examples:**
+
+Single-Turn format (evaluates a continuous string):
+
+.. code-block:: json
+
+    {"text": "Like really quickly and then they run off.", "context_audio_filepath": "speaker_1.wav", "audio_filepath": "audio_1.wav"}
+
+Multi-Turn format (evaluates sequential conversational turns, padded incrementally):
+
+.. code-block:: json
+
+    {"text": ["Yes.", "Sure.", "Right.", "I get what you’re saying."], "context_audio_filepath": "speaker_2.wav", "audio_filepath": "audio_2.wav"}
+
+**Running the Evaluation/Inference Script:**
+
+.. code-block:: bash
+
+    python examples/speechlm2/duplex_eartts_eval.py \
+        --config-path=conf/ \
+        --config-name=duplex_eartts.yaml \
+        ++checkpoint_path=/path/to/duplex_eartts/model.ckpt \
+        ++datasets_json_path=/path/to/evalset_config.jsonl \
+        ++out_dir=/path/to/output/audio_samples/ \
+        ++user_custom_speaker_reference=/path/to/optional_override_speaker.wav
+
+The script will decode the text, apply the target speaker conditioning, generate the resulting audio waveforms into `out_dir`, and compute ASR intelligibility metrics (CER/WER) on the generated speech.
+
+NemotronVoiceChat
+*****************
+
+You can evaluate and run full-duplex inference using the `NemotronVoiceChat` pipeline. This model natively chains the `DuplexSTTModel` with the `DuplexEARTTS` speech decoder for an end-to-end response:
+
+.. code-block:: python
+
+    import torch
+    import soundfile as sf
+    from nemo.collections.audio.parts.utils.transforms import resample
+    import nemo.collections.speechlm2 as slm
+
+    model = slm.models.NemotronVoiceChat.from_pretrained("path/to/pretrained_checkpoint").eval()
+
+    # Load user audio prompt
+    audio_path = "path/to/user_audio.wav"
+    audio_signal, sample_rate = sf.read(audio_path)
+    audio_signal = torch.tensor(audio_signal).unsqueeze(0)
+
+    # Resample to the source_sample_rate (usually 16kHz for STT perception)
+    if sample_rate != 16000:
+        audio_signal = resample(audio_signal, sample_rate, 16000)
+        sample_rate = 16000
+
+    # Prepare audio for model
+    audio_signal = audio_signal.to(model.device)
+    audio_len = torch.tensor([audio_signal.shape[1]], device=model.device)
+
+    # (Optional) Load an explicit speaker reference audio to condition the agent's voice
+    # speaker_audio, _ = sf.read("path/to/speaker_reference.wav")
+    # speaker_audio = torch.tensor(speaker_audio).unsqueeze(0).to(model.device)
+    # speaker_len = torch.tensor([speaker_audio.shape[1]], device=model.device)
+
+    # Note: If an explicit audio reference is not passed into `offline_inference`, 
+    # the model relies on the internal config parameters:
+    # 1. model.cfg.inference_speaker_name (Highest priority preset, e.g., 'Megan')
+    # 2. model.cfg.inference_speaker_reference (Fallback audio file path)
+        
+    # Run full offline inference
+    results = model.offline_inference(
+        input_signal=audio_signal,
+        input_signal_lens=audio_len,
+        # speaker_audio=speaker_audio,       # Pass speaker reference if available
+        # speaker_audio_lens=speaker_len
+    )
+
+    # Decode the predicted text and generated speech waveform
+    generated_text = results["text"][0]
+    generated_speech = results["audio"][0]
+    
+    print(f"Agent response: {generated_text}")
+    # generated_speech can now be saved or played (sampled at model.target_sample_rate)
+    
+
+Training a Model
+----------------
+
+This example demonstrates how to train a SALM model. 
+
+.. note::
+   **NemotronVoiceChat is an inference-only class.** It does not implement a `training_step` and cannot be trained using the pipeline below. To update its underlying capabilities, you must train the `DuplexSTTModel` and `DuplexEARTTS` models independently.
+
+.. code-block:: python
+
+    from omegaconf import OmegaConf
+    import torch
+    from lightning.pytorch import Trainer
+    from lightning.pytorch.strategies import ModelParallelStrategy
+    
+    import nemo.collections.speechlm2 as slm
+    from nemo.collections.speechlm2.data import SALMDataset, DataModule
+    from nemo.utils.exp_manager import exp_manager
+    
+    # Load configuration
+    config_path = "path/to/config.yaml"  # E.g., from examples/speechlm2/conf/salm.yaml
+    cfg = OmegaConf.load(config_path)
+    
+    # Initialize PyTorch Lightning trainer
+    trainer = Trainer(
+        max_steps=100000,
+        accelerator="gpu",
+        devices=1,
+        precision="bf16-true",
+        strategy=ModelParallelStrategy(data_parallel_size=2, tensor_parallel_size=1),
+        limit_train_batches=1000,
+        val_check_interval=1000,
+        use_distributed_sampler=False,
+        logger=False,
+        enable_checkpointing=False,
+    )
+    
+    # Set up experiment manager for logging
+    exp_manager(trainer, cfg.get("exp_manager", None))
+    
+    # Initialize model with configuration
+    model = slm.models.SALM(OmegaConf.to_container(cfg.model, resolve=True))
+    
+    # Create dataset and datamodule
+    dataset = SALMDataset(tokenizer=model.tokenizer)
+    datamodule = DataModule(cfg.data, tokenizer=model.tokenizer, dataset=dataset)
+    
+    # Train the model
+    trainer.fit(model, datamodule)
+
+Example Using Command-Line Training Script
+------------------------------------------
+
+Alternatively, you can train a model using the provided training scripts in the examples directory:
+
+.. code-block:: bash
+
+    # Train a SALM model
+    python examples/speechlm2/salm_train.py \
+      --config-path=examples/speechlm2/conf \
+      --config-name=salm
+
+    # For SALM inference/evaluation
+    python examples/speechlm2/salm_eval.py \
+      pretrained_name=/path/to/checkpoint \
+      inputs=/path/to/test_manifest \
+      batch_size=64 \
+      max_new_tokens=128 \
+      output_manifest=generations.jsonl
+
+To train the SALMAutomodel variant (with NeMo Automodel backend), use the ``salm_automodel`` config:
+
+.. code-block:: bash
+
+    # Train SALMAutomodel with NVIDIA Nemotron Nano V3 MoE backbone on 8 GPUs
+    torchrun --nproc_per_node=8 examples/speechlm2/salm_train.py \
+      --config-name=salm_automodel \
+      model.pretrained_llm=nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16
+
+The ``salm_automodel.yaml`` config sets ``model.use_nemo_automodel: true``, which selects the
+``SALMAutomodel`` class. This variant supports ``AutomodelParallelStrategy`` for FSDP2/TP/EP
+parallelism and MoE optimizations (Grouped GEMM, DeepEP).
+
+For more detailed information on training at scale, model parallelism, and SLURM-based training, see :doc:`training and scaling <training_and_scaling>`.
+
+Collection Structure
+--------------------
+
+The speechlm2 collection is organized into the following key components:
+
+- **Models**: Contains implementations of DuplexS2SModel, DuplexS2SSpeechDecoderModel, DuplexSTTModel, SALM, SALMAutomodel, DuplexEARTTS, and the inference-only NemotronVoiceChat.
+- **Modules**: Contains audio perception and speech generation modules.
+- **Data**: Includes dataset classes and data loading utilities.
+
+SpeechLM2 Documentation
+-----------------------
+
+For more information, see additional sections in the SpeechLM2 docs:
+
+.. toctree::
+   :maxdepth: 1
+
+   models
+   datasets
+   configs
+   training_and_scaling

docs/source/speechlm2/models.rst

@@ -0,0 +1,341 @@
+Models
+======
+
+The Duplex Speech-to-Speech (S2S) collection consists of several model architectures designed to enable conversational AI systems with speech understanding and generation capabilities. These models combine audio perception components with language models and speech synthesis to create end-to-end speech-to-speech systems.
+
+Core Model Architectures
+------------------------
+
+The collection includes the following core model architectures:
+
+
+DuplexEARTTS
+^^^^^^^^^^^^
+
+DuplexEARTTS is a streaming text-to-speech model designed for duplex speech-to-speech systems. It focuses on low-latency, fully streamable speech generation by converting text tokens into audio representations in real time.
+
+The architecture is based on the Streaming TTS model proposed in `Audio Flamingo 3 <https://arxiv.org/abs/2507.08128>`_, with several extensions for duplex interaction:
+
+* **Gated fusion of text and audio representations**: (`GatedProjectedSumRMSNorm`), enabling better multimodal integration.
+* **Subword-aware embeddings**: (`SubwordFlagEmbedding`) to improve pronunciation for words composed of multiple text tokens.
+* **Custom BOS/EOS embeddings**: (`BOSEOSEmbedding`) for interruption-aware, multi-turn duplex generation.
+
+
+Key components:
+
+* **RVQVAEModel**: An RVQ-based neural audio codec that compresses speech into discrete acoustic tokens using a convolutional encoder and reconstructs high-quality audio via a convolutional decoder.
+* **RVQEARTTSModel**: A streaming speech generation model that predicts multiple RVQ codebooks in parallel using a Mixture-of-Gaussians (MoG) prediction head. It produces audio tokens autoregressively from text representations with minimal latency.
+
+DuplexEARTTS is particularly useful for:
+* Duplex speech-to-speech systems requiring interruption-aware synthesis.
+* Low-latency text-to-speech generation.
+* Real-time conversational agents with streamed audio output.
+
+
+SALM (Speech-Augmented Language Model)
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+SALM is a speech-augmented language model that integrates an audio perception module with a pre-trained LLM. The model is designed to understand speech input and generate text responses.
+This is an implementation of the `SALM paper <https://arxiv.org/abs/2310.09424>`_.
+
+Key components:
+
+* **Audio Perception Module**: Processes audio inputs and converts them into embeddings that can be understood by the LLM. This module leverages a pre-trained ASR model's encoder.
+* **LLM**: A pre-trained large language model that receives the audio embeddings and generates appropriate text responses.
+* **Audio Locator Tag**: A special token that serves as a placeholder in the input text, which gets replaced with the audio embeddings during processing.
+
+SALM is particularly useful for:
+* Speech-to-text applications where high-quality text generation is required
+* Speech understanding tasks that benefit from the contextual understanding of an LLM
+* Applications that need to handle mixed text and speech inputs
+
+SALMAutomodel (NeMo Automodel variant)
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+SALMAutomodel is a variant of SALM that replaces HuggingFace Transformers with
+`NeMo Automodel <https://github.com/NVIDIA-NeMo/Automodel>`_ for the LLM backbone.
+It provides the same speech-augmented language model capabilities with additional
+benefits for distributed training and inference — especially with Mixture-of-Experts
+(MoE) models like `NVIDIA Nemotron Nano V3 <https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16>`_.
+
+Key differences from SALM:
+
+* **Deferred initialization**: The LLM and perception module are not loaded until
+  ``configure_model()`` is called, enabling memory-efficient distributed loading
+  where each GPU only loads its own shard.
+* **Native LoRA**: Uses NeMo Automodel's built-in LoRA instead of HuggingFace PEFT.
+  LoRA is applied before FSDP2 sharding for correct meta-device handling.
+* **AutomodelParallelStrategy**: Integrates with a custom Lightning strategy that
+  delegates device mesh creation to NeMo Automodel, supporting FSDP2, TP, PP, CP,
+  EP (MoE), and HSDP.
+* **MoE optimizations**: NeMo Automodel provides first-class support for
+  Mixture-of-Experts architectures with **Grouped GEMM** (fused expert computation
+  for higher throughput) and **DeepEP** (Deep Expert Parallelism for efficient
+  all-to-all expert routing across GPUs). This makes it possible to efficiently
+  train Speech LLMs using MoE backbones like NVIDIA Nemotron Nano V3 (30B total,
+  3B active parameters).
+* **Embedding stays in LLM**: Unlike SALM (which moves ``embed_tokens`` out of the
+  LLM to avoid FSDP/TP hook issues), SALMAutomodel keeps embeddings inside the LLM
+  and uses ``F.embedding`` with explicit DTensor handling.
+* **Encoder chunking**: Long audio inputs can be split into fixed-duration chunks
+  for the speech encoder, batched through the perception forward, and concatenated
+  back into one embedding sequence before the LLM forward. This is controlled by
+  ``model.encoder_chunk_size_seconds`` in ``salm_automodel.yaml``; set it to
+  ``null`` to encode each audio row directly. The same knob is also available in
+  ``SALM`` via ``model.encoder_chunk_size_seconds`` in ``salm.yaml`` (defaults to
+  ``null`` there to preserve existing behavior).
+
+SALMAutomodel is particularly useful for:
+
+* Efficient training of Speech LLMs with MoE backbones (e.g., Nemotron Nano V3)
+* Large-scale distributed training with advanced parallelism strategies (FSDP2 with EP for MoE, TP)
+* Models that benefit from native Automodel LoRA integration
+* Inference with model parallelism (TP/EP)
+
+DuplexS2SModel
+^^^^^^^^^^^^^^
+
+The DuplexS2SModel extends the SALM architecture to enable full speech-to-speech capabilities, adding the ability to generate speech output.
+
+Key components:
+
+* **Audio Perception Module**: Processes audio inputs into embeddings for the LLM.
+* **LLM**: Processes audio embeddings and generates text or audio token responses.
+* **Audio Codec**: Converts discrete audio tokens into speech waveforms.
+* **Audio Head**: Predicts audio tokens for speech generation.
+* **Embed Audio Tokens**: Embedding layers for audio token representation.
+
+This model is particularly useful for:
+* Complete conversational agents that can listen and respond with speech
+* Voice assistants and interactive dialogue systems
+* Applications requiring natural-sounding spoken responses
+
+DuplexS2SSpeechDecoderModel
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+This model focuses on the speech generation aspect of the duplex system, optimizing the decoder for high-quality speech output.
+
+Key components:
+
+* **Audio Perception Module**: Similar to other models, processes audio into embeddings.
+* **Speech Decoder**: A specialized component for generating high-quality speech from text or audio representations.
+* **TransformerARSpeechDecoder**: An autoregressive transformer-based decoder for speech generation.
+
+This model is particularly useful for:
+* Applications focusing on high-quality speech synthesis
+* Systems where speech generation quality is the primary concern
+* Specialized voice output applications
+
+DuplexSTTModel
+^^^^^^^^^^^^^^
+
+This model focuses on speech-to-text conversion in duplex conversations, processing both user speech and generating agent text responses.
+
+Key components:
+
+* **Audio Perception Module**: Processes audio inputs into embeddings for the LLM.
+* **LLM**: Processes audio embeddings and generates text responses.
+* **Tokenizer**: Converts LLM outputs into readable text.
+
+This model is particularly useful for:
+* Speech-to-text applications in conversational settings
+* Duplex systems where text responses are needed instead of speech
+* Applications requiring transcript generation from spoken dialogue
+
+
+NemotronVoiceChat
+^^^^^^^^^^^^^^^^^
+
+NemotronVoiceChat is an **inference-only**, end-to-end Duplex Speech-to-Speech pipeline. It achieves full-duplex conversational capabilities by seamlessly merging the `DuplexSTTModel` with the `DuplexEARTTS` model. 
+
+Because it is designed exclusively for evaluation, offline inference, and validation workflows (no training step is implemented), it is highly optimized for executing the full perception-generation-synthesis loop.
+
+Key components:
+
+* **DuplexSTTModel**: Handles the streaming audio perception and text response generation.
+* **DuplexEARTTS**: Serves as the autoregressive speech decoder, generating high-fidelity audio from the STT model's text tokens in a streamable fashion.
+
+This model is particularly useful for:
+* End-to-end evaluation of the complete speech-to-speech pipeline.
+* Offline speech-to-speech inference workflows.
+
+
+Model Components
+----------------
+
+Audio Perception Module
+^^^^^^^^^^^^^^^^^^^^^^^
+
+The audio perception module is responsible for converting speech signals into embeddings that can be processed by language models. It typically consists of:
+
+1. **Preprocessor**: Converts raw audio waveforms into spectral features
+2. **Encoder**: Processes these features to create meaningful representations
+3. **Modality Adapter**: Adapts the encoder outputs to be compatible with the LLM's input space
+
+Speech Generation
+^^^^^^^^^^^^^^^^^
+
+Speech generation components convert text or token representations back into speech. The collection offers:
+
+1. **TransformerARSpeechDecoder**: An autoregressive transformer-based speech decoder
+2. **Audio Codec Integration**: Works with audio codecs to generate natural speech from discrete tokens
+3. **DuplexEARTTS**: A ready-to-use duplex text-to-speech model that supports user interruption via a special text interruption token. The model integrates an RVQ-based audio codec with a streaming speech generation module to enable low-latency, real-time synthesis.
+
+Implementation Details
+----------------------
+
+The DuplexS2SModel implementation contains several key methods that handle different aspects of the model's functionality:
+
+Model Initialization
+^^^^^^^^^^^^^^^^^^^^
+
+The constructor (`__init__`) initializes the following components:
+
+1. **Pretrained ASR encoder/perception**: Loads a pretrained NeMo ASR model and adapts it for audio perception
+2. **LLM**: Loads a pretrained language model using HuggingFace's AutoModel
+3. **Audio codec**: Loads a pretrained NeMo audio codec for speech generation
+4. **Token prediction heads**: Adds separate heads for text token and audio token prediction
+
+Forward Method
+^^^^^^^^^^^^^^
+
+The `forward` method:
+
+1. Accepts input representations (sum of audio perception and text embedding outputs)
+2. Runs an offline forward pass through the language model
+3. Generates logits from both text and audio token prediction heads
+4. Returns these logits for loss computation
+
+Training Step
+^^^^^^^^^^^^^
+
+The `training_step` method:
+
+1. Builds input representations using the `prepare_inputs` method
+2. Runs the `forward` method to get text and audio logits
+3. Computes losses for both text and audio tokens
+4. Logs training metrics (loss, learning rate, etc.)
+5. Returns the loss for backpropagation
+
+Prepare Inputs
+^^^^^^^^^^^^^^
+
+The `prepare_inputs` method:
+
+1. Processes source audio through the perception module (with gradients enabled)
+2. Processes target audio through the audio codec (with gradients disabled)
+3. Truncates source/target audio and target text sequences to have the same sequence length, if needed
+4. Performs additional truncation for sequence lengths to be divisible by tensor parallelism world size (if enabled)
+5. Returns a dictionary with input and label tensors for training
+
+.. code-block:: python
+
+    def prepare_inputs(self, batch):
+        # Process source audio through perception
+        audio_embs, audio_emb_lens = self.perception(
+            input_signal=batch["source_audio"], 
+            input_signal_length=batch["source_audio_lens"]
+        )
+        
+        # Process target audio through codec (no gradients)
+        with torch.no_grad():
+            target_audio_tokens, target_audio_token_lens = self.audio_codec.encode(batch["target_audio"], batch["target_audio_lens"])
+        
+        # Truncate sequences if needed
+        # ... (truncation logic)
+
+        # Embed text tokens and combine them with audio representations
+        # ... (text embedding logic)
+
+        # Return processed inputs and labels
+        return {
+            "audio_embeds": audio_embs,
+            "input_embeds": input_embeds,
+            "attention_mask": attention_mask,
+            "target_text_ids": target_text_ids,
+            "target_audio_ids": target_audio_ids,
+        }
+
+Validation
+^^^^^^^^^^
+
+The validation process:
+
+1. Clears GPU memory to avoid OOM issues
+2. Loads a scoring ASR model into GPU for evaluation
+3. Initializes metric aggregation for each dataset
+4. Processes each validation dataset separately
+5. Computes BLEU scores for text and ASR-decoded audio
+6. Logs and clears metrics after validation is complete
+
+Scaling Support
+---------------
+
+The DuplexS2SModel includes a `configure_model` method that sets up model parallelism for large-scale training. This method:
+
+1. Detects the parallelism strategy from the trainer's device mesh
+2. Applies Fully Sharded Data Parallel (FSDP) sharding to appropriate modules
+3. Applies Tensor Parallelism (TP) and Sequence Parallelism (SP) when configured
+4. Handles model-specific adaptations for different LLM architectures
+
+The scaling approach supports:
+
+* Pure FSDP2 for distributing parameters across GPUs
+* Pure TP/SP for splitting computation across GPUs
+* 2D parallelism combining both approaches
+
+Pretrained Model Usage
+----------------------
+
+All models in the speechlm2 collection can be instantiated from pretrained checkpoints:
+
+.. code-block:: python
+
+    import nemo.collections.speechlm2 as slm
+    
+    # Load SALM model
+    salm_model = slm.models.SALM.from_pretrained("path/to/checkpoint")
+    
+    # Load SALMAutomodel
+    salm_automodel = slm.models.SALMAutomodel.from_pretrained("path/to/checkpoint")
+
+    # Load DuplexS2SModel
+    duplex_model = slm.models.DuplexS2SModel.from_pretrained("path/to/checkpoint")
+
+    # Load DuplexS2SSpeechDecoderModel
+    speech_decoder_model = slm.models.DuplexS2SSpeechDecoderModel.from_pretrained("path/to/checkpoint")
+
+    # Load DuplexSTTModel
+    stt_model = slm.models.DuplexSTTModel.from_pretrained("path/to/checkpoint")
+
+    # Load DuplexEARTTS
+    ear_tts_model = slm.models.DuplexEARTTS.from_pretrained("path/to/checkpoint")
+
+    # Load NemotronVoiceChat (Inference Only)
+    voicechat_model = slm.models.NemotronVoiceChat.from_pretrained("path/to/checkpoint")
+
+Fine-Tuning from a Checkpoint
+------------------------------
+
+To fine-tune a model starting from a previous training checkpoint (with a fresh
+optimizer and step counter), set ``model.init_from_checkpoint`` in the config:
+
+.. code-block:: python
+
+    # Via Hydra override:
+    # ++model.init_from_checkpoint=/path/to/step=6375.ckpt
+
+This loads only model weights from the checkpoint — optimizer state, LR scheduler,
+and training step are discarded. Supports DCP directories (from FSDP2/TP training),
+HuggingFace directories, and single-file ``.ckpt`` files. Works with both ``SALM``
+and ``SALMAutomodel``.
+
+See :ref:`fine-tuning-from-checkpoint` in the configuration documentation for full
+details.
+
+Model Configuration
+-------------------
+
+All models in this collection use a configuration-based approach, where a YAML configuration file specifies the model architecture, components, and training parameters. See the :doc:`configurations documentation <configs>` for details on these configuration files.
+
+For information about scaling and training these models at scale, see the :doc:`training and scaling documentation <training_and_scaling>`.

docs/source/speechlm2/training_and_scaling.rst

@@ -0,0 +1,332 @@
+Training and Scaling
+====================
+
+This page provides detailed information on training speechlm2 models, including setup requirements, running experiments at scale, debugging, and parallelism strategies.
+
+Running Experiments
+-------------------
+
+The speechlm2 collection includes several scripts to facilitate running experiments, especially on SLURM-based clusters.
+
+SLURM Job Submission
+^^^^^^^^^^^^^^^^^^^^
+
+For training on SLURM clusters, use the following workflow:
+
+.. code-block:: bash
+
+    # Submit 8 consecutive jobs with random seeds
+    scripts/speechlm2/auto_launcher_with_seed.sh -n8 s2s_tinyllama_repro.sub
+
+The ``auto_launcher_with_seed.sh`` script:
+
+1. Generates a random seed for each submitted job
+2. Leverages ``shard_seed="randomized"`` in Lhotse to ensure each data parallel rank is seeded differently
+3. Ensures each tensor parallel rank is seeded identically
+
+SLURM Submission Script
+^^^^^^^^^^^^^^^^^^^^^^^
+
+Example ``s2s_tinyllama_repro.sub`` script:
+
+.. code-block:: bash
+
+    #!/bin/bash
+    #SBATCH --job-name=s2s_training
+    #SBATCH --nodes=4
+    #SBATCH --ntasks-per-node=8
+    #SBATCH --gres=gpu:8
+    #SBATCH --time=24:00:00
+    #SBATCH --exclusive
+    #SBATCH --output=s2s_tinyllama_repro_%j.out
+
+    # Check that the global random seed base is provided
+    if [ -z "$1" ]; then
+      echo "Usage: $0 <global_random_seed_base>"
+      exit 1
+    fi
+    SEED=${1}
+
+    EXP_NAME="s2s_training"
+    RESULTS_DIR="results/${EXP_NAME}"
+
+    srun --ntasks=${SLURM_NTASKS} --ntasks-per-node=${SLURM_NTASKS_PER_NODE} \
+      python -u examples/speechlm2/s2s_duplex_train.py \
+      --config-path=/path/to/config/dir \
+      --config-name=s2s_training.yaml \
+      exp_manager.name=${EXP_NAME} \
+      exp_manager.wandb_logger_kwargs.name=${EXP_NAME} \
+      trainer.num_nodes=$SLURM_JOB_NUM_NODES \
+      exp_manager.explicit_log_dir=${RESULTS_DIR} \
+      data.train_ds.seed=$SEED \
+      data.validation_ds.seed=$SEED
+
+
+Configuration Files
+^^^^^^^^^^^^^^^^^^^
+
+The main configuration file (``s2s_training.yaml``) contains all model, training, and data parameters. See :doc:`configs` for more details. It's recommended to copy and modify this file rather than overriding options in the SLURM script to maintain versioning and configuration clarity.
+
+Debugging
+---------
+
+Running Locally with torchrun
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+For local debugging and profiling, use ``torchrun``:
+
+.. code-block:: bash
+
+    # Run with 4 GPUs locally
+    torchrun --nproc_per_node=4 examples/speechlm2/s2s_duplex_train.py \
+      --config-path=/path/to/config/dir \
+      --config-name=s2s_training.yaml
+
+Scaling Strategies
+------------------
+
+The speechlm2 collection includes support for model parallelism to scale training to large models across multiple GPUs.
+
+Model Parallel Strategies
+^^^^^^^^^^^^^^^^^^^^^^^^^
+
+The collection supports multiple parallelism strategies:
+
+1. **Fully Sharded Data Parallel (FSDP2)**: Distributes model parameters across GPUs
+2. **Tensor Parallelism (TP)**: Splits individual tensors across GPUs
+3. **Sequence Parallelism (SP)**: Splits sequence processing across GPUs
+4. **2D Parallelism**: Combination of FSDP2 with TP/SP
+
+AutomodelParallelStrategy (SALMAutomodel)
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+For ``SALMAutomodel``, the collection provides ``AutomodelParallelStrategy`` which
+delegates device mesh creation and parallelism to NeMo Automodel. This strategy
+supports FSDP2, TP, PP, CP, EP (MoE), and HSDP.
+
+.. code-block:: yaml
+
+    trainer:
+      strategy:
+        _target_: nemo.collections.speechlm2.parts.parallel.AutomodelParallelStrategy
+        dp_size: null       # inferred from world_size / other dims
+        dp_replicate_size: 1  # HSDP replication group size
+        tp_size: 1
+        pp_size: 1
+        cp_size: 1
+        ep_size: 8          # Expert parallelism for MoE models
+
+        # Activation checkpointing — two independent knobs:
+        activation_checkpointing_llm: false         # LLM transformer blocks
+        activation_checkpointing_perception: false  # speech encoder layers
+
+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).
+
+The speech encoder / perception module currently only supports FSDP2 (controlled via ``dp_size``).
+
+Activation Checkpointing
+""""""""""""""""""""""""
+
+``AutomodelParallelStrategy`` exposes two independent activation-checkpointing knobs:
+
+* ``activation_checkpointing_llm`` — single switch covering both the non-EP FSDP2
+  path (forces ``FSDP2Config.activation_checkpointing=True``) and the EP/MoE
+  parallelizer path (passed through as a separate runtime arg). Use this for
+  MoE LLMs whether ``ep_size`` is 1 or larger.
+* ``activation_checkpointing_perception`` — wraps each transformer layer in
+  ``perception.encoder.layers`` (and the Conformer ``pre_encode`` front-end when
+  it isn't a bare ``nn.Linear``) with ``checkpoint_wrapper`` *before* FSDP2
+  sharding. Implemented in ``AudioPerceptionModule.set_activation_checkpointing``.
+
+Both default to ``false``. Toggle them independently to trade compute for
+memory at either end of the model. They are SALMAutomodel-specific knobs (the
+HF Transformers SALM path uses HuggingFace's own gradient-checkpointing API).
+
+.. note::
+   Expert Parallelism (EP) reuses the FSDP2 data-parallel axis (``dp_size``).
+   Dense layers are sharded via FSDP2, while MoE expert layers use EP for
+   all-to-all expert routing — both operate on the same set of GPUs.
+   Setting ``ep_size`` controls how many GPUs participate in expert routing;
+   it does not add a separate dimension.
+
+Training with MoE LLM Backbones
+""""""""""""""""""""""""""""""""
+
+SALMAutomodel enables efficient training of Speech LLMs with Mixture-of-Experts
+backbones like `NVIDIA Nemotron Nano V3 <https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16>`_
+(30B total parameters, 3B active). NeMo Automodel provides two key MoE
+optimizations:
+
+* **Grouped GEMM**: Fuses all expert computations within a single MoE layer into
+  one batched matrix multiplication, maximizing GPU utilization and throughput.
+* **DeepEP** (Deep Expert Parallelism): An efficient all-to-all communication
+  primitive for routing tokens to experts across GPUs, significantly reducing the
+  communication overhead of Expert Parallelism.
+
+Example: training SALMAutomodel with Nemotron Nano V3 on 8 GPUs with EP=8:
+
+.. code-block:: bash
+
+    torchrun --nproc_per_node=8 examples/speechlm2/salm_train.py \
+      --config-name=salm_automodel \
+      model.pretrained_llm=nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 \
+      trainer.strategy.ep_size=8
+
+For distributed inference, launch with ``torchrun``:
+
+.. code-block:: bash
+
+    torchrun --nproc_per_node=8 examples/speechlm2/salm_eval.py \
+      pretrained_name=path/to/checkpoint \
+      inputs=path/to/manifest \
+      ep_size=2
+
+Packed Sequences (THD)
+""""""""""""""""""""""
+
+``SALMAutomodel`` supports an opt-in packed-sequence (``THD``) training and
+validation path that concatenates per-utterance text + audio embeddings into
+a single flat ``[T_total, H]`` sequence with a ``cu_seqlens`` index, instead
+of right-padding into the standard ``[B, T_max, H]`` (``BSHD``) layout. TE's
+varlen FlashAttention then operates segment-by-segment without ever attending
+across utterances, and Mamba's ``seq_idx`` is derived from the same
+``cu_seqlens`` so SSM state resets at document boundaries.
+
+For variable-length speech batches the padding overhead is substantial — the
+``BSHD`` layout pays ``B * (T_max - T_avg)`` wasted compute per minibatch,
+``THD`` pays only the per-utterance rounding to a multiple of ``2*cp_size``
+(needed for TE's CP DualChunkSwap pattern). Throughput improvement scales
+with the variance of utterance lengths.
+
+Enable per-batch:
+
+.. code-block:: yaml
+
+    model:
+      packed_sequences: true   # opt-in; default false (BSHD)
+      automodel_backend:
+        attn: te                # THD path requires TE attention
+
+When ``packed_sequences`` is unset, the existing BSHD path is used unchanged.
+Generate / inference always uses BSHD (it doesn't go through ``prepare_inputs``).
+
+Context Parallelism (CP)
+""""""""""""""""""""""""
+
+``SALMAutomodel`` supports context parallelism for long-audio training on
+hybrid Mamba/attention LLMs (e.g. Nemotron-V3). CP shards the sequence
+dimension across GPUs so per-rank activations and KV-cache memory scale as
+``T / cp_size`` instead of ``T``; attention layers go through TE's
+DualChunkSwap pattern and Mamba mixers go through hidden-parallel
+all-to-all (``MambaContextParallel`` in NeMo Automodel).
+
+Enable via the strategy:
+
+.. code-block:: yaml
+
+    trainer:
+      strategy:
+        _target_: nemo.collections.speechlm2.parts.parallel.AutomodelParallelStrategy
+        cp_size: 2          # context parallel size; must divide num_heads of every Mamba block
+        ep_size: 2          # may share the same ranks as CP
+
+**The THD packed-sequence path is the only supported configuration under
+CP.** Each utterance is its own attention segment and the per-utterance
+sequence rounding aligns naturally with CP's ``2*cp_size`` requirement.
+
+.. warning::
+   **BSHD + CP is not supported.** TE's fused-attention CP path supports
+   ``causal`` but not ``padding_causal``, so the right-pad mask must be
+   dropped before the LLM. With the mask dropped, pad K/V leak into
+   real-token attention through the causal mask and the gradient through
+   the LoRA / projection parameters becomes ``NaN`` after the first
+   optimizer step (validated empirically: BSHD + CP=2 + EP=2 on a 2-GPU
+   run produces ``loss=4.62`` at step 1 then ``loss=nan`` from step 2
+   onwards). This is independent of the TE/cuDNN backward issue
+   documented below — setting ``NVTE_FUSED_ATTN=0`` does not fix it.
+   Set ``model.packed_sequences: true`` to use the THD path instead.
+
+.. note::
+   **CP-safe data loading is automatic.** The speechlm2 datamodule wraps
+   the Lhotse loader in
+   :class:`~nemo.collections.common.data.lhotse.broadcasting.BroadcastingDataLoader`,
+   so under CP/TP every batch is constructed once on the DP source rank
+   (``cp_rank == 0`` and ``tp_rank == 0``) and broadcast to its sub-mesh
+   peers. This eliminates per-rank Lhotse non-determinism (``concurrent_bucketing``,
+   worker scheduling jitter, etc.) as a source of NCCL deadlocks under CP.
+   See :doc:`/dataloaders` for the standalone API.
+
+.. note::
+   **TE/THD exploding-gradients workaround on some GPUs.** On certain GPU
+   architectures (notably Blackwell ``sm_120``), the cuDNN backend that
+   TransformerEngine 2.14 picks for ``qkv_format="thd"`` with
+   ``attn_mask_type="padding_causal"`` returns correct forward activations
+   but gradients amplified 8×–960× per layer. Compounded across the LLM's
+   attention stack this drives gradients to ``1e22``-magnitudes at step 0,
+   the gradient-clip-by-norm computes ``1.0 / inf = 0``, and Adam's moments
+   eventually NaN. Force TE to dispatch FlashAttention instead of cuDNN by
+   setting ``NVTE_FUSED_ATTN=0`` in the launcher environment (requires
+   ``flash-attn`` to be installed for your GPU arch). The FlashAttention
+   THD/``padding_causal`` backward is gradient-correct on the same shapes.
+
+To configure parallelism, modify the ``trainer.strategy`` section in your YAML config:
+
+.. code-block:: yaml
+
+    trainer:
+      strategy:
+        _target_: nemo.core.ModelParallelStrategy
+        find_unused_parameters: False
+        data_parallel: 1   # World size for data parallelism (FSDP2)
+        tensor_parallel: 8  # World size for tensor parallelism
+      devices: 8
+      num_nodes: 1
+      accelerator: gpu
+      precision: bf16-true
+
+The model's ``configure_model`` method automatically sets up the appropriate parallelization based on this configuration.
+
+FSDP2 Configuration (HF Automodel)
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+For Fully Sharded Data Parallel training:
+
+1. Set ``data_parallel`` to the number of GPUs you want to use for data parallelism
+2. Set ``tensor_parallel`` to 1 (disabled)
+
+FSDP2 shards the model parameters across GPUs, all-gathers them for forward/backward passes, and then de-allocates after computation. This allows training of larger models with limited GPU memory.
+See `PyTorch FSDP2 <https://pytorch.org/docs/stable/distributed.fsdp.fully_shard.html>`_ for more details.
+
+Tensor Parallelism Configuration (HF Automodel)
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+For Tensor Parallelism:
+
+1. Set ``tensor_parallel`` to the number of GPUs you want to use for tensor parallelism
+2. Set ``data_parallel`` to 1 (or higher for 2D parallelism)
+
+The ``parallelize_module`` function applies a parallelization plan to specific model components, like splitting attention heads or embedding dimensions across GPUs.
+See `PyTorch TP <https://pytorch.org/docs/stable/distributed.tensor.parallel.html>`_ for more details.
+
+Implementation Details
+----------------------
+
+The core implementation of model parallelism is in the ``configure_model`` method of the model classes. Key aspects include:
+
+1. **Module Sharding**: Calling ``fully_shard`` on modules to distribute parameters across data parallel ranks
+2. **Parallelization Plans**: Creating and applying plans that specify how different layers should be parallelized
+3. **Model-Specific Adaptations**: Handling architectural differences between different LLMs
+
+Advanced Usage
+--------------
+
+Script Customization
+^^^^^^^^^^^^^^^^^^^^
+
+When customizing the training scripts, keep these points in mind:
+
+1. **Path Overrides**: Override paths in the YAML configuration files with your own, as needed
+2. **W&B Keys**: Update Weights & Biases API keys in configuration files
+3. **Batch Size Tuning**: Adjust batch size based on your GPU memory and model size