chore: import upstream snapshot with attribution

examples/speechlm2/conf/duplex_eartts.yaml

@@ -0,0 +1,202 @@
+model:
+  pretrained_lm_name: "nvidia/NVIDIA-Nemotron-Nano-9B-v2"
+  pretrained_audio_codec: ???  # to be released
+  pretrained_tts_model: null
+  scoring_asr: stt_en_fastconformer_transducer_large # used only in validation/evaluation
+  trust_remote_code: false
+
+  # Regexp (re.compile) patterns matching parameters to be frozen.
+  freeze_params:
+    - "^audio_codec\\..+$" # Keep audio codec frozen as it only provides supervision for training.
+    - "^embed_tokens\\..+$" # Keep embed_tokens frozen as done in eartts
+
+  prevent_freeze_params: []  # Use to make specific submodules trainable; overrides freeze_params
+
+  # set custom text eos/bos/pad tokens
+  bos_token: "<s>"
+  eos_token: "</s>"
+  pad_token: "<SPECIAL_12>"
+
+  # inference params
+  inference_guidance_scale: 0.5
+  inference_noise_scale: 0.8
+  inference_top_p_or_k: 0.8
+  inference_guidance_enabled: true
+
+
+  optimizer:
+    _target_: torch.optim.AdamW
+    lr: 4e-05
+    betas: [0.9, 0.98]
+    weight_decay: 0
+    foreach: true # set to false if having issues with tensor-parallelism
+
+  lr_scheduler:
+    _target_: nemo.core.optim.lr_scheduler.InverseSquareRootAnnealing
+    warmup_steps: 2500
+    min_lr: 1e-6
+    max_steps: ${trainer.max_steps}
+
+  codec_config:
+    latent_size: 512
+    n_fft: 16
+    hop_length: 4
+    base_hidden_size: 384
+    channel_mult:
+      - 1
+      - 2
+      - 4
+    rates:
+      - 7
+      - 7
+      - 9
+    num_blocks: 3
+    kernel_size: 7
+    groups: 1
+    codebook_size: 1024
+    num_quantizers: 31
+    wav_to_token_ratio: 1764
+
+  tts_config:
+    # extra configs added
+    use_gated_fusion_for_text_audio: true
+    disable_eos_prediction: true # disable eos prediction
+    use_bos_eos_emb: true
+    use_subword_flag_emb: true
+    num_delay_speech_tokens: 2
+    # EAR-TTS configs
+    backbone_type: gemma3_text
+    backbone_model_class: null
+    backbone_config_class: null
+    backbone_config:
+      hidden_size: 1152
+      intermediate_size: 4608
+      num_hidden_layers: 28
+      num_attention_heads: 16
+      num_key_value_heads: 16
+      head_dim: 72
+      attention_dropout: 0.1
+      use_cache: false
+    latent_size: 512
+    codebook_size: 1024
+    num_quantizers: 31
+    context_hidden_size: null
+    cas_config:
+      backbone_type: t5gemma
+      backbone_model_class: null
+      backbone_config_class: null
+      backbone_config:
+        is_encoder_decoder: false
+        encoder:
+          hidden_size: 1152
+          intermediate_size: 4608
+          num_hidden_layers: 1
+          num_attention_heads: 16
+          num_key_value_heads: 16
+          head_dim: 72
+          use_cache: false
+          attention_dropout: 0.1
+    mog_head_config:
+      intermediate_size: 4608
+      num_layers: 3
+      low_rank: 64
+      num_predictions: 1024
+      min_log_std: -4.0
+      eps: 1e-06
+    p_uncond: 0.1
+    label_smoothing: 0.01
+    max_training_rate: 0.8
+    quantizer_dropout: 0.5
+    random_target_masking: false
+    exponent: 3.0
+trainer:
+  devices: -1
+  accelerator: gpu
+  num_nodes: 1
+  precision: 32
+  logger: False # logger provided by exp_manager
+  enable_checkpointing: False
+  use_distributed_sampler: False
+  max_steps: 1000000
+  val_check_interval: 2000
+  limit_train_batches: ${trainer.val_check_interval}  # an "epoch"
+  limit_val_batches: 2
+  log_every_n_steps: 20
+  num_sanity_val_steps: 0
+  gradient_clip_val: 1.0
+  accumulate_grad_batches: 1
+  strategy:
+    _target_: lightning.pytorch.strategies.DDPStrategy
+    gradient_as_bucket_view: true
+    find_unused_parameters: false
+
+data:
+  # data loader configs
+  add_text_bos_and_eos_in_each_turn: true
+  add_audio_prompt_after_description: true
+  audio_prompt_duration: 3.0
+  frame_length: 0.08
+  source_sample_rate: 22050
+  target_sample_rate: 22050
+  input_roles: ["user", "User"]
+  output_roles: ["agent", "Assistant", "assistant","Agent"]
+
+  train_ds:
+    sample_rate: ${data.target_sample_rate}
+    input_cfg:
+      - type: lhotse_shar
+        shar_path: ???
+    seed: 42
+    shard_seed: "randomized"
+    num_workers: 2
+    batch_size: 4
+    # Optional bucketing:
+    # batch_size: null
+    # 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:
+    # The entries under 'datasets' are a list of separate dataloaders.
+    # The structure is <dataset-name>: {<dataloader-dict-config>}
+    # They inherit all settings from validation_ds, but can individually override them.
+    datasets:
+      val_set_0:  # rename to your dataset name, add more as needed
+        shar_path: ???
+    sample_rate: ${data.target_sample_rate}
+    batch_size: 1
+    seed: 42
+    shard_seed: "randomized"
+
+exp_manager:
+   exp_dir: null
+   explicit_log_dir: duplex_eartts_results/
+   name: eartts
+   create_tensorboard_logger: false
+   create_checkpoint_callback: true
+   use_datetime_version: true
+   max_time_per_run: 00:03:50:00
+
+   resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc.
+   # you need to set these two to True to continue the training
+   resume_if_exists: true
+   resume_ignore_no_checkpoint: true
+
+   # You may use this section to create a W&B logger
+   create_wandb_logger: false
+   wandb_logger_kwargs:
+     name: development-run
+     project: duplex_eartts
+     resume: true
+
+   checkpoint_callback_params:
+     filename: "{step}"
+     monitor: val_asr_bleu
+     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

examples/speechlm2/conf/duplex_stt.yaml

@@ -0,0 +1,122 @@
+model:
+  # Every name/path here starting with 'pretrained' is used to initialize the model weights.
+  pretrained_llm: TinyLlama/TinyLlama_v1.1
+  pretrained_asr: stt_en_fastconformer_hybrid_large_streaming_80ms
+  scoring_asr: stt_en_fastconformer_transducer_large  # used only in validation/evaluation
+
+  pretrained_weights: True  # When False, we use pretrained_name to load the architecture, but with random init
+  trust_remote_code: false
+
+  prevent_freeze_params: []  # Use to make specific submodules trainable; overrides freeze_params
+
+  text_loss_weight: 1.0  # STT model only predicts text, no audio generation
+
+  perception:
+     target:  nemo.collections.speechlm2.modules.perception.AudioPerceptionModule
+     modality_adapter:
+       _target_: nemo.collections.speechlm2.modules.perception.IdentityConnector
+       d_model: 1024
+
+  optimizer:
+    _target_: torch.optim.AdamW
+    lr: 3e-4
+    betas: [0.9, 0.98]
+    weight_decay: 0
+    foreach: true # set to false if having issues with tensor-parallelism
+
+  lr_scheduler:
+    _target_: nemo.core.optim.lr_scheduler.CosineAnnealing
+    warmup_steps: 0  #2500
+    min_lr: 1e-6
+    max_steps: ${trainer.max_steps}
+
+trainer:
+  devices: -1
+  accelerator: gpu
+  num_nodes: 1
+  precision: bf16-true
+  logger: False # logger provided by exp_manager
+  enable_checkpointing: False
+  use_distributed_sampler: False
+  max_steps: 1000000
+  limit_train_batches: 100  # "epoch" size
+  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:
+    # Replace DDPStrategy with ModelParallelStrategy to enable model parallelism
+    _target_: lightning.pytorch.strategies.DDPStrategy
+    gradient_as_bucket_view: true
+    find_unused_parameters: true
+    # _target_: lightning.pytorch.strategies.ModelParallelStrategy
+    # tensor_parallel_size: 1
+    # data_parallel_size: 2
+
+data:
+  frame_length: 0.08
+  source_sample_rate: 16000
+  input_roles: ["user", "User"]
+  output_roles: ["agent", "Assistant"]
+
+  train_ds:
+    sample_rate: ${data.source_sample_rate}
+    input_cfg:
+      - type: lhotse_shar
+        shar_path: ???
+    seed: 42
+    shard_seed: "randomized"
+    num_workers: 2
+    batch_size: 4
+    # Optional bucketing:
+    # batch_size: null
+    # 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:
+    # The entries under 'datasets' are a list of separate dataloaders.
+    # The structure is <dataset-name>: {<dataloader-dict-config>}
+    # They inherit all settings from validation_ds, but can individually override them.
+    datasets:
+      val_set_0:  # rename to your dataset name, add more as needed
+        shar_path: ???
+    sample_rate: ${data.source_sample_rate}
+    batch_size: 1
+    seed: 42
+    shard_seed: "randomized"
+
+exp_manager:
+   exp_dir: null
+   explicit_log_dir: duplex_stt_results/
+   name: speechlm2
+   create_tensorboard_logger: false
+   create_checkpoint_callback: true
+   use_datetime_version: true
+   max_time_per_run: 00:03:50:00
+
+   resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc.
+   # you need to set these two to True to continue the training
+   resume_if_exists: true
+   resume_ignore_no_checkpoint: true
+
+   # You may use this section to create a W&B logger
+   create_wandb_logger: false
+   wandb_logger_kwargs:
+     name: development-run
+     project: speechlm2_duplex_stt
+     resume: true
+
+   checkpoint_callback_params:
+     filename: "{step}"
+     monitor: val_bleu
+     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

examples/speechlm2/conf/nemotron_voicechat.yaml

@@ -0,0 +1,57 @@
+checkpoint_path: null # Path to the pre-trained NemotronVoiceChat checkpoint for evaluation
+model:
+  scoring_asr: stt_en_fastconformer_transducer_large # ASR model used to transcribe generated audio for ASR-BLEU computation
+  inference_speaker_reference: null # Path to an audio file used to clone/condition the TTS voice. Set to "null" if using a preset name below.
+  inference_speaker_name: Megan # Preset speaker identifier. If provided, this overrides `inference_speaker_reference`.
+  
+  stt:
+    model:
+      # evaluation params
+      eval_text_turn_taking: true # Enables evaluation of turn-taking and text prediction accuracy in the Duplex STT model
+
+  speech_generation:
+    model:
+      # inference params for the Duplex EAR-TTS module
+      inference_guidance_scale: 0.2 # Classifier-Free Guidance (CFG) scale for conditioning the audio generation
+      inference_noise_scale: 0.001 # Sampling temperature/noise for MoG
+      inference_top_p_or_k: 0.95 # Nucleus sampling (top-p) or top-k threshold for token selection
+      inference_guidance_enabled: true # Toggle to enable/disable Classifier-Free Guidance
+      inference_force_speech_silence_on_eos: true # Forces the model to output silence tokens once the End-Of-Sequence (EOS) token is generated
+
+trainer:
+  devices: -1 # Number of GPUs to use (-1 uses all available)
+  accelerator: gpu # Hardware accelerator type
+  num_nodes: 1 # Number of compute nodes
+  precision: 32 # Mixed precision setting (16-bit) for faster, memory-efficient inference
+  logger: False # Disabled here because NeMo's `exp_manager` handles logging
+  limit_val_batches: 1.0 # Fraction of the validation dataset to use (1.0 = use the entire dataset)
+  log_every_n_steps: 20 # Frequency of logging metrics to the console/wandb
+  use_distributed_sampler: false # Disable distributed sampler
+  strategy:
+    _target_: lightning.pytorch.strategies.DDPStrategy # Distributed Data Parallel strategy for multi-GPU inference
+    gradient_as_bucket_view: true # Memory optimization for DDP
+    find_unused_parameters: true # Required if parts of the model (like text-only branches) don't receive gradients/usage
+
+data:
+  frame_length: 0.08 # Duration of a single audio frame in seconds (80ms)
+  source_sample_rate: 16000 # Sample rate of the input/user audio prompts (16 kHz)
+  target_sample_rate: 22050 # Sample rate of the generated output speech (22.05 kHz)
+  input_roles: ["user", "User"] # Conversation roles mapped to the input prompt
+  output_roles: ["agent", "Assistant", "assistant","Agent"] # Conversation roles the model is tasked with generating
+
+  validation_ds:
+    datasets:
+      evaluation_set:
+        shar_path: /lustre/fsw/portfolios/llmservice/users/kevinhu/duplex/ultrachat_v2/shar_duplex/manifest_000020 # Path to the Lhotse WebDataset tar shards manifest
+
+    sample_rate: ${data.target_sample_rate} # Audio will be resampled to this rate if necessary
+    batch_size: 4 # Number of samples processed per GPU during evaluation
+    seed: 42 # Random seed for reproducibility
+    shard_seed: "randomized" # Ensures distributed workers get different data shards
+
+exp_manager:
+   explicit_log_dir: nemotron_voicechat_log_dir/ # Root directory where evaluation metrics, JSON logs, and generated audio will be saved
+   name: nemotron-voicechat-eval # Name of the experiment
+   create_tensorboard_logger: false # Toggle for TensorBoard logging
+   create_checkpoint_callback: false # Enables the checkpoint callback module
+   use_datetime_version: true # Appends a timestamp to the log directory name

examples/speechlm2/conf/s2s_duplex.yaml

@@ -0,0 +1,162 @@
+model:
+  # Every name/path here starting with 'pretrained' is used to initialize the model weights.
+  pretrained_llm: TinyLlama/TinyLlama_v1.1
+  pretrained_audio_codec: ???  # to be released
+  pretrained_asr: stt_en_fastconformer_hybrid_large_streaming_80ms
+  scoring_asr: stt_en_fastconformer_transducer_large  # used only in validation/evaluation
+
+  pretrained_weights: True  # When False, we use pretrained_name to load the architecture, but with random init
+
+  # Regexp (re.compile) patterns matching parameters to be frozen.
+  freeze_params:
+    - "^audio_codec\\..+$"  # Keep audio codec frozen as it only provides supervision for training.
+  prevent_freeze_params: []  # Use to make specific submodules trainable; overrides freeze_params
+
+  audio_loss_weight: 4
+  text_loss_weight: 3
+
+  # Note: Uncomment the block below to enable LoRA on LLM via HuggingFace PEFT library.
+  #   It will automatically freeze LLM parameters even if freeze_params was unused,
+  #   and prevent freezing any parameter that has the string '.lora_' in its name.
+  # lora:
+  #   task_type: CAUSAL_LM
+  #   r: 8
+  #   lora_alpha: 32
+  #   lora_dropout: 0.1
+
+  perception:
+     target:  nemo.collections.speechlm2.modules.perception.AudioPerceptionModule
+     modality_adapter:
+       _target_: nemo.collections.asr.modules.ConformerEncoder
+       feat_in: 512
+       feat_out: -1 # you may set it if you need different output size other than the default d_model
+       n_layers: 2
+       d_model: 512
+       subsampling: dw_striding # vggnet, striding, stacking or stacking_norm, dw_striding
+       subsampling_factor: 1 # must be power of 2 for striding and vggnet
+       subsampling_conv_channels: 256 # set to -1 to make it equal to the d_model
+       causal_downsampling: true
+       ff_expansion_factor: 4
+       self_attention_model: rel_pos # rel_pos or abs_pos
+       n_heads: 8 # may need to be lower for smaller d_models
+       # [left, right] specifies the number of steps to be seen from left and right of each step in self-attention
+       att_context_size: [70, 1] # -1 means unlimited context
+       att_context_style: chunked_limited # regular or chunked_limited
+       xscaling: true # scales up the input embeddings by sqrt(d_model)
+       untie_biases: true # unties the biases of the TransformerXL layers
+       pos_emb_max_len: 5000
+       conv_kernel_size: 9
+       conv_norm_type: layer_norm # batch_norm or layer_norm or groupnormN (N specifies the number of groups)
+       # conv_context_size can be"causal" or a list of two integers while conv_context_size[0]+conv_context_size[1]+1==conv_kernel_size
+       # null means [(kernel_size-1)//2, (kernel_size-1)//2], and 'causal' means [(kernel_size-1), 0]
+       conv_context_size: causal
+       ### regularization
+       dropout: 0 # The dropout used in most of the Conformer Modules
+       dropout_pre_encoder: 0 # The dropout used before the encoder
+       dropout_emb: 0.0 # The dropout used for embeddings
+       dropout_att: 0 # The dropout for multi-headed attention modules
+
+  optimizer:
+    _target_: torch.optim.AdamW
+    lr: 3e-4
+    betas: [0.9, 0.98]
+    weight_decay: 0
+    foreach: true # set to false if having issues with tensor-parallelism
+
+  lr_scheduler:
+    _target_: nemo.core.optim.lr_scheduler.CosineAnnealing
+    warmup_steps: 0
+    min_lr: 1e-6
+    max_steps: ${trainer.max_steps}
+
+trainer:
+  devices: -1
+  accelerator: gpu
+  num_nodes: 1
+  precision: bf16-true
+  logger: False # logger provided by exp_manager
+  enable_checkpointing: False
+  use_distributed_sampler: False
+  max_steps: 1000000
+  limit_train_batches: 100  # "epoch" size
+  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:
+    # Replace DDPStrategy with ModelParallelStrategy to enable model parallelism
+    _target_: lightning.pytorch.strategies.DDPStrategy
+    gradient_as_bucket_view: true
+    find_unused_parameters: true
+    # _target_: lightning.pytorch.strategies.ModelParallelStrategy
+    # tensor_parallel_size: 1
+    # data_parallel_size: 2
+
+data:
+  frame_length: 0.08
+  source_sample_rate: 16000
+  target_sample_rate: 22050
+  input_roles: ["user", "User"]
+  output_roles: ["agent", "Assistant"]
+
+  train_ds:
+    sample_rate: ${data.target_sample_rate}
+    input_cfg:
+      - type: lhotse_shar
+        shar_path: ???  # needs to be specified
+    seed: 42
+    shard_seed: "randomized"
+    num_workers: 2
+#    batch_size: 4
+    # Optional bucketing:
+    batch_size: null
+    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:
+    # The entries under 'datasets' are a list of separate dataloaders.
+    # The structure is <dataset-name>: {<dataloader-dict-config>}
+    # They inherit all settings from validation_ds, but can individually override them.
+    datasets:
+      val_set_0:  # rename to your dataset name, add more as needed
+        shar_path: ???  # needs to be specified
+    sample_rate: ${data.target_sample_rate}
+    batch_size: 1
+    seed: 42
+    shard_seed: "randomized"
+
+exp_manager:
+   exp_dir: null
+   explicit_log_dir: s2s_results/
+   name: speechlm2
+   create_tensorboard_logger: false
+   create_checkpoint_callback: true
+   use_datetime_version: true
+   max_time_per_run: 00:03:50:00
+
+   resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc.
+   # you need to set these two to True to continue the training
+   resume_if_exists: true
+   resume_ignore_no_checkpoint: true
+
+   # You may use this section to create a W&B logger
+   create_wandb_logger: false
+   wandb_logger_kwargs:
+     name: development-run
+     project: speechlm2
+     resume: true
+
+   checkpoint_callback_params:
+     filename: "{step}"
+     monitor: val_asr_bleu
+     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

examples/speechlm2/conf/s2s_duplex_speech_decoder.yaml

@@ -0,0 +1,182 @@
+model:
+  # Every name/path here starting with 'pretrained' is used to initialize the model weights.
+  pretrained_llm: TinyLlama/TinyLlama_v1.1
+  pretrained_audio_codec: ???  # to be released
+  pretrained_asr: stt_en_fastconformer_hybrid_large_streaming_80ms
+  scoring_asr: stt_en_fastconformer_transducer_large  # used only in validation/evaluation
+
+  pretrained_weights: True  # When False, we use pretrained_name to load the architecture, but with random init
+
+  # Regexp (re.compile) patterns matching parameters to be frozen.
+  freeze_params:
+    - "^audio_codec\\..+$"  # Keep audio codec frozen as it only provides supervision for training.
+  prevent_freeze_params: []  # Use to make specific submodules trainable; overrides freeze_params
+
+  audio_loss_weight: 4
+  text_loss_weight: 3
+
+  # Note: Uncomment the block below to enable LoRA on LLM via HuggingFace PEFT library.
+  #   It will automatically freeze LLM parameters even if freeze_params was unused,
+  #   and prevent freezing any parameter that has the string '.lora_' in its name.
+  # lora:
+  #   task_type: CAUSAL_LM
+  #   r: 8
+  #   lora_alpha: 32
+  #   lora_dropout: 0.1
+
+  perception:
+     target:  nemo.collections.speechlm2.modules.perception.AudioPerceptionModule
+     modality_adapter:
+       _target_: nemo.collections.asr.modules.ConformerEncoder
+       feat_in: 512
+       feat_out: -1 # you may set it if you need different output size other than the default d_model
+       n_layers: 2
+       d_model: 512
+       subsampling: dw_striding # vggnet, striding, stacking or stacking_norm, dw_striding
+       subsampling_factor: 1 # must be power of 2 for striding and vggnet
+       subsampling_conv_channels: 256 # set to -1 to make it equal to the d_model
+       causal_downsampling: true
+       ff_expansion_factor: 4
+       self_attention_model: rel_pos # rel_pos or abs_pos
+       n_heads: 8 # may need to be lower for smaller d_models
+       # [left, right] specifies the number of steps to be seen from left and right of each step in self-attention
+       att_context_size: [70, 1] # -1 means unlimited context
+       att_context_style: chunked_limited # regular or chunked_limited
+       xscaling: true # scales up the input embeddings by sqrt(d_model)
+       untie_biases: true # unties the biases of the TransformerXL layers
+       pos_emb_max_len: 5000
+       conv_kernel_size: 9
+       conv_norm_type: layer_norm # batch_norm or layer_norm or groupnormN (N specifies the number of groups)
+       # conv_context_size can be"causal" or a list of two integers while conv_context_size[0]+conv_context_size[1]+1==conv_kernel_size
+       # null means [(kernel_size-1)//2, (kernel_size-1)//2], and 'causal' means [(kernel_size-1), 0]
+       conv_context_size: causal
+       ### regularization
+       dropout: 0 # The dropout used in most of the Conformer Modules
+       dropout_pre_encoder: 0 # The dropout used before the encoder
+       dropout_emb: 0.0 # The dropout used for embeddings
+       dropout_att: 0 # The dropout for multi-headed attention modules
+
+  speech_decoder:
+    n_layers: 12
+    d_model: 768
+    d_ffn: 3072
+    sa_n_heads: 12
+    kernel_size: 3
+    p_dropout: 0.1
+    p_dropout_out: 0.0
+    has_xattn: false
+    xa_d_memory: 768
+    xa_n_heads: 12
+    is_causal: true
+    apply_norm_to_cond: true
+    apply_norm_out: true
+    max_length_causal_mask: 5000
+    cond_on_prev_audio_tokens: True
+    detach_input: False
+    use_learnable_pos_emb: True
+
+  optimizer:
+    _target_: torch.optim.AdamW
+    lr: 3e-4
+    betas: [0.9, 0.98]
+    weight_decay: 0
+    foreach: true # set to false if having issues with tensor-parallelism
+
+  lr_scheduler:
+#    _target_: nemo.core.optim.lr_scheduler.InverseSquareRootAnnealing
+    _target_: nemo.core.optim.lr_scheduler.CosineAnnealing
+    warmup_steps: 0  #2500
+    min_lr: 1e-6
+    max_steps: ${trainer.max_steps}
+
+trainer:
+  devices: -1
+  accelerator: gpu
+  num_nodes: 1
+  precision: bf16-true
+  logger: False # logger provided by exp_manager
+  enable_checkpointing: False
+  use_distributed_sampler: False
+  max_steps: 1000000
+  limit_train_batches: 100  # "epoch" size
+  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:
+    # Replace DDPStrategy with ModelParallelStrategy to enable model parallelism
+    _target_: lightning.pytorch.strategies.DDPStrategy
+    gradient_as_bucket_view: true
+    find_unused_parameters: true
+    # _target_: lightning.pytorch.strategies.ModelParallelStrategy
+    # tensor_parallel_size: 1
+    # data_parallel_size: 2
+
+data:
+  frame_length: 0.08
+  source_sample_rate: 16000
+  target_sample_rate: 22050
+  input_roles: ["user", "User"]
+  output_roles: ["agent", "Assistant"]
+
+  train_ds:
+    sample_rate: ${data.target_sample_rate}
+    input_cfg:
+      - type: lhotse_shar
+        shar_path: ???
+    seed: 42
+    shard_seed: "randomized"
+    num_workers: 2
+    batch_size: 4
+    # Optional bucketing:
+    # batch_size: null
+    # 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:
+    # The entries under 'datasets' are a list of separate dataloaders.
+    # The structure is <dataset-name>: {<dataloader-dict-config>}
+    # They inherit all settings from validation_ds, but can individually override them.
+    datasets:
+      val_set_0:  # rename to your dataset name, add more as needed
+        shar_path: ???
+    sample_rate: ${data.target_sample_rate}
+    batch_size: 1
+    seed: 42
+    shard_seed: "randomized"
+
+exp_manager:
+   exp_dir: null
+   explicit_log_dir: s2s_sdv2_results/
+   name: speechlm2
+   create_tensorboard_logger: false
+   create_checkpoint_callback: true
+   use_datetime_version: true
+   max_time_per_run: 00:03:50:00
+
+   resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc.
+   # you need to set these two to True to continue the training
+   resume_if_exists: true
+   resume_ignore_no_checkpoint: true
+
+   # You may use this section to create a W&B logger
+   create_wandb_logger: false
+   wandb_logger_kwargs:
+     name: development-run
+     project: speechlm2_speech_decoder
+     resume: true
+
+   checkpoint_callback_params:
+     filename: "{step}"
+     monitor: val_asr_bleu
+     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

examples/speechlm2/conf/salm.yaml

@@ -0,0 +1,184 @@
+model:
+  # Every name/path here starting with 'pretrained' is used to initialize the model weights.
+  pretrained_llm: Qwen/Qwen3-1.7B
+  pretrained_asr: nvidia/canary-1b-flash
+
+  pretrained_weights: True  # When False, we use pretrained_name to load the architecture, but with random init
+
+  # 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. Set to null to train from pretrained_llm/asr.
+  init_from_checkpoint: null
+
+  # Regexp (re.compile) patterns matching parameters to be frozen.
+  freeze_params:
+    # Frozen LLM
+    - "^llm\\..+$"  # LLM
+    - "^embed_tokens\\..+$"  # LLM embedding is moved
+    # Frozen pretrained ASR (only the modality adapter layers are trainable)
+    - "^perception\\.preprocessor\\..+$"
+    - "^perception\\.encoder\\..+$"
+  prevent_freeze_params: []  # Use to make specific submodules trainable; overrides freeze_params
+
+  prompt_format: qwen
+  audio_locator_tag: "<|audioplaceholder|>"  # placeholder token for audio turn is expected
+  # Optional: split audios longer than this before the speech encoder forward, then
+  # concatenate the encoded chunks back into one sequence before passing them to the LLM.
+  # Leave as null (default) to encode each audio row directly and preserve existing behavior.
+  encoder_chunk_size_seconds: null
+
+  # Note: Uncomment the block below to enable LoRA on LLM via HuggingFace PEFT library.
+  #   It will automatically freeze LLM parameters even if freeze_params was unused,
+  #   and prevent freezing any parameter that has the string '.lora_' in its name.
+  # lora:
+  #   task_type: CAUSAL_LM
+  #   r: 128
+  #   lora_alpha: 256
+  #   lora_dropout: 0.01
+  #   # target_modules are only necessary if the `pretrained_llm` is not yet registered in PEFT library
+  #   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
+    
+    # Optional Rotary Time Embedding (RoTE, https://arxiv.org/abs/2410.12109): encodes absolute time into the encoder output
+    # features (at the modality-adapter entrance). Omit this block to disable (default).
+    # rote:
+    #   _target_: nemo.collections.speechlm2.modules.rote.RotaryTimeEmbedding
+    #   dim: 1024             # must match the encoder output feature dim (must be even)
+    #   theta: 1200.0         # base of the geometric frequency bank
+    #   rotary_fraction: 0.2  # fraction of dim to rotate
+    
+    # spec_augment:
+    #   _target_: nemo.collections.asr.modules.SpectrogramAugmentation
+    #   freq_masks: 2 # set to zero to disable it
+    #   time_masks: 10 # set to zero to disable it
+    #   freq_width: 27
+    #   time_width: 5  # 5 frames = 50ms
+
+  optimizer:
+    _target_: torch.optim.AdamW
+    lr: 5e-4
+    betas: [0.9, 0.98]
+    weight_decay: 1e-3
+    foreach: true # set to false if having issues with tensor-parallelism
+
+  lr_scheduler:
+    _target_: nemo.core.optim.lr_scheduler.CosineAnnealing
+    warmup_steps: 1000
+    min_lr: 1e-6
+    max_steps: ${trainer.max_steps}
+
+trainer:
+  devices: -1
+  accelerator: gpu
+  num_nodes: 1
+  precision: bf16-true
+  logger: False # logger provided by exp_manager
+  enable_checkpointing: False
+  use_distributed_sampler: False
+  max_steps: 100000
+  limit_train_batches: 5000  # "epoch" size
+  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:
+    # Replace DDPStrategy with ModelParallelStrategy to enable model parallelism
+    _target_: lightning.pytorch.strategies.DDPStrategy
+    gradient_as_bucket_view: true
+    find_unused_parameters: true
+    # _target_: lightning.pytorch.strategies.ModelParallelStrategy
+    # tensor_parallel_size: 1
+    # data_parallel_size: 8  # This is FSDP2
+
+data:
+  train_ds:
+    sample_rate: 16000
+    prompt_format: ${model.prompt_format}
+    token_equivalent_duration: 0.08
+    input_cfg:
+      - type: lhotse_as_conversation
+        cuts_path: ???  # needs to be set
+        audio_locator_tag: ${model.audio_locator_tag}
+        tags:
+          # Uncomment below line to include a system prompt for supported models (e.g. Llama; not Qwen).
+          # system_prompt: "some system prompt"
+          context: "Repeat after me, typing in lowercase."
+    seed: 42
+    shuffle: true
+    shard_seed: "randomized"
+    num_workers: 1
+    batch_size: 4
+    # Optional bucketing:
+    # batch_size: null
+    # use_bucketing: true
+    # use_multimodal_sampling: true
+    # measure_total_length: true
+    # Note: `batch_tokens`, `bucket_duration_bins`, and `max_tokens` all represent tokens as
+    #  the sum of input audio frames and output text tokens. Number of audio frames is
+    #  calculated using `token_equivalent_duration`.
+    # batch_tokens: 4000
+    # max_tokens: 2048
+    # bucket_duration_bins: [64, 128, 256, 384, 512, 768, 1024, 1280, 1536, 2048]
+    # num_buckets: 10
+    # bucket_buffer_size: 5000
+
+  validation_ds:
+    # The entries under 'datasets' are a list of separate dataloaders.
+    # The structure is <dataset-name>: {<dataloader-dict-config>}
+    # They inherit all settings from validation_ds, but can individually override them.
+    prompt_format: ${model.prompt_format}
+    token_equivalent_duration: 0.08
+    datasets:
+      val_set_0:  # rename to your dataset name, add more as needed
+        input_cfg:
+          - type: lhotse_as_conversation
+            cuts_path: ???  # needs to be set
+            audio_locator_tag: ${model.audio_locator_tag}
+            tags:
+              # Uncomment below line to include a system prompt for supported models (e.g. Llama; not Qwen).
+              # system_prompt: "some system prompt"
+              context: "Repeat after me, typing in lowercase."
+    sample_rate: 16000
+    batch_size: 1
+    seed: 42
+    shard_seed: "randomized"
+
+exp_manager:
+   exp_dir: null
+   explicit_log_dir: salm_results/
+   name: salm
+   create_tensorboard_logger: false
+   create_checkpoint_callback: true
+   use_datetime_version: true
+   max_time_per_run: 00:03:50:00
+
+   resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc.
+   # you need to set these two to True to continue the training
+   resume_if_exists: true
+   resume_ignore_no_checkpoint: true
+
+   # You may use this section to create a W&B logger
+   create_wandb_logger: false
+   wandb_logger_kwargs:
+     name: development-run
+     project: salm
+     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

examples/speechlm2/conf/salm_asr_decoder.yaml

@@ -0,0 +1,165 @@
+model:
+  # Every name/path here starting with 'pretrained' is used to initialize the model weights.
+  pretrained_llm: Qwen/Qwen3-1.7B
+  pretrained_asr: nvidia/canary-1b-flash
+
+  pretrained_weights: True  # When False, we use pretrained_name to load the architecture, but with random init
+
+  # Regexp (re.compile) patterns matching parameters to be frozen.
+  freeze_params:
+    # Frozen LLM
+    - "^llm\\..+$"  # LLM
+    - "^embed_tokens\\..+$"  # LLM embedding is moved
+    # Frozen pretrained ASR (only the modality adapter layers are trainable)
+    - "^perception\\.preprocessor\\..+$"
+    - "^perception\\.encoder\\..+$"
+  prevent_freeze_params: []  # Use to make specific submodules trainable; overrides freeze_params
+
+  prompt_format: qwen
+  audio_locator_tag: "<|audioplaceholder|>"  # placeholder token for audio turn is expected
+
+  # Note: Uncomment the block below to enable LoRA on LLM via HuggingFace PEFT library.
+  #   It will automatically freeze LLM parameters even if freeze_params was unused,
+  #   and prevent freezing any parameter that has the string '.lora_' in its name.
+  # lora:
+  #   task_type: CAUSAL_LM
+  #   r: 128
+  #   lora_alpha: 256
+  #   lora_dropout: 0.01
+  #   # target_modules are only necessary if the `pretrained_llm` is not yet registered in PEFT library
+  #   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
+     # spec_augment:
+     #   _target_: nemo.collections.asr.modules.SpectrogramAugmentation
+     #   freq_masks: 2 # set to zero to disable it
+     #   time_masks: 10 # set to zero to disable it
+     #   freq_width: 27
+     #   time_width: 5  # 5 frames = 50ms
+
+  optimizer:
+    _target_: torch.optim.AdamW
+    lr: 5e-4
+    betas: [0.9, 0.98]
+    weight_decay: 1e-3
+    foreach: true # set to false if having issues with tensor-parallelism
+
+  lr_scheduler:
+    _target_: nemo.core.optim.lr_scheduler.CosineAnnealing
+    warmup_steps: 1000
+    min_lr: 1e-6
+    max_steps: ${trainer.max_steps}
+
+trainer:
+  devices: -1
+  accelerator: gpu
+  num_nodes: 1
+  precision: bf16-true
+  logger: False # logger provided by exp_manager
+  enable_checkpointing: False
+  use_distributed_sampler: False
+  max_steps: 100000
+  limit_train_batches: 5000  # "epoch" size
+  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:
+    # Replace DDPStrategy with ModelParallelStrategy to enable model parallelism
+    _target_: lightning.pytorch.strategies.DDPStrategy
+    gradient_as_bucket_view: true
+    find_unused_parameters: true
+    # _target_: lightning.pytorch.strategies.ModelParallelStrategy
+    # tensor_parallel_size: 1
+    # data_parallel_size: 8  # This is FSDP2
+
+data:
+  train_ds:
+    sample_rate: 16000
+    prompt_format: ${model.prompt_format}
+    token_equivalent_duration: 0.08
+    input_cfg:
+      - type: lhotse_as_conversation
+        cuts_path: ???  # needs to be set
+        audio_locator_tag: ${model.audio_locator_tag}
+        tags:
+          # Uncomment below line to include a system prompt for supported models (e.g. Llama; not Qwen).
+          # system_prompt: "some system prompt"
+          context: "Repeat after me, typing in lowercase."
+    seed: 42
+    shuffle: true
+    shard_seed: "randomized"
+    num_workers: 1
+    batch_size: 4
+    # Optional bucketing:
+    # batch_size: null
+    # use_bucketing: true
+    # use_multimodal_sampling: true
+    # measure_total_length: true
+    # Note: `batch_tokens`, `bucket_duration_bins`, and `max_tokens` all represent tokens as
+    #  the sum of input audio frames and output text tokens. Number of audio frames is
+    #  calculated using `token_equivalent_duration`.
+    # batch_tokens: 4000
+    # max_tokens: 2048
+    # bucket_duration_bins: [64, 128, 256, 384, 512, 768, 1024, 1280, 1536, 2048]
+    # num_buckets: 10
+    # bucket_buffer_size: 5000
+
+  validation_ds:
+    # The entries under 'datasets' are a list of separate dataloaders.
+    # The structure is <dataset-name>: {<dataloader-dict-config>}
+    # They inherit all settings from validation_ds, but can individually override them.
+    prompt_format: ${model.prompt_format}
+    token_equivalent_duration: 0.08
+    datasets:
+      val_set_0:  # rename to your dataset name, add more as needed
+        input_cfg:
+          - type: lhotse_as_conversation
+            cuts_path: ???  # needs to be set
+            audio_locator_tag: ${model.audio_locator_tag}
+            tags:
+              # Uncomment below line to include a system prompt for supported models (e.g. Llama; not Qwen).
+              # system_prompt: "some system prompt"
+              context: "Repeat after me, typing in lowercase."
+    sample_rate: 16000
+    batch_size: 1
+    seed: 42
+    shard_seed: "randomized"
+
+exp_manager:
+   exp_dir: null
+   explicit_log_dir: salm_results/
+   name: salm
+   create_tensorboard_logger: false
+   create_checkpoint_callback: true
+   use_datetime_version: true
+   max_time_per_run: 00:03:50:00
+
+   resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc.
+   # you need to set these two to True to continue the training
+   resume_if_exists: true
+   resume_ignore_no_checkpoint: true
+
+   # You may use this section to create a W&B logger
+   create_wandb_logger: false
+   wandb_logger_kwargs:
+     name: development-run
+     project: salm
+     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

examples/speechlm2/conf/salm_asr_decoder_multilayerproj.yaml

@@ -0,0 +1,159 @@
+model:
+  # Every name/path here starting with 'pretrained' is used to initialize the model weights.
+  pretrained_llm: nvidia/Llama-3.1-Nemotron-Nano-8B-v1
+  pretrained_asr: nvidia/parakeet-tdt-0.6b-v2
+
+  pretrained_weights: True  # When False, we use pretrained_name to load the architecture, but with random init
+
+  # Regexp (re.compile) patterns matching parameters to be frozen.
+  freeze_params:
+    # Frozen LLM
+    - "^llm\\..+$"  # LLM
+    - "^embed_tokens\\..+$"  # LLM embedding is moved
+    # Frozen pretrained ASR (only the modality adapter layers are trainable)
+    - "^perception\\.preprocessor\\..+$"
+    - "^perception\\.encoder\\..+$"
+  prevent_freeze_params: []  # Use to make specific submodules trainable; overrides freeze_params
+
+  prompt_format: llama3
+  audio_locator_tag: "<|audioplaceholder|>"  # placeholder token for audio turn is expected
+
+  # Note: Uncomment the block below to enable LoRA on LLM via HuggingFace PEFT library.
+  #   It will automatically freeze LLM parameters even if freeze_params was unused,
+  #   and prevent freezing any parameter that has the string '.lora_' in its name.
+  # lora:
+  #   task_type: CAUSAL_LM
+  #   r: 128
+  #   lora_alpha: 256
+  #   lora_dropout: 0.01
+
+  perception:
+     target: nemo.collections.speechlm2.modules.perception.AudioTranscriptionPerceptionModule
+     output_dim: 4096
+     modality_adapter:
+       _target_: nemo.collections.asr.modules.QformerConnector
+       target_layer_ids: [0, 5, 11, 17, 23]
+       input_dim: 1024
+       output_dim: 4096
+
+  optimizer:
+    _target_: torch.optim.AdamW
+    lr: 5e-4
+    betas: [0.9, 0.98]
+    weight_decay: 1e-3
+    foreach: true # set to false if having issues with tensor-parallelism
+
+  lr_scheduler:
+    _target_: nemo.core.optim.lr_scheduler.CosineAnnealing
+    warmup_steps: 1000
+    min_lr: 1e-6
+    max_steps: ${trainer.max_steps}
+
+trainer:
+  devices: -1
+  accelerator: gpu
+  num_nodes: 1
+  precision: bf16-true
+  logger: False # logger provided by exp_manager
+  enable_checkpointing: False
+  use_distributed_sampler: False
+  max_steps: 100000
+  limit_train_batches: 5000  # "epoch" size
+  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:
+    # Replace DDPStrategy with ModelParallelStrategy to enable model parallelism
+    _target_: lightning.pytorch.strategies.DDPStrategy
+    gradient_as_bucket_view: true
+    find_unused_parameters: true
+    # _target_: lightning.pytorch.strategies.ModelParallelStrategy
+    # tensor_parallel_size: 1
+    # data_parallel_size: 8  # This is FSDP2
+
+data:
+  train_ds:
+    sample_rate: 16000
+    prompt_format: ${model.prompt_format}
+    token_equivalent_duration: 0.08
+    input_cfg:
+      - type: lhotse_as_conversation
+        cuts_path: ???  # needs to be set
+        audio_locator_tag: ${model.audio_locator_tag}
+        tags:
+          # Uncomment below line to include a system prompt for supported models (e.g. Llama; not Qwen).
+          # system_prompt: "some system prompt"
+          context: "Repeat after me, typing in lowercase."
+    seed: 42
+    shuffle: true
+    shard_seed: "randomized"
+    num_workers: 1
+    batch_size: 4
+    # Optional bucketing:
+    # batch_size: null
+    # use_bucketing: true
+    # use_multimodal_sampling: true
+    # measure_total_length: true
+    # Note: `batch_tokens`, `bucket_duration_bins`, and `max_tokens` all represent tokens as
+    #  the sum of input audio frames and output text tokens. Number of audio frames is
+    #  calculated using `token_equivalent_duration`.
+    # batch_tokens: 4000
+    # max_tokens: 2048
+    # bucket_duration_bins: [64, 128, 256, 384, 512, 768, 1024, 1280, 1536, 2048]
+    # num_buckets: 10
+    # bucket_buffer_size: 5000
+
+  validation_ds:
+    # The entries under 'datasets' are a list of separate dataloaders.
+    # The structure is <dataset-name>: {<dataloader-dict-config>}
+    # They inherit all settings from validation_ds, but can individually override them.
+    prompt_format: ${model.prompt_format}
+    token_equivalent_duration: 0.08
+    datasets:
+      val_set_0:  # rename to your dataset name, add more as needed
+        input_cfg:
+          - type: lhotse_as_conversation
+            cuts_path: ???  # needs to be set
+            audio_locator_tag: ${model.audio_locator_tag}
+            tags:
+              # Uncomment below line to include a system prompt for supported models (e.g. Llama; not Qwen).
+              # system_prompt: "some system prompt"
+              context: "Repeat after me, typing in lowercase."
+    sample_rate: 16000
+    batch_size: 1
+    seed: 42
+    shard_seed: "randomized"
+
+exp_manager:
+   exp_dir: null
+   explicit_log_dir: salm_results/
+   name: salm
+   create_tensorboard_logger: false
+   create_checkpoint_callback: true
+   use_datetime_version: true
+   max_time_per_run: 00:03:50:00
+
+   resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc.
+   # you need to set these two to True to continue the training
+   resume_if_exists: true
+   resume_ignore_no_checkpoint: true
+
+   # You may use this section to create a W&B logger
+   create_wandb_logger: false
+   wandb_logger_kwargs:
+     name: development-run
+     project: salm
+     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

examples/speechlm2/conf/salm_asr_decoder_qformer.yaml

@@ -0,0 +1,164 @@
+model:
+  # Every name/path here starting with 'pretrained' is used to initialize the model weights.
+  pretrained_llm: nvidia/Llama-3.1-Nemotron-Nano-8B-v1
+  pretrained_asr: nvidia/parakeet-tdt-0.6b-v2
+
+  pretrained_weights: True  # When False, we use pretrained_name to load the architecture, but with random init
+
+  # Regexp (re.compile) patterns matching parameters to be frozen.
+  freeze_params:
+    # Frozen LLM
+    - "^llm\\..+$"  # LLM
+    - "^embed_tokens\\..+$"  # LLM embedding is moved
+    # Frozen pretrained ASR (only the modality adapter layers are trainable)
+    - "^perception\\.preprocessor\\..+$"
+    - "^perception\\.encoder\\..+$"
+  prevent_freeze_params: []  # Use to make specific submodules trainable; overrides freeze_params
+
+  prompt_format: llama3
+  audio_locator_tag: "<|audioplaceholder|>"  # placeholder token for audio turn is expected
+
+  # Note: Uncomment the block below to enable LoRA on LLM via HuggingFace PEFT library.
+  #   It will automatically freeze LLM parameters even if freeze_params was unused,
+  #   and prevent freezing any parameter that has the string '.lora_' in its name.
+  # lora:
+  #   task_type: CAUSAL_LM
+  #   r: 128
+  #   lora_alpha: 256
+  #   lora_dropout: 0.01
+
+  perception:
+     target: nemo.collections.speechlm2.modules.perception.AudioTranscriptionPerceptionModule
+     output_dim: 4096
+     modality_adapter:
+       _target_: nemo.collections.asr.modules.QformerConnector
+       prompt_size: 64
+       target_layer_ids: [0, 5, 11, 17, 23]
+       qformer_num_hidden_layers: 6
+       encoder_config:
+         d_model: 1024
+         encoder_attention_heads: 8
+       llm_config:
+         hidden_size: 4096
+
+  optimizer:
+    _target_: torch.optim.AdamW
+    lr: 5e-4
+    betas: [0.9, 0.98]
+    weight_decay: 1e-3
+    foreach: true # set to false if having issues with tensor-parallelism
+
+  lr_scheduler:
+    _target_: nemo.core.optim.lr_scheduler.CosineAnnealing
+    warmup_steps: 1000
+    min_lr: 1e-6
+    max_steps: ${trainer.max_steps}
+
+trainer:
+  devices: -1
+  accelerator: gpu
+  num_nodes: 1
+  precision: bf16-true
+  logger: False # logger provided by exp_manager
+  enable_checkpointing: False
+  use_distributed_sampler: False
+  max_steps: 100000
+  limit_train_batches: 5000  # "epoch" size
+  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:
+    # Replace DDPStrategy with ModelParallelStrategy to enable model parallelism
+    _target_: lightning.pytorch.strategies.DDPStrategy
+    gradient_as_bucket_view: true
+    find_unused_parameters: true
+    # _target_: lightning.pytorch.strategies.ModelParallelStrategy
+    # tensor_parallel_size: 1
+    # data_parallel_size: 8  # This is FSDP2
+
+data:
+  train_ds:
+    sample_rate: 16000
+    prompt_format: ${model.prompt_format}
+    token_equivalent_duration: 0.08
+    input_cfg:
+      - type: lhotse_as_conversation
+        cuts_path: ???  # needs to be set
+        audio_locator_tag: ${model.audio_locator_tag}
+        tags:
+          # Uncomment below line to include a system prompt for supported models (e.g. Llama; not Qwen).
+          # system_prompt: "some system prompt"
+          context: "Repeat after me, typing in lowercase."
+    seed: 42
+    shuffle: true
+    shard_seed: "randomized"
+    num_workers: 1
+    batch_size: 4
+    # Optional bucketing:
+    # batch_size: null
+    # use_bucketing: true
+    # use_multimodal_sampling: true
+    # measure_total_length: true
+    # Note: `batch_tokens`, `bucket_duration_bins`, and `max_tokens` all represent tokens as
+    #  the sum of input audio frames and output text tokens. Number of audio frames is
+    #  calculated using `token_equivalent_duration`.
+    # batch_tokens: 4000
+    # max_tokens: 2048
+    # bucket_duration_bins: [64, 128, 256, 384, 512, 768, 1024, 1280, 1536, 2048]
+    # num_buckets: 10
+    # bucket_buffer_size: 5000
+
+  validation_ds:
+    # The entries under 'datasets' are a list of separate dataloaders.
+    # The structure is <dataset-name>: {<dataloader-dict-config>}
+    # They inherit all settings from validation_ds, but can individually override them.
+    prompt_format: ${model.prompt_format}
+    token_equivalent_duration: 0.08
+    datasets:
+      val_set_0:  # rename to your dataset name, add more as needed
+        input_cfg:
+          - type: lhotse_as_conversation
+            cuts_path: ???  # needs to be set
+            audio_locator_tag: ${model.audio_locator_tag}
+            tags:
+              # Uncomment below line to include a system prompt for supported models (e.g. Llama; not Qwen).
+              # system_prompt: "some system prompt"
+              context: "Repeat after me, typing in lowercase."
+    sample_rate: 16000
+    batch_size: 1
+    seed: 42
+    shard_seed: "randomized"
+
+exp_manager:
+   exp_dir: null
+   explicit_log_dir: salm_results/
+   name: salm
+   create_tensorboard_logger: false
+   create_checkpoint_callback: true
+   use_datetime_version: true
+   max_time_per_run: 00:03:50:00
+
+   resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc.
+   # you need to set these two to True to continue the training
+   resume_if_exists: true
+   resume_ignore_no_checkpoint: true
+
+   # You may use this section to create a W&B logger
+   create_wandb_logger: false
+   wandb_logger_kwargs:
+     name: development-run
+     project: salm
+     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

examples/speechlm2/conf/salm_automodel.yaml

@@ -0,0 +1,272 @@
+model:
+  # Every name/path here starting with 'pretrained' is used to initialize the model weights.
+  pretrained_llm: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16
+  pretrained_asr: nvidia/canary-1b-v2
+
+  pretrained_weights: True  # When False, we use pretrained_name to load the architecture, but with random init
+
+  # 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. Set to null to train from pretrained_llm/asr.
+  init_from_checkpoint: null
+
+  # Set to true to use SALMAutomodel (NeMo Automodel backend) instead of SALM (HF Transformers backend).
+  use_nemo_automodel: true
+
+  # Regexp (re.compile) patterns matching parameters to be frozen.
+  freeze_params:
+    # Frozen LLM (embed_tokens stays inside llm, so this pattern covers it too)
+    - "^llm\\..+$"
+    # Frozen pretrained ASR (only the modality adapter layers are trainable).
+    # The second alternation matches the MultiLayerProjectionConnector path where the encoder
+    # is wrapped inside ``perception.encoder_multilayer.encoder`` instead of ``perception.encoder``.
+    - "^perception\\.preprocessor\\..+$"
+    - "^perception\\.encoder(_multilayer\\.encoder)?\\..+$"
+  prevent_freeze_params: []  # Use to make specific submodules trainable; overrides freeze_params
+
+  prompt_format: nemotron-nano-v3
+  audio_locator_tag: "<|audio|>"  # placeholder token for audio turn is expected
+  # Split audios longer than this before the speech encoder forward, then concatenate
+  # the encoded chunks back into one sequence before passing them to the LLM.
+  # Set to null to disable encoder chunking and encode each audio row directly.
+  encoder_chunk_size_seconds: 30.0
+
+  # Uncomment the block below to enable LoRA on the LLM via Automodel.
+  # LoRA parameters are kept trainable even when the LLM is frozen.
+  # lora:
+  #   dim: 128
+  #   alpha: 256
+  #   dropout: 0.01
+  #   target_modules: ["q_proj", "v_proj"]
+  #   # match_all_linear: false
+  #   # exclude_modules: []
+  #   # use_dora: false
+  #   # dropout_position: post
+  #   # lora_A_init: xavier
+  #   # use_triton: false
+
+  # MoE auxiliary load balancing loss coefficient. Set > 0 to enable.
+  # Gradients are injected automatically during backward — does not change reported CE loss.
+  # Typical range: 0.001 - 0.01.
+  aux_loss_coeff: 0.0
+
+  # MoE router/gate trainability. When true, unfreezes Gate.weight so the
+  # router can adapt to the new data distribution during fine-tuning.
+  # Default false — pretrained routing is kept frozen.
+  train_gate: false
+
+  # MoE expert balance/utilization metrics logging.
+  moe_metrics:
+    enabled: true
+    mode: brief            # "brief" or "detailed" (adds per-layer breakdowns)
+    detailed_every_steps: null  # null = every step when mode=detailed
+    top_k_experts: 5       # top/bottom utilization experts to report
+
+  # torch.compile configuration. When enabled, the LLM is compiled via
+  # Automodel's CompileConfig. Set dynamic=true for variable sequence lengths.
+  # compile:
+  #   enabled: false
+  #   mode: default          # "default", "reduce-overhead", or "max-autotune"
+  #   fullgraph: false       # Compile the full computation graph
+  #   dynamic: true          # Enable dynamic shapes (recommended for variable-length audio)
+  #   backend: null           # Compilation backend (null = inductor)
+  #   dynamo_cache_size_limit: 256  # Triton compilation cache limit
+
+  # Automodel backend dispatch. Selects the kernel/backend for each major module
+  # in the LLM (attention, linear, rms_norm, MoE experts/dispatcher). Defaults
+  # come from Automodel's BackendConfig and auto-select TE/DeepEP when available;
+  # override here to pin a specific backend (e.g. attn=sdpa to bypass TE).
+  # automodel_backend:
+  #   attn: te                  # "te" | "sdpa" | "flex"
+  #   linear: te                # "torch" | "te"
+  #   rms_norm: torch_fp32      # "torch" | "torch_fp32" | "te"
+  #   rope_fusion: true         # Fused RoPE (requires TE)
+  #   experts: torch_mm         # MoE expert GEMM: "torch" | "te" | "gmm" | "torch_mm"
+  #   dispatcher: deepep        # MoE token dispatcher: "torch" | "deepep" | "hybridep" | "uccl_ep"
+  #   dispatcher_num_sms: 20    # SM count for DeepEP/UCCL-EP kernels
+  #   fake_balanced_gate: false # Replace learned Gate with balanced fake gate (debug/bench)
+  #   fake_gate_noise: 0.0      # [0, 1] — noise for FakeBalancedGate routing
+  #   enable_hf_state_dict_adapter: true
+  #   enable_fsdp_optimizations: false
+  #   gate_precision: null      # e.g. "float32" to force fp32 gate compute
+  #   te_fp8: null              # {recipe: "current"} or {recipe: "block"} to enable TE FP8
+  #                             # (requires linear=te or experts=te)
+
+  # Pin the SDPA kernel list used when automodel_backend.attn=sdpa. Accepts
+  # strings from: "flash_attention", "efficient_attention", "math", "cudnn_attention".
+  # None = auto-select based on CP / activation checkpointing. Set to
+  # ["flash_attention"] to force FA2 and error out if unavailable.
+  # sdpa_method: null
+
+  perception:
+    target: nemo.collections.speechlm2.modules.perception.AudioPerceptionModule
+    output_dim: 4096
+    modality_adapter:
+      _target_: nemo.collections.speechlm2.modules.perception.IdentityConnector
+      d_model: 1024
+      
+    # Optional Rotary Time Embedding (RoTE, https://arxiv.org/abs/2410.12109): encodes absolute time into the encoder output
+    # features (at the modality-adapter entrance). Omit this block to disable (default).
+    # rote:
+    #   _target_: nemo.collections.speechlm2.modules.rote.RotaryTimeEmbedding
+    #   dim: 1024             # must match the encoder output feature dim (must be even)
+    #   theta: 1200.0         # base of the geometric frequency bank
+    #   rotary_fraction: 0.2  # fraction of dim to rotate
+    
+    # spec_augment:
+    #   _target_: nemo.collections.asr.modules.SpectrogramAugmentation
+    #   freq_masks: 2 # set to zero to disable it
+    #   time_masks: 10 # set to zero to disable it
+    #   freq_width: 27
+    #   time_width: 5  # 5 frames = 50ms
+
+  optimizer:
+    _target_: flashoptim.FlashAdamW
+    lr: 1e-4
+    betas: [0.9, 0.999]
+    weight_decay: 5e-2
+
+  lr_scheduler:
+    _target_: nemo.core.optim.lr_scheduler.CosineAnnealing
+    warmup_steps: 1000
+    min_lr: 1e-6
+    max_steps: ${trainer.max_steps}
+
+trainer:
+  devices: 8
+  accelerator: gpu
+  num_nodes: 1
+  precision: bf16-flash
+  logger: False # logger provided by exp_manager
+  enable_checkpointing: False
+  use_distributed_sampler: False
+  max_steps: 100000
+  limit_train_batches: 5000  # "epoch" size
+  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 and supports
+    # FSDP2, TP, PP, CP, EP (MoE), and HSDP. 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).
+    _target_: nemo.collections.speechlm2.parts.parallel.AutomodelParallelStrategy
+
+    # --- Parallelism dimensions ---
+    dp_size: null       # Data parallel size (null = inferred from world_size / other dims)
+    dp_replicate_size: 1  # HSDP replication group size (>1 enables hybrid sharding)
+    tp_size: 1          # Tensor parallel size
+    pp_size: 1          # Pipeline parallel size
+    cp_size: 1          # Context parallel size
+    ep_size: 8          # Expert parallel size (for MoE models)
+
+    # --- Activation checkpointing ---
+    # Two independent knobs: one for the LLM, one for the perception encoder.
+    # ``activation_checkpointing_llm`` is a single switch that covers both the
+    # non-EP FSDP2 path (via FSDP2Config.activation_checkpointing) and the
+    # EP/MoE parallelizer path.
+    # ``activation_checkpointing_perception`` wraps each layer in ``perception.encoder.layers`` 
+    # with ``checkpoint_wrapper`` before FSDP2 sharding.
+    activation_checkpointing_llm: false
+    activation_checkpointing_perception: false
+
+    # --- FSDP2 distributed config (plain dict, resolved to FSDP2Config automatically) ---
+    # distributed_config:
+    #   sequence_parallel: false    # Enable sequence parallelism (requires tp_size > 1)
+    #   # offload_policy:           # Uncomment to enable CPU offloading
+    #   #   _target_: torch.distributed.fsdp.CPUOffloadPolicy
+
+    # --- MoE config (plain dict, resolved to MoEParallelizerConfig automatically) ---
+    # moe_config:
+    #   ignore_router_for_ac: false       # Selective AC: save router outputs, recompute the rest
+    #   reshard_after_forward: false       # Reshard params after forward (saves memory, more comms)
+    #   lm_head_precision: null            # Override LM head precision (e.g., "float32" for stability)
+    #   wrap_outer_model: true             # Apply FSDP to the outer model wrapper
+
+data:
+  train_ds:
+    sample_rate: 16000
+    prompt_format: ${model.prompt_format}
+    token_equivalent_duration: 0.08
+    input_cfg:
+      - type: lhotse_as_conversation
+        cuts_path: ???  # needs to be set
+        audio_locator_tag: ${model.audio_locator_tag}
+        tags:
+          # Uncomment below line to include a system prompt for supported models (e.g. Llama; not Qwen).
+          # system_prompt: "some system prompt"
+          context: ""
+    seed: 42
+    shuffle: true
+    shard_seed: "randomized"
+    num_workers: 1
+    batch_size: 4
+    # Optional bucketing:
+    # batch_size: null
+    # use_bucketing: true
+    # use_multimodal_sampling: true
+    # measure_total_length: true
+    # Note: `batch_tokens`, `bucket_duration_bins`, and `max_tokens` all represent tokens as
+    #  the sum of input audio frames and output text tokens. Number of audio frames is
+    #  calculated using `token_equivalent_duration`.
+    # batch_tokens: 4000
+    # max_tokens: 2048
+    # bucket_duration_bins: [64, 128, 256, 384, 512, 768, 1024, 1280, 1536, 2048]
+    # num_buckets: 10
+    # bucket_buffer_size: 5000
+
+  validation_ds:
+    # The entries under 'datasets' are a list of separate dataloaders.
+    # The structure is <dataset-name>: {<dataloader-dict-config>}
+    # They inherit all settings from validation_ds, but can individually override them.
+    prompt_format: ${model.prompt_format}
+    token_equivalent_duration: 0.08
+    datasets:
+      val_set_0:  # rename to your dataset name, add more as needed
+        input_cfg:
+          - type: lhotse_as_conversation
+            cuts_path: ???  # needs to be set
+            audio_locator_tag: ${model.audio_locator_tag}
+            tags:
+              # Uncomment below line to include a system prompt for supported models (e.g. Llama; not Qwen).
+              # system_prompt: "some system prompt"
+              context: "Transcribe the following:"
+    sample_rate: 16000
+    batch_size: 1
+    seed: 42
+    shard_seed: "randomized"
+
+exp_manager:
+   exp_dir: null
+   explicit_log_dir: salm_results/
+   name: salm
+   create_tensorboard_logger: false
+   create_checkpoint_callback: true
+   use_datetime_version: true
+   max_time_per_run: 00:03:50:00
+
+   resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc.
+   # you need to set these two to True to continue the training
+   resume_if_exists: true
+   resume_ignore_no_checkpoint: true
+
+   # You may use this section to create a W&B logger
+   create_wandb_logger: false
+   wandb_logger_kwargs:
+     name: development-run
+     project: salm
+     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

examples/speechlm2/conf/salm_automodel_pee.yaml

@@ -0,0 +1,305 @@
+model:
+  # Every name/path here starting with 'pretrained' is used to initialize the model weights.
+  pretrained_llm: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16
+  pretrained_asr: nvidia/canary-1b-v2
+
+  pretrained_weights: True  # When False, we use pretrained_name to load the architecture, but with random init
+
+  # 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. Set to null to train from pretrained_llm/asr.
+  init_from_checkpoint: null
+
+  # Set to true to use SALMAutomodel (NeMo Automodel backend) instead of SALM (HF Transformers backend).
+  use_nemo_automodel: true
+
+  # Regexp (re.compile) patterns matching parameters to be frozen.
+  # PEE recipe: freeze the LLM and the Sortformer diarizer expert; keep the ASR
+  # Conformer encoder (perception.encoder.asr_encoder) and the fusion layers
+  # (perception.encoder.asr_norm / diar_norm) trainable.
+  freeze_params:
+    # Frozen LLM (embed_tokens stays inside llm, so this pattern covers it too)
+    - "^llm\\..+$"
+    # Frozen speech-feature preprocessor (mel front-end).
+    - "^perception\\.preprocessor\\..+$"
+    # Frozen Sortformer diarizer inside the Parallel Expert Encoder. The ASR encoder
+    # (perception.encoder.asr_encoder) and fusion norms deliberately stay trainable.
+    - "^perception\\.encoder\\.diarization_model\\..+$"
+  prevent_freeze_params: []  # Use to make specific submodules trainable; overrides freeze_params
+
+  prompt_format: nemotron-nano-v3
+  audio_locator_tag: "<|audio|>"  # placeholder token for audio turn is expected
+  # Split audios longer than this before the speech encoder forward, then concatenate
+  # the encoded chunks back into one sequence before passing them to the LLM.
+  # Set to null to disable encoder chunking and encode each audio row directly.
+  encoder_chunk_size_seconds: 60.0
+
+  # ─── Parallel Expert Encoder (PEE) options ──────────────────────────────────
+  # PEE swaps the perception encoder for a ParallelExpertEncoder bundle (streaming
+  # Sortformer diarizer + Canary ASR encoder), letting SALM emit <spk:N>-tagged
+  # multi-speaker (SOT) transcripts. The three blocks below define the PEE recipe;
+  # remove them (and data.multispeaker_cfg) to fall back to plain single-speaker SALM.
+
+  # Path to a ParallelExpertEncoderPT .nemo bundle, mounted onto perception.encoder
+  # at init. The bundle is not yet published to HuggingFace, so point this at a
+  # locally exported .nemo to enable PEE; null keeps the pretrained_asr encoder.
+  pe_encoder_path: null  # Plug-in PEE nemo checkpoint.
+
+  # Resolve the native <spk:N> speaker-token ids from the LLM tokenizer. Requires
+  # the patched "-spk" tokenizer (e.g. from patch_nano_v3_speaker_tokens.py) where
+  # <spk:0>..<spk:{max_speakers-1}> occupy contiguous ids starting at base_token_id.
+  speaker_tokens:
+    enable: true
+    template: "<spk:{i}>"   # speaker-token string format
+    max_speakers: 10        # number of contiguous <spk:N> entries to resolve
+    base_token_id: 100      # expected token id of <spk:0>
+
+  # Auxiliary Latent Speaker Supervision (LSS) loss. Its mere presence enables it
+  # (no enable flag); speaker_token_ids is auto-injected from speaker_tokens above.
+  # include_ce_loss MUST stay False — SALM already computes CE in loss_parallel().
+  lss_loss:
+    _target_: nemo.collections.common.losses.latent_speaker_supervision_loss.LatentSpeakerSupervisionLoss
+    speaker_loss_weight: 1.0
+    speaker_label_smoothing: 0.0
+    per_speaker_normalization: true
+
+  # Uncomment the block below to enable LoRA on the LLM via Automodel.
+  # LoRA parameters are kept trainable even when the LLM is frozen.
+  # lora:
+  #   dim: 128
+  #   alpha: 256
+  #   dropout: 0.01
+  #   target_modules: ["q_proj", "v_proj"]
+  #   # match_all_linear: false
+  #   # exclude_modules: []
+  #   # use_dora: false
+  #   # dropout_position: post
+  #   # lora_A_init: xavier
+  #   # use_triton: false
+
+  # MoE auxiliary load balancing loss coefficient. Set > 0 to enable.
+  # Gradients are injected automatically during backward — does not change reported CE loss.
+  # Typical range: 0.001 - 0.01.
+  aux_loss_coeff: 0.0
+
+  # MoE router/gate trainability. When true, unfreezes Gate.weight so the
+  # router can adapt to the new data distribution during fine-tuning.
+  # Default false — pretrained routing is kept frozen.
+  train_gate: false
+
+  # MoE expert balance/utilization metrics logging.
+  moe_metrics:
+    enabled: true
+    mode: brief            # "brief" or "detailed" (adds per-layer breakdowns)
+    detailed_every_steps: null  # null = every step when mode=detailed
+    top_k_experts: 5       # top/bottom utilization experts to report
+
+  # torch.compile configuration. When enabled, the LLM is compiled via
+  # Automodel's CompileConfig. Set dynamic=true for variable sequence lengths.
+  # compile:
+  #   enabled: false
+  #   mode: default          # "default", "reduce-overhead", or "max-autotune"
+  #   fullgraph: false       # Compile the full computation graph
+  #   dynamic: true          # Enable dynamic shapes (recommended for variable-length audio)
+  #   backend: null           # Compilation backend (null = inductor)
+  #   dynamo_cache_size_limit: 256  # Triton compilation cache limit
+
+  # Automodel backend dispatch. Selects the kernel/backend for each major module
+  # in the LLM (attention, linear, rms_norm, MoE experts/dispatcher). Defaults
+  # come from Automodel's BackendConfig and auto-select TE/DeepEP when available;
+  # override here to pin a specific backend (e.g. attn=sdpa to bypass TE).
+  # automodel_backend:
+  #   attn: te                  # "te" | "sdpa" | "flex"
+  #   linear: te                # "torch" | "te"
+  #   rms_norm: torch_fp32      # "torch" | "torch_fp32" | "te"
+  #   rope_fusion: true         # Fused RoPE (requires TE)
+  #   experts: torch_mm         # MoE expert GEMM: "torch" | "te" | "gmm" | "torch_mm"
+  #   dispatcher: deepep        # MoE token dispatcher: "torch" | "deepep" | "hybridep" | "uccl_ep"
+  #   dispatcher_num_sms: 20    # SM count for DeepEP/UCCL-EP kernels
+  #   fake_balanced_gate: false # Replace learned Gate with balanced fake gate (debug/bench)
+  #   fake_gate_noise: 0.0      # [0, 1] — noise for FakeBalancedGate routing
+  #   enable_hf_state_dict_adapter: true
+  #   enable_fsdp_optimizations: false
+  #   gate_precision: null      # e.g. "float32" to force fp32 gate compute
+  #   te_fp8: null              # {recipe: "current"} or {recipe: "block"} to enable TE FP8
+  #                             # (requires linear=te or experts=te)
+
+  # Pin the SDPA kernel list used when automodel_backend.attn=sdpa. Accepts
+  # strings from: "flash_attention", "efficient_attention", "math", "cudnn_attention".
+  # None = auto-select based on CP / activation checkpointing. Set to
+  # ["flash_attention"] to force FA2 and error out if unavailable.
+  # sdpa_method: null
+
+  perception:
+     target: nemo.collections.speechlm2.modules.perception.AudioPerceptionModule
+     output_dim: 4096
+     modality_adapter:
+       _target_: nemo.collections.speechlm2.modules.perception.IdentityConnector
+       d_model: 1024
+     # spec_augment:
+     #   _target_: nemo.collections.asr.modules.SpectrogramAugmentation
+     #   freq_masks: 2 # set to zero to disable it
+     #   time_masks: 10 # set to zero to disable it
+     #   freq_width: 27
+     #   time_width: 5  # 5 frames = 50ms
+
+  optimizer:
+    _target_: flashoptim.FlashAdamW
+    lr: 1e-4
+    betas: [0.9, 0.999]
+    weight_decay: 5e-2
+
+  lr_scheduler:
+    _target_: nemo.core.optim.lr_scheduler.CosineAnnealing
+    warmup_steps: 1000
+    min_lr: 1e-6
+    max_steps: ${trainer.max_steps}
+
+trainer:
+  devices: 8
+  accelerator: gpu
+  num_nodes: 1
+  precision: bf16-flash
+  logger: False # logger provided by exp_manager
+  enable_checkpointing: False
+  use_distributed_sampler: False
+  max_steps: 100000
+  limit_train_batches: 5000  # "epoch" size
+  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 and supports
+    # FSDP2, TP, PP, CP, EP (MoE), and HSDP. 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).
+    _target_: nemo.collections.speechlm2.parts.parallel.AutomodelParallelStrategy
+
+    # --- Parallelism dimensions ---
+    dp_size: null       # Data parallel size (null = inferred from world_size / other dims)
+    dp_replicate_size: 1  # HSDP replication group size (>1 enables hybrid sharding)
+    tp_size: 1          # Tensor parallel size
+    pp_size: 1          # Pipeline parallel size
+    cp_size: 1          # Context parallel size
+    ep_size: 8          # Expert parallel size (for MoE models)
+
+    # --- Activation checkpointing ---
+    # Two independent knobs: one for the LLM, one for the perception encoder.
+    # ``activation_checkpointing_llm`` is a single switch that covers both the
+    # non-EP FSDP2 path (via FSDP2Config.activation_checkpointing) and the
+    # EP/MoE parallelizer path.
+    # ``activation_checkpointing_perception`` wraps each layer in ``perception.encoder.layers`` 
+    # with ``checkpoint_wrapper`` before FSDP2 sharding.
+    activation_checkpointing_llm: false
+    activation_checkpointing_perception: false
+
+    # --- FSDP2 distributed config (plain dict, resolved to FSDP2Config automatically) ---
+    # distributed_config:
+    #   sequence_parallel: false    # Enable sequence parallelism (requires tp_size > 1)
+    #   # offload_policy:           # Uncomment to enable CPU offloading
+    #   #   _target_: torch.distributed.fsdp.CPUOffloadPolicy
+
+    # --- MoE config (plain dict, resolved to MoEParallelizerConfig automatically) ---
+    # moe_config:
+    #   ignore_router_for_ac: false       # Selective AC: save router outputs, recompute the rest
+    #   reshard_after_forward: false       # Reshard params after forward (saves memory, more comms)
+    #   lm_head_precision: null            # Override LM head precision (e.g., "float32" for stability)
+    #   wrap_outer_model: true             # Apply FSDP to the outer model wrapper
+
+data:
+  # RTTM/SOT speaker-activity targets for ParallelExpertEncoder training. Active for
+  # the PEE recipe; requires cuts/manifests that carry rttm_filepath entries and SOT
+  # (<spk:N>-tagged) text. Comment out to train SALM without speaker targets.
+  multispeaker_cfg:
+    num_speakers: 4                            # max speakers per sample (matches PEE n_spk)
+    sample_rate: ${data.train_ds.sample_rate}
+    window_stride: 0.01                        # preprocessor hop in seconds
+    subsampling_factor: 8                      # encoder output stride (mel -> target frames)
+    no_rttm_to_ones: true                      # cuts without RTTM -> single full-duration speaker
+
+  train_ds:
+    sample_rate: 16000
+    prompt_format: ${model.prompt_format}
+    token_equivalent_duration: 0.08
+    input_cfg:
+      - type: lhotse_as_conversation
+        cuts_path: ???  # needs to be set
+        audio_locator_tag: ${model.audio_locator_tag}
+        tags:
+          # Uncomment below line to include a system prompt for supported models (e.g. Llama; not Qwen).
+          # system_prompt: "some system prompt"
+          context: ""
+    seed: 42
+    shuffle: true
+    shard_seed: "randomized"
+    num_workers: 1
+    batch_size: 4
+    # Optional bucketing:
+    # batch_size: null
+    # use_bucketing: true
+    # use_multimodal_sampling: true
+    # measure_total_length: true
+    # Note: `batch_tokens`, `bucket_duration_bins`, and `max_tokens` all represent tokens as
+    #  the sum of input audio frames and output text tokens. Number of audio frames is
+    #  calculated using `token_equivalent_duration`.
+    # batch_tokens: 4000
+    # max_tokens: 2048
+    # bucket_duration_bins: [64, 128, 256, 384, 512, 768, 1024, 1280, 1536, 2048]
+    # num_buckets: 10
+    # bucket_buffer_size: 5000
+
+  validation_ds:
+    # The entries under 'datasets' are a list of separate dataloaders.
+    # The structure is <dataset-name>: {<dataloader-dict-config>}
+    # They inherit all settings from validation_ds, but can individually override them.
+    prompt_format: ${model.prompt_format}
+    token_equivalent_duration: 0.08
+    datasets:
+      val_set_0:  # rename to your dataset name, add more as needed
+        input_cfg:
+          - type: lhotse_as_conversation
+            cuts_path: ???  # needs to be set
+            audio_locator_tag: ${model.audio_locator_tag}
+            tags:
+              # Uncomment below line to include a system prompt for supported models (e.g. Llama; not Qwen).
+              # system_prompt: "some system prompt"
+              context: "Transcribe the following:"
+    sample_rate: 16000
+    batch_size: 1
+    seed: 42
+    shard_seed: "randomized"
+
+exp_manager:
+   exp_dir: null
+   explicit_log_dir: salm_results/
+   name: salm
+   create_tensorboard_logger: false
+   create_checkpoint_callback: true
+   use_datetime_version: true
+   max_time_per_run: 00:03:50:00
+
+   resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc.
+   # you need to set these two to True to continue the training
+   resume_if_exists: true
+   resume_ignore_no_checkpoint: true
+
+   # You may use this section to create a W&B logger
+   create_wandb_logger: false
+   wandb_logger_kwargs:
+     name: development-run
+     project: salm
+     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

examples/speechlm2/duplex_eartts_eval.py

@@ -0,0 +1,692 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Evaluation script for Duplex EARTTS models following MagpieTTS evaluation recipe.
+
+Args:
+    config-path (str):
+        Path to the directory containing the YAML configuration file.
+
+    config-name (str):
+        Name of the YAML configuration file.
+
+    checkpoint_path (str):
+        Path to the Duplex EARTTS checkpoint file.
+
+    datasets_json_path (str):
+        Path to a JSONL (JSON Lines) file describing the evaluation dataset.
+        Each line must be a valid JSON object representing one sample.
+
+        Supported formats:
+
+        ----------------------------------------------------------------------
+        1) SINGLE-TURN FORMAT
+        ----------------------------------------------------------------------
+        "text" is a string.
+
+        Example:
+        {"text": "Like really quickly and they go haha and then they run off.",
+         "context_audio_filepath": "speaker_1.wav",
+         "audio_filepath": "audio_1.wav"}
+
+        {"text": "Sure. Okay.",
+         "context_audio_filepath": "speaker_2.wav",
+         "audio_filepath": "audio_2.wav"}
+
+        ----------------------------------------------------------------------
+        2) MULTI-TURN FORMAT
+        ----------------------------------------------------------------------
+        "text" is a list of utterances (List[str]).
+        Each element represents one conversational turn. The model will
+        tokenize and pad each segment sequentially.
+
+        Example:
+        {"text": ["Okay yeah.", "Yeah.", "Right.", "I get what you’re saying.", "That makes sense."],
+         "context_audio_filepath": "speaker_1.wav",
+         "audio_filepath": "dummy_blank_audio_mt_0001.wav"}
+
+        {"text": ["Okay.", "Really?", "Yeah, okay.", "I didn’t know that.", "That’s interesting."],
+         "context_audio_filepath": "speaker_2.wav",
+         "audio_filepath": "audio_2.wav"}
+
+        ----------------------------------------------------------------------
+        FIELD DESCRIPTIONS
+        ----------------------------------------------------------------------
+
+        text:
+            Either:
+                - str (single-turn)
+                - List[str] (multi-turn)
+
+        context_audio_filepath:
+            Path to the reference speaker audio used for conditioning.
+            This can be overridden by setting:
+                ++user_custom_speaker_reference=<path>
+
+        audio_filepath:
+            Output audio file name.
+            This is used only as the base filename for saving generated audio
+            inside `out_dir`. The file does NOT need to exist beforehand.
+
+    out_dir (str):
+        Directory where generated audio samples will be saved.
+
+    inference_dtype (str, optional):
+        Target dtype used during inference. This controls the precision
+        of model weights and operations.
+
+        Supported values:
+            - "float32" (default)
+            - "float16"
+            - "bfloat16"
+
+        Notes:
+            - If set to a lower precision (e.g., float16), the model weights
+              and/or execution dtype will be adjusted accordingly.
+            - Internally mapped via `getattr(torch, inference_dtype)`.
+
+    keep_codec_original_dtype (bool, optional):
+        Controls whether the audio codec module keeps its original dtype
+        when `inference_dtype` is not float32.
+
+        If True (default):
+            - Only the TTS backbone (`model.tts_model`) is cast to the target dtype.
+            - The codec remains in its original precision (typically float32).
+            - Useful to isolate precision effects and avoid degradation from
+              codec quantization.
+
+        If False:
+            - The entire model (including codec) is cast to `inference_dtype`.
+            - `model.audio_codec_run_dtype` is also set accordingly.
+
+    debug_dtype (bool, optional):
+        Enables runtime inspection of tensor dtypes flowing through the model.
+
+        If True:
+            - Forward hooks are attached to all leaf modules.
+            - During the first batch, dtype usage statistics are collected
+              and logged.
+            - Outputs include:
+                - Per-module-group dtype distribution
+                - Example module names per dtype
+
+Usage:
+    # Example with fp32 inference
+        python duplex_eartts_eval.py \
+            --config-path=conf/ \
+            --config-name=duplex_eartts.yaml \
+            ++checkpoint_path=duplex_eartts_results/duplex_eartts/model.ckpt \
+            ++datasets_json_path=/path/to/evalset_config.jsonl \
+            ++out_dir=duplex_eartts_results/duplex_eartts/audio_samples/dummy_dataset
+
+    # Example with fp16 inference and dtype debugging
+        python duplex_eartts_eval.py \
+            --config-path=conf/ \
+            --config-name=duplex_eartts.yaml \
+            ++checkpoint_path=duplex_eartts_results/duplex_eartts/model.ckpt \
+            ++datasets_json_path=/path/to/evalset_config.jsonl \
+            ++out_dir=uplex_eartts_results/duplex_eartts/audio_samples/dummy_dataset \
+            ++inference_dtype=float16 \
+            ++keep_codec_original_dtype=True \
+            ++debug_dtype=True
+"""
+
+import json
+import os
+from functools import partial
+
+import librosa
+import soundfile as sf
+import torch
+from torch.nn.utils.rnn import pad_sequence
+from torch.utils.data import DataLoader, Dataset
+
+from nemo.collections.audio.parts.utils.transforms import resample
+
+torch.set_float32_matmul_precision("medium")
+torch.backends.cudnn.allow_tf32 = True
+torch.backends.cuda.matmul.allow_tf32 = True
+from contextlib import nullcontext
+
+from omegaconf import OmegaConf
+
+from nemo.collections.speechlm2.models.duplex_ear_tts import DuplexEARTTS
+from nemo.collections.speechlm2.parts.metrics.asr_cer_wer import Intelligibility
+from nemo.collections.speechlm2.parts.metrics.secs import SECS
+from nemo.collections.speechlm2.parts.precision import fp32_precision
+from nemo.core.config import hydra_runner
+from nemo.utils import logging
+
+# Use .get() to avoid crashing when running a single GPU without torchrun
+if torch.cuda.is_available():
+    torch.cuda.set_device(int(os.environ.get("LOCAL_RANK", 0)))
+
+
+def attach_dtype_counter(model):
+    """
+    Attaches forward hooks to all leaf modules of a model to track the dtype
+    of their outputs during inference.
+
+    This utility is designed for debugging precision behavior, especially when
+    using mixed precision or reduced precision (fp16 / bf16).
+
+    Behavior:
+        - Registers a forward hook on each leaf module (modules with no children).
+        - For each forward pass, records the dtype of the module output.
+        - Aggregates statistics grouped by top-level module name.
+        - Stores a few example module class names per dtype.
+
+    Returns:
+        handles (List[RemovableHandle]):
+            List of hook handles. These must be removed manually to avoid
+            memory leaks or performance degradation.
+
+        stats (Dict[str, Dict[str, int]]):
+            Nested dictionary containing dtype counts per module group.
+            Structure:
+                stats[module_group][dtype] = count
+
+            Example:
+                {
+                    "tts_model": {
+                        "torch.float16": 120,
+                        "torch.float32": 0,
+                        "torch.bfloat16": 0,
+                        "other": 2
+                    }
+                }
+
+        examples (Dict[str, Dict[str, List[str]]]):
+            Stores up to 3 example module class names per dtype per group.
+            Useful for quickly identifying which layers are running in
+            unexpected precision.
+
+    Notes:
+        - Only inspects outputs (not inputs or parameters).
+        - Dtype is inferred from the first tensor found in the output.
+        - Non-floating dtypes are categorized as "other".
+        - Grouping is based on the top-level module name (prefix before first dot).
+
+    Typical usage:
+        handles, stats, examples = attach_dtype_counter(model)
+
+        # Run inference ...
+
+        for h in handles:
+            h.remove()
+    """
+    handles = []
+
+    # structure: stats[module_group][dtype] = count
+    stats = {}
+    examples = {}
+
+    def is_leaf(module):
+        return len(list(module.children())) == 0
+
+    def get_dtype(x):
+        if torch.is_tensor(x):
+            return str(x.dtype)
+        elif isinstance(x, (list, tuple)):
+            for t in x:
+                if torch.is_tensor(t):
+                    return str(t.dtype)
+        return "other"
+
+    def get_module_group(name):
+        # top-level module (before first dot)
+        return name.split(".")[0] if "." in name else name
+
+    def hook_fn(name):
+        def fn(module, inputs, outputs):
+            dtype = get_dtype(outputs)
+            if dtype not in ["torch.float16", "torch.bfloat16", "torch.float32"]:
+                dtype = "other"
+
+            group = get_module_group(name)
+
+            if group not in stats:
+                stats[group] = {
+                    "torch.float16": 0,
+                    "torch.bfloat16": 0,
+                    "torch.float32": 0,
+                    "other": 0,
+                }
+                examples[group] = {
+                    "torch.float16": [],
+                    "torch.bfloat16": [],
+                    "torch.float32": [],
+                    "other": [],
+                }
+
+            stats[group][dtype] += 1
+
+            # store a few examples per dtype per group
+            if len(examples[group][dtype]) < 3:
+                examples[group][dtype].append(module.__class__.__name__)
+
+        return fn
+
+    for name, module in model.named_modules():
+        if is_leaf(module):
+            handles.append(module.register_forward_hook(hook_fn(name)))
+
+    return handles, stats, examples
+
+
+def report_dtype_stats(handles, stats, examples):
+    """
+    Cleans up monitoring hooks and logs a detailed report of the tensor precisions
+    (dtypes) observed during the model forward pass.
+
+    This function should be called after at least one inference iteration has
+    completed while hooks are attached. It removes the hooks to prevent
+    performance overhead and prints a structured summary of which module groups
+    executed in which dtypes.
+
+    Args:
+        handles (List[torch.utils.hooks.RemovableHandle]): The list of hooks
+            returned by `attach_dtype_counter`.
+        stats (Dict): Nested dictionary containing dtype counts per module group.
+        examples (Dict): Dictionary containing example module names for each
+            observed dtype.
+    """
+    for h in handles:
+        h.remove()
+
+    logging.info("\n=== DTYPE USAGE PER MODULE ===")
+
+    for group, group_stats in stats.items():
+        total = sum(group_stats.values())
+        if total == 0:
+            continue
+
+        logging.info(f"\n--- {group} ---")
+        for dtype, count in group_stats.items():
+            if count > 0:
+                logging.info(f"{dtype}: {count} ({100*count/total:.2f}%)")
+
+    logging.info("\n=== EXAMPLES ===")
+    for group, group_examples in examples.items():
+        logging.info(f"\n--- {group} ---")
+        for dtype, mods in group_examples.items():
+            if mods:
+                logging.info(f"{dtype}: {mods}")
+
+
+class EvalJSONLDataset(Dataset):
+    """
+    Standard PyTorch Dataset for reading JSONL evaluation files.
+    """
+
+    def __init__(self, file_path):
+        self.samples = []
+        with open(file_path, "r", encoding="utf-8") as f:
+            for line_idx, line in enumerate(f, 1):
+                line = line.strip()
+                if not line:
+                    continue
+                try:
+                    self.samples.append(json.loads(line))
+                except json.JSONDecodeError as e:
+                    raise ValueError(f"Invalid JSON on line {line_idx}: {e}")
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, idx):
+        return self.samples[idx]
+
+
+def collate_and_tokenize_custom(
+    batch,
+    model,
+    extra_duration_thrshould=1.3,
+    sample_rate=22050,
+    root_path=None,
+    add_beginning_pad_tokens=False,
+    add_eos=False,
+    pad_factor_text_speech=10,
+    force_interruption=False,
+):
+    tokenized_list = []
+
+    # --- TEXT TOKENIZATION ---
+    for s in batch:
+        text_data = s["text"]
+
+        # Check if text is a list (New Logic)
+        if isinstance(text_data, list):
+            # Start with BOS
+            full_ids = []
+
+            for segment in text_data:
+                # Tokenize segment
+                seg_ids = [model.tokenizer.bos]
+                seg_ids = seg_ids + model.tokenizer.text_to_ids(segment)
+                seg_len = len(seg_ids)
+
+                # Calculate pad length (pad_factor_text_speechx the size of the text)
+                pad_len = seg_len * pad_factor_text_speech
+
+                # Construct: text + 4x pads
+                # We extend the list with the tokens and then the pad tokens
+                pad_ids = [model.text_pad_id] * pad_len
+
+                if force_interruption:
+                    fname = s["audio_filepath"]
+                    no_ext = fname.split(".")[0]
+                    sample_id = int(no_ext.split("_")[-1])
+
+                    case = sample_id % 3  # 0,1,2 -> ~33% each
+
+                    if case == 0:
+                        # 33%: emulate interruption where text was not fully processed
+                        # (no pad eos placement at all)
+                        if len(seg_ids) >= 2:
+                            seg_ids[-2] = model.text_eos_id
+                            seg_ids[-1] = model.text_pad_id
+                        else:
+                            # fallback: if seg_ids is too short, emulate with pad EOS at 0
+                            pad_ids[0] = model.text_eos_id
+                    elif case == 1:
+                        # 33%: put EOS at pad index 6 - so 0.5 seconds after the whole text was processed
+                        eos_idx = min(6, len(pad_ids) - 1)
+                        pad_ids[eos_idx] = model.text_eos_id
+                    else:
+                        # 33%: put EOS at pad index 0
+                        eos_idx = 0
+                        pad_ids[eos_idx] = model.text_eos_id
+                else:
+                    if (
+                        add_eos
+                    ):  # add eos in the end of the paddding sequence keep 70% for the speech and the rest for after EOS
+                        eos_idx = int(len(pad_ids) * 0.7)
+                        pad_ids[eos_idx] = model.text_eos_id
+
+                full_ids.extend(seg_ids)
+                full_ids.extend(pad_ids)
+
+            tokenized_list.append(torch.as_tensor(full_ids, dtype=torch.long))
+
+        else:
+            # Standard String Handling
+            tokenized_list.append(
+                torch.as_tensor([model.tokenizer.bos] + model.tokenizer.text_to_ids(text_data), dtype=torch.long)
+            )
+
+    if add_beginning_pad_tokens:
+        pad_len = 25
+        prefix = torch.full((pad_len,), model.text_pad_id, dtype=torch.long)
+        for i in range(len(tokenized_list)):
+            tokenized_list[i] = torch.cat([prefix, tokenized_list[i]])
+
+    # Pad the text sequences (batch-wise)
+    input_ids = pad_sequence(tokenized_list, batch_first=True, padding_value=model.text_pad_id)
+
+    # load the target audio if available
+    audio_list = []
+    audio_lengths = []
+    target_num_frames = []
+
+    for i, s in enumerate(batch):
+        # Load Context Audio
+        audio_path = s["context_audio_filepath"]
+        if root_path is not None:
+            audio_path = os.path.join(root_path, audio_path)
+
+        # Safety check for context audio presence, though usually required
+        if os.path.exists(audio_path):
+            wav, sr = librosa.load(audio_path, sr=sample_rate, mono=True)
+            wav = torch.as_tensor(wav, dtype=torch.float32)
+        else:
+            # Fallback if context missing (optional safety)
+            wav = torch.zeros(1, dtype=torch.float32)
+
+        audio_list.append(wav)
+        audio_lengths.append(len(wav))
+
+        # Handle Target Audio / Duration
+        tdur_audio_path = s["audio_filepath"]
+        if root_path is not None:
+            tdur_audio_path = os.path.join(root_path, tdur_audio_path)
+
+        # Check availability
+        if tdur_audio_path and os.path.exists(tdur_audio_path):
+            wav_dur, sr_ = librosa.load(tdur_audio_path, sr=sample_rate, mono=True)
+            tdur = wav_dur.shape[0] // model.target_samples_per_frame
+            target_num_frames.append(tdur * extra_duration_thrshould)
+        else:
+            # Audio not available: Derive size from text channel
+            # We follow the 4x approach logic here to determine frames.
+            # If text was a list, it already has physical pads (1 + 4 ratio).
+            # We map 1 token roughly to 1 frame (or whatever the model scale is).
+            # Assuming 1 token ~ 1 frame in the model's alignment, we just take the input length.
+
+            current_text_len = len(tokenized_list[i])
+
+            if isinstance(s["text"], list):
+                # The text tokens are already physically padded 10x.
+                # Target frames should match this structure exactly.
+                target_num_frames.append(current_text_len)
+            else:
+                # If text was a string (no physical pads added), but audio is missing,
+                # we simulate the 4x duration expansion (1 part text, 4 parts silence = 5x total).
+                target_num_frames.append(current_text_len * 5)
+
+    # audio padding
+    max_audio_len = max(audio_lengths)
+    B = len(audio_lengths)
+
+    padded_audio = torch.zeros((B, max_audio_len), dtype=torch.float32)
+
+    for i, wav in enumerate(audio_list):
+        padded_audio[i, : len(wav)] = wav
+
+    # Keep on CPU
+    audio_lengths = torch.tensor(audio_lengths, dtype=torch.long)
+
+    # Expand text length to match expected output speech duration
+    B, L = input_ids.shape
+    target_len = int(max(target_num_frames))
+
+    # Ensure target_len is at least as long as the input text
+    # (prevents truncation if calc was slightly off)
+    target_len = max(target_len, L)
+
+    padded_input_ids = torch.full((B, target_len), fill_value=model.text_pad_id, dtype=input_ids.dtype)
+
+    # Copy the actual tokens (which might already contain list-based padding)
+    padded_input_ids[:, :L] = input_ids
+
+    # If text is a list ["Hi", "There"], join it into "Hi There"
+    collapsed_raw_text = [" ".join(s["text"]) if isinstance(s["text"], list) else s["text"] for s in batch]
+
+    return {
+        "input_ids": padded_input_ids,
+        "raw_text": collapsed_raw_text,
+        "context_audio": padded_audio,
+        "context_audio_lengths": audio_lengths,
+        "target_audio_paths": [s["audio_filepath"] for s in batch],
+        "target_num_frames": target_num_frames,
+    }
+
+
+@hydra_runner(config_path="conf", config_name="duplex_eartts")
+def inference(cfg):
+    OmegaConf.resolve(cfg)
+
+    distributed = int(os.environ.get("WORLD_SIZE", "1")) > 1
+    if distributed and not torch.distributed.is_initialized():
+        torch.distributed.init_process_group(backend="nccl")
+
+    # Dynamically determine the correct GPU for this process
+    if torch.cuda.is_available():
+        local_rank = int(os.environ.get("LOCAL_RANK", 0))
+        target_device = torch.device(f"cuda:{local_rank}")
+    else:
+        target_device = torch.device("cpu")
+
+    torch.set_float32_matmul_precision("medium")
+    torch.backends.cudnn.allow_tf32 = True
+    torch.backends.cuda.matmul.allow_tf32 = True
+
+    target_dtype = getattr(torch, cfg.get("inference_dtype", "float32"))
+    if target_dtype != torch.float32:
+        torch.set_default_dtype(target_dtype)
+
+    if cfg.get("checkpoint_path", None):
+        model = DuplexEARTTS.load_from_checkpoint(
+            cfg.checkpoint_path, cfg=OmegaConf.to_container(cfg, resolve=True), map_location=target_device
+        ).eval()
+    else:
+        raise ValueError("For evaluation, you must provide `cfg.checkpoint_path`.")
+
+    if target_dtype != torch.float32:
+        if cfg.get("keep_codec_original_dtype", True):
+            model.tts_model.to(dtype=target_dtype)
+            model.ensures_codec_target_dtype()  # ensures that codec is in the right precision
+        else:
+            model.audio_codec_run_dtype = target_dtype
+            model.to(dtype=target_dtype)
+
+    if cfg.get("debug_dtype", False):
+        handles, stats, examples = attach_dtype_counter(model)
+
+    with fp32_precision():
+        intelligibility = Intelligibility("stt_en_fastconformer_transducer_large", reuse_asr_hyps=False).reset()
+        secs_metric = SECS("titanet_large").reset()
+
+    # Initialize the Dataset
+    eval_dataset = EvalJSONLDataset(cfg.datasets_json_path)
+
+    # Use partial to bind the model and config parameters to the collate function
+    collate_fn = partial(
+        collate_and_tokenize_custom,
+        model=model,
+        extra_duration_thrshould=1.5,
+        sample_rate=model.target_sample_rate,
+        root_path=cfg.audio_dir,
+        add_beginning_pad_tokens=cfg.get("add_beginning_pad_tokens", True),
+        add_eos=cfg.get("add_eos", True),
+        pad_factor_text_speech=cfg.get("pad_factor_text_speech", 10),
+        force_interruption=cfg.get("force_interruption", False),
+    )
+
+    # Initialize the DataLoader
+    dataloader = DataLoader(
+        dataset=eval_dataset,
+        batch_size=cfg.batch_size,
+        collate_fn=collate_fn,
+        num_workers=cfg.get("num_workers", 4),
+        pin_memory=True,
+        shuffle=False,
+        drop_last=False,
+    )
+
+    if cfg.get("user_custom_speaker_reference", None):
+        wav, sr = librosa.load(cfg.model.inference_speaker_reference, sr=model.target_sample_rate, mono=True)
+        speaker_wav = torch.as_tensor(wav, dtype=target_dtype).unsqueeze(0).to(model.device)
+
+    # Iterate over the DataLoader
+    for batch_id, inputs in enumerate(dataloader):
+
+        # Move required tensors to the GPU immediately
+        inputs["input_ids"] = inputs["input_ids"].to(model.device)
+        inputs["context_audio"] = inputs["context_audio"].to(model.device)
+        inputs["context_audio_lengths"] = inputs["context_audio_lengths"].to(model.device)
+
+        if cfg.get("user_custom_speaker_reference", None):
+            inputs["context_audio"] = speaker_wav.expand(inputs["input_ids"].size(0), *speaker_wav.shape[1:])
+            inputs["context_audio_lengths"][:] = speaker_wav.size(-1)
+
+        with torch.no_grad():
+            model.set_init_inputs(
+                speaker_audio=inputs["context_audio"],
+                speaker_audio_lens=inputs["context_audio_lengths"],
+                system_prompt=cfg.get("inference_system_prompt", None),
+            )
+            init_inputs = model.get_init_inputs(B=inputs["input_ids"].size(0))
+
+            audio, audio_len = model.offline_inference(
+                next_subword_ids=inputs["input_ids"],
+                task="custom",
+                init_inputs=init_inputs,
+            )
+
+        if cfg.get("debug_dtype", False) and batch_id == 0:
+            report_dtype_stats(handles, stats, examples)
+
+        with fp32_precision():
+            audio = audio.float()
+
+            # reset audio len to the actual size removing extra long audio padding
+            audio_len = (
+                torch.tensor(inputs["target_num_frames"], device=audio.device) * model.target_samples_per_frame
+            ).int()
+
+            # resample audio to the asr sampling rate
+            metric_audio_pred = resample(audio, model.target_sample_rate, 16000)
+            metric_audio_pred_lens = (audio_len / model.target_sample_rate * 16000).to(torch.long)
+
+            intelligibility.update(
+                name="dataset",
+                refs=inputs["raw_text"],
+                pred_audio=metric_audio_pred,
+                pred_audio_lens=metric_audio_pred_lens,
+                asr_hyps=None,
+            )
+
+            secs_metric.update(
+                name="dataset",
+                target_audio=resample(inputs["context_audio"], model.target_sample_rate, 16000),
+                target_audio_lens=(inputs["context_audio_lengths"] / model.target_sample_rate * 16000).to(torch.long),
+                pred_audio=metric_audio_pred,
+                pred_audio_lens=metric_audio_pred_lens,
+            )
+
+            # save audio to cfg.out_dir
+            os.makedirs(cfg.out_dir, exist_ok=True)
+            audio = audio.detach().cpu().float()
+            audio_len = audio_len.cpu()
+
+            for i in range(audio.size(0)):
+                wav = audio[i, : audio_len[i]].numpy()
+                # Use original target audio filename
+                target_path = inputs["target_audio_paths"][i]
+                base_name = os.path.basename(target_path)
+                out_path = os.path.join(cfg.out_dir, base_name)
+
+                sf.write(
+                    out_path,
+                    wav,
+                    samplerate=model.target_sample_rate,
+                )
+
+                logging.info(f"Saved: {out_path}")
+
+    with fp32_precision():
+        logging.info("\n--- Evaluation Metrics ---")
+        cer_wer = intelligibility.compute()
+        for k, m in cer_wer.items():
+            logging.info(f"Intelligibility - {k}: {m}")
+
+        secs_scores = secs_metric.compute()
+        for k, m in secs_scores.items():
+            logging.info(f"SECS - {k}: {m}")
+
+
+if __name__ == "__main__":
+    inference()

examples/speechlm2/duplex_eartts_train.py

@@ -0,0 +1,75 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import datetime
+import os
+
+import torch
+from lightning.pytorch import Trainer
+from omegaconf import OmegaConf
+
+from nemo.collections.speechlm2 import DataModule, DuplexEARTTSDataset
+from nemo.collections.speechlm2.models.duplex_ear_tts import DuplexEARTTS
+from nemo.collections.speechlm2.parts.pretrained import load_checkpoint, set_model_dict_for_partial_init
+from nemo.core.config import hydra_runner
+from nemo.utils.exp_manager import exp_manager
+from nemo.utils.trainer_utils import resolve_trainer_cfg
+
+torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
+
+
+@hydra_runner(config_path="conf", config_name="duplex_eartts")
+def train(cfg):
+    OmegaConf.resolve(cfg)
+    torch.distributed.init_process_group(
+        backend="nccl", timeout=datetime.timedelta(seconds=int(cfg.trainer.strategy.get("timeout", 3600)))
+    )
+    torch.set_float32_matmul_precision("medium")
+    torch.backends.cudnn.allow_tf32 = True
+    trainer = Trainer(**resolve_trainer_cfg(cfg.trainer))
+    log_dir = exp_manager(trainer, cfg.get("exp_manager", None))
+    OmegaConf.save(cfg, log_dir / "exp_config.yaml")
+
+    with trainer.init_module():
+        model = DuplexEARTTS(OmegaConf.to_container(cfg, resolve=True))
+
+        # load pretrained tts checkpoint if available
+        if model.cfg.get("pretrained_tts_model", None):
+            checkpoint_state = load_checkpoint(model.cfg.pretrained_tts_model)
+            checkpoint_state = set_model_dict_for_partial_init(checkpoint_state, model.tts_model.state_dict())
+            model.tts_model.load_state_dict(checkpoint_state, strict=True)
+
+        # load pretrained checkpoint and rescale the weights if needed
+        if model.cfg.get("pretrained_model", None):
+            model.restore_from_pretrained_checkpoint(model.cfg.pretrained_model)
+
+    dataset = DuplexEARTTSDataset(
+        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,
+        add_text_bos_and_eos_in_each_turn=cfg.data.get("add_text_bos_and_eos_in_each_turn", True),
+        add_audio_prompt=cfg.data.get("add_audio_prompt", True),
+        audio_prompt_duration=cfg.data.get("audio_prompt_duration", 3),
+        num_delay_speech_tokens=cfg.model.get("num_delay_speech_tokens", 2),
+        add_system_prompt=cfg.model.get("use_system_prompt", False),
+        ignore_data_system_prompt=cfg.model.get("ignore_data_system_prompt", False),
+    )
+    datamodule = DataModule(cfg.data, tokenizer=model.tokenizer, dataset=dataset)
+    trainer.fit(model, datamodule)
+
+
+if __name__ == "__main__":
+    train()

examples/speechlm2/nemotron_voicechat_eval.py

@@ -0,0 +1,171 @@
+# 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.
+
+"""
+Evaluation script for NemotronVoiceChat models.
+
+This script runs validation for a NemotronVoiceChat checkpoint using a 
+Duplex S2S/STT-style Lhotse dataset. It evaluates the full speech-to-speech 
+pipeline, including both the Duplex STT and Duplex TTS models.
+
+Metrics
+-------
+During validation, the script computes:
+- Text BLEU score (reference text vs predicted text)
+- ASR BLEU score (reference text vs ASR-transcribed generated speech)
+
+The ASR model used for scoring is defined by the configuration parameter:
+    model.scoring_asr
+This model is used to transcribe generated speech and compute BLEU-based 
+speech consistency metrics. The specific ASR checkpoint is fully controlled 
+via config, in the same way as other parameters such as:
+    exp_manager.explicit_log_dir
+
+Arguments
+---------
+For a complete configuration reference, please look at the example config located at:
+    examples/speechlm2/conf/nemotron_voicechat.yaml
+
+cfg : omegaconf.DictConfig
+    The main Hydra configuration object defining the evaluation parameters. 
+    It is expected to contain the following top-level configurations:
+    
+    checkpoint_path (str | null)
+        Path to the pre-trained NemotronVoiceChat checkpoint for evaluation.
+
+    model (DictConfig)
+        Model-specific settings encompassing both STT and TTS subsystems. Key parameters include:
+        * scoring_asr (str): The ASR model name/path used to evaluate generated speech (e.g., 'stt_en_fastconformer_transducer_large').
+        * inference_speaker_reference (str | null): Path to the reference audio used to condition the speaker's voice. Set to `null` if using `inference_speaker_name`.
+        * inference_speaker_name (str): Named speaker identifier (e.g., 'Megan'); overrides `inference_speaker_reference`.
+        * stt (DictConfig): Sub-config for the `DuplexSTTModel` (e.g., `eval_text_turn_taking`).
+        * speech_generation (DictConfig): Sub-config for the `DuplexEARTTS` model. Includes codec configs, EAR-TTS backbone, and inference behavior like `inference_guidance_scale` (CFG) and `inference_noise_scale` (sampling temperature for the MoG head).
+
+    data (DictConfig)
+        Configuration for the data pipelines, datasets, and DataModule. Key parameters include:
+        * source_sample_rate (int): Sample rate of the input/user audio (e.g., 16000).
+        * target_sample_rate (int): Sample rate of the generated output audio (e.g., 22050).
+        * frame_length (float): Duration of audio frames in seconds (e.g., 0.08).
+        * input_roles (list): Conversation roles mapped to the input prompt (e.g., ["user", "User"]).
+        * output_roles (list): Conversation roles targeted for model generation (e.g., ["agent", "Assistant"]).
+        * validation_ds (DictConfig): Paths and settings for the Lhotse validation shards (e.g., `shar_path`, `batch_size`). Note that the data format for `data.validation_ds.evaluation_set` must follow the `duplexs2s-dataset-structure`. For detailed specifications, see: https://docs.nvidia.com/nemo/speech/nightly/speechlm2/datasets.html#duplexs2s-dataset-structure
+
+    exp_manager (DictConfig)
+        Experiment manager configurations for logging. Must include:
+        * name (str): Experiment name (e.g., 'nemotron-voicechat-eval').
+        * explicit_log_dir (str): The root directory where output artifacts and metric logs are saved.
+
+    trainer (DictConfig)
+        PyTorch Lightning Trainer parameters dictating hardware usage. Key settings include `devices`, `num_nodes`, `limit_val_batches` (fraction of dataset to evaluate), and `precision` (e.g., 32).
+
+Example Run
+-----------
+You can run the evaluation script and override parameters dynamically using Hydra command-line flags. 
+Here is an example execution using dummy paths:
+
+    python /path/to/nemo/examples/speechlm2/nemotron_voicechat_eval.py \
+        --config-path=examples/speechlm2/conf/ \
+        --config-name=nemotron_voicechat.yaml \
+        exp_manager.name="Nemotron_VoiceChat_Eval" \
+        ++model.stt.model.eval_text_turn_taking=True \
+        ++checkpoint_path="/path/to/nemotron_voicechat_ckpt/" \
+        ++model.inference_speaker_reference=null \
+        ++model.inference_speaker_name="Megan" \
+        ++model.speech_generation.model.inference_guidance_scale=0.2 \
+        ++model.speech_generation.model.inference_guidance_enabled=True \
+        ++model.speech_generation.model.inference_top_p_or_k=0.95 \
+        ++model.speech_generation.model.inference_noise_scale=0.001 \
+        trainer.num_nodes=1 \
+        exp_manager.explicit_log_dir="/path/to/results_dir/Nemotron_VoiceChat_Eval/" \
+        data.validation_ds.batch_size=2 \
+        data.validation_ds.datasets.evaluation_set.shar_path="/path/to/validation_dataset/" \
+        ++trainer.limit_val_batches=1.0 \
+        ++trainer.precision=32 \
+        data.validation_ds.seed=42
+
+Outputs
+-------
+All generated artifacts are saved under:
+    exp_manager.explicit_log_dir + "/validation_logs"
+
+The script:
+- Saves generated audio files
+- Saves per-utterance logs in JSON format via `ResultsLogger`
+- Saves predicted text, target text, and ASR-transcribed speech
+
+Each validation example is exported as a JSON entry with the following format:
+{
+    "target_text": "...",
+    "pred_text": "...",
+    "speech_pred_transcribed": "...",
+    "audio_path": "pred_wavs/example.wav"
+}
+
+Where:
+    target_text: Ground-truth target text.
+    pred_text: Text predicted by the STT/S2S model.
+    speech_pred_transcribed: Transcription of the generated speech using the ASR model defined by `model.scoring_asr`.
+    audio_path: Relative path to the generated waveform inside exp_manager.explicit_log_dir.
+"""
+
+import os
+
+import torch
+from lightning.pytorch import Trainer
+from omegaconf import OmegaConf
+
+from nemo.collections.speechlm2 import DataModule, DuplexSTTDataset
+from nemo.collections.speechlm2.models.nemotron_voicechat import NemotronVoiceChat
+from nemo.core.config import hydra_runner
+from nemo.utils.exp_manager import exp_manager
+from nemo.utils.trainer_utils import resolve_trainer_cfg
+
+torch.set_float32_matmul_precision("medium")
+torch.backends.cudnn.allow_tf32 = True
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
+
+
+@hydra_runner(config_path="conf", config_name="nemotron_voicechat")
+def inference(cfg):
+    OmegaConf.resolve(cfg)
+    torch.distributed.init_process_group(backend="nccl")
+    torch.set_float32_matmul_precision("medium")
+    torch.backends.cudnn.allow_tf32 = True
+    trainer = Trainer(**resolve_trainer_cfg(cfg.trainer))
+    log_dir = exp_manager(trainer, cfg.get("exp_manager", None))
+
+    with trainer.init_module():
+        # instanciate and load the model using from_pretrained
+        model = NemotronVoiceChat.from_pretrained(cfg.checkpoint_path).eval()
+
+    # update model internal configs using the new configs
+    model.full_cfg.merge_with(cfg)
+    model.cfg.merge_with(cfg.model)
+    OmegaConf.save(model.full_cfg, log_dir / "exp_config.yaml")
+    model.validation_save_path = os.path.join(log_dir, "validation_logs")
+
+    dataset = DuplexSTTDataset(
+        tokenizer=model.stt_model.tokenizer,
+        frame_length=cfg.data.frame_length,
+        source_sample_rate=cfg.data.source_sample_rate,
+        input_roles=cfg.data.input_roles,
+        output_roles=cfg.data.output_roles,
+    )
+    datamodule = DataModule(cfg.data, tokenizer=model.stt_model.tokenizer, dataset=dataset)
+    trainer.validate(model, datamodule)
+
+
+if __name__ == "__main__":
+    inference()

examples/speechlm2/nemotron_voicechat_to_hf.py

@@ -0,0 +1,300 @@
+# 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.
+"""
+Utility script for exporting Duplex Speech-to-Chat models to HuggingFace format.
+
+This exporter supports:
+- Standard PyTorch checkpoints (non-distributed)
+- Distributed checkpoints (FSDP / TP / torch.distributed checkpoint directories)
+- safetensors checkpoints
+
+The script loads STT and TTS checkpoints, constructs a joint NemotronVoiceChat
+model instance, and saves a HuggingFace-compatible checkpoint that can be
+loaded via:
+
+    NemotronVoiceChat.from_pretrained(path)
+
+Arguments
+---------
+tts_ckpt_path : str
+    Path to the TTS checkpoint.
+    - Can be a standard PyTorch checkpoint (.ckpt or .pt)
+    - Can be a directory containing distributed checkpoints (FSDP/TP)
+    - Can be a safetensors file
+
+tts_ckpt_config : str
+    Path to the experiment configuration used to instantiate the TTS model.
+    This configuration defines architecture, hyperparameters, and data settings
+    required to reconstruct the model before loading weights.
+
+stt_ckpt_path : str
+    Path to the STT (speech-to-text) checkpoint.
+    - Supports PyTorch checkpoints
+    - Supports distributed checkpoints
+    - Supports safetensors
+
+stt_ckpt_config : str
+    Path to the experiment configuration for the STT model.
+    Used to reconstruct the STT module before loading weights.
+
+output_dir : str
+    Directory where the HuggingFace-compatible checkpoint will be stored.
+    After export, the model can be loaded via:
+
+        NemotronVoiceChat.from_pretrained(output_dir)
+
+dtype : str (optional, default="float32")
+    Target dtype for storing parameters.
+    Typical values: "float32", "float16", "bfloat16".
+    Controls the precision of saved weights.
+
+register_speaker_dict : Dict[str, str] (optional)
+    Dictionary mapping speaker names to reference audio paths.
+    Used for speaker registration and voice cloning at inference time.
+
+    Example:
+        {
+            "Megan": "/path/to/megan_reference.wav",
+            "Emma": "/path/to/emma_reference.wav"
+        }
+
+    Each entry:
+    - Key: Speaker identifier (string)
+    - Value: Path to an audio file containing the speaker’s voice
+
+    The exporter loads the reference audio, resamples it to the model’s target
+    sample rate, and registers it as an audio prompt latent so the model can
+    generate speech in that speaker’s voice.
+
+reinit_audio_prompt_frozen_projection : bool (optional, default=False)
+    If True, reinitializes the frozen audio prompt projection layer.
+
+    Purpose:
+    - Disables voice cloning effects during inference
+    - Useful when exporting models where voice cloning is not desired
+
+    When enabled:
+    - The projection matrix is replaced with random weights
+    - Speaker conditioning is effectively disabled
+
+Notes
+-----
+- Distributed checkpoints are detected when `tts_ckpt_path` or `stt_ckpt_path`
+  points to a directory containing checkpoint shards.
+- safetensors checkpoints are supported when the path ends with `.safetensors`.
+- Non-distributed PyTorch checkpoints are loaded normally with state_dict mapping.
+"""
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Dict
+
+import torch
+from omegaconf import OmegaConf
+
+from nemo.collections.audio.parts.utils.transforms import resample
+from nemo.collections.speechlm2.models.duplex_ear_tts import load_audio_librosa
+from nemo.collections.speechlm2.models.nemotron_voicechat import NemotronVoiceChat
+from nemo.core.config import hydra_runner
+from nemo.utils import logging
+
+
+@dataclass
+class HfExportConfig:
+    """
+    Configuration for exporting Speech-to-Chat/TTS models to HuggingFace format.
+
+    Attributes:
+        tts_ckpt_path:
+            Path to the TTS checkpoint (PyTorch Lightning ckpt or directory).
+        tts_ckpt_config:
+            Path to the experiment config used to instantiate the TTS model.
+        stt_ckpt_path:
+            Path to the STT checkpoint (PyTorch Lightning ckpt or directory).
+        stt_ckpt_config:
+            Path to the experiment config used to instantiate the STT model.
+        output_dir:
+            Directory where the HuggingFace-compatible checkpoint will be stored.
+        dtype:
+            Target dtype for storing parameters (default: float32).
+        register_speaker_dict:
+            Dictionary mapping speaker names to reference audio paths.
+        reinit_audio_prompt_frozen_projection:
+            If True, reinitialize the frozen projection to disable voice cloning.
+    """
+
+    tts_ckpt_path: str
+    tts_ckpt_config: str
+    stt_ckpt_path: str
+    stt_ckpt_config: str
+    output_dir: str
+    dtype: str = "float32"
+    register_speaker_dict: Dict[str, str] = field(default_factory=dict)
+    reinit_audio_prompt_frozen_projection: bool = False
+
+
+def load_checkpoint_by_module(model: torch.nn.Module, checkpoint_path: str, module_name: str):
+    """
+    Load checkpoint weights into a submodule of the model.
+
+    Supports:
+        - Distributed checkpoints (FSDP/TP)
+        - safetensors checkpoints
+        - standard PyTorch checkpoints
+
+    Args:
+        model:
+            Parent torch module.
+        checkpoint_path:
+            Path to checkpoint or distributed checkpoint directory.
+        module_name:
+            Name of the submodule (e.g., "stt_model" or "tts_model").
+    """
+
+    module = getattr(model, module_name)
+
+    # Distributed checkpoint (FSDP/TP)
+    if Path(checkpoint_path).is_dir():
+        from torch.distributed.checkpoint import load
+
+        state_dict = {"state_dict": module.state_dict()}
+        load(state_dict, checkpoint_id=checkpoint_path)
+        module.load_state_dict(state_dict["state_dict"])
+        return
+
+    # safetensors checkpoint
+    if ".safetensors" in checkpoint_path:
+        if hasattr(model, "init_from_safetensors_ckpt"):
+            model.init_from_safetensors_ckpt(checkpoint_path, prefix=f"{module_name}.")
+            return
+        raise RuntimeError("Model does not support safetensors loader.")
+
+    # Standard PyTorch checkpoint
+    ckpt_data = torch.load(checkpoint_path, map_location="cpu")
+    sd = ckpt_data.get("state_dict", ckpt_data)
+
+    try:
+        module.load_state_dict(sd, strict=False)
+    except Exception:
+        # Fallback: filter keys matching module prefix
+        prefix = module_name + "."
+        filtered = {k[len(prefix) :]: v for k, v in sd.items() if k.startswith(prefix)}
+        module.load_state_dict(filtered, strict=False)
+
+
+@hydra_runner(config_name="HfExportConfig", schema=HfExportConfig)
+def main(cfg: HfExportConfig):
+    """
+    Entry point for exporting Nemotron VoiceChat checkpoint.
+
+    Steps:
+        1. Load STT and TTS experiment configs.
+        2. Instantiate NemotronVoiceChat model.
+        3. Load STT/TTS checkpoints.
+        4. Optionally register speaker audio prompts.
+        5. Save HuggingFace-compatible checkpoint.
+
+    Args:
+        cfg:
+            Hydra configuration object containing export parameters.
+    """
+
+    # Load STT configuration
+    stt_model_cfg = OmegaConf.load(cfg.stt_ckpt_config)
+
+    # Prevent model from reloading pretrained perception module during export.
+    # Some STT configs define these fields.
+    # If it exists, we set it to None so the exporter uses the checkpoint weights only.
+    if hasattr(stt_model_cfg.model, "pretrained_perception_from_s2s"):
+        stt_model_cfg.model.pretrained_perception_from_s2s = None
+    if hasattr(stt_model_cfg.model, "pretrained_s2s_model"):
+        stt_model_cfg.model.pretrained_s2s_model = None
+
+    # Load TTS experiment configuration
+    tts_model_cfg = OmegaConf.load(cfg.tts_ckpt_config)
+
+    # Disable codec model reloading during export.
+    # Some recipes reference an external pretrained codec.
+    # Setting this to None forces the exporter to use checkpoint weights only,
+    # avoiding additional model downloads or initialization.
+    if hasattr(tts_model_cfg.model, "pretrained_codec_model"):
+        tts_model_cfg.model.pretrained_codec_model = None
+
+    # Joint model configuration
+    model_cfg = {
+        "model": {
+            "scoring_asr": "stt_en_fastconformer_transducer_large",
+            "stt": stt_model_cfg,
+            "speech_generation": tts_model_cfg,
+        },
+        "data": {
+            "frame_length": 0.08,
+            "source_sample_rate": stt_model_cfg.data.source_sample_rate,
+            "target_sample_rate": tts_model_cfg.data.target_sample_rate,
+        },
+        "exp_manager": {"explicit_log_dir": " "},
+        "torch_dtype": cfg.dtype,
+    }
+
+    model_cfg = OmegaConf.create(model_cfg)
+
+    # Instantiate model
+    model = NemotronVoiceChat(OmegaConf.to_container(model_cfg, resolve=True))
+
+    # Load checkpoints
+    load_checkpoint_by_module(model, cfg.stt_ckpt_path, "stt_model")
+    load_checkpoint_by_module(model, cfg.tts_ckpt_path, "tts_model")
+
+    # Register inference speakers (voice cloning)
+    if cfg.register_speaker_dict:
+        model.tts_model.to(model.device)
+        for speaker_name, audio_path in cfg.register_speaker_dict.items():
+            speaker_audio, sr = load_audio_librosa(audio_path)
+            speaker_audio = resample(speaker_audio, sr, model.tts_model.target_sample_rate).to(model.device)
+
+            speaker_audio_lens = (
+                torch.tensor([speaker_audio.size(1)]).long().repeat(speaker_audio.size(0)).to(model.device)
+            )
+
+            model.tts_model.set_audio_prompt_lantent(
+                speaker_audio,
+                speaker_audio_lens,
+                system_prompt=None,
+                batch_size=1,
+                name=speaker_name,
+            )
+            logging.info(f"Speaker {speaker_name} registered!")
+
+    # Optionally reinitialize projection (disables voice cloning)
+    if cfg.reinit_audio_prompt_frozen_projection:
+        D = model.tts_model.tts_model.hidden_size
+        model.tts_model.tts_model.audio_prompt_projection_W.copy_(
+            torch.randn(
+                D,
+                D,
+                device=model.tts_model.tts_model.audio_prompt_projection_W.device,
+                dtype=model.tts_model.tts_model.audio_prompt_projection_W.dtype,
+            )
+        )
+        logging.info("Audio frozen projection reinitialized!")
+
+    # Cast model to target dtype and save HuggingFace checkpoint
+    model = model.to(getattr(torch, cfg.dtype))
+    model.save_pretrained(cfg.output_dir, config=OmegaConf.to_container(model_cfg, resolve=True))
+
+    logging.info(f"HuggingFace-compatible checkpoint saved at: {cfg.output_dir}")
+
+
+if __name__ == "__main__":
+    main()

examples/speechlm2/s2s_duplex_speech_decoder_train.py

@@ -0,0 +1,55 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+
+import torch
+from lightning.pytorch import Trainer
+from omegaconf import OmegaConf
+
+from nemo.collections.speechlm2 import DataModule, DuplexS2SDataset, DuplexS2SSpeechDecoderModel
+from nemo.core.config import hydra_runner
+from nemo.utils.exp_manager import exp_manager
+from nemo.utils.trainer_utils import resolve_trainer_cfg
+
+torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
+
+
+@hydra_runner(config_path="conf", config_name="s2s_duplex_speech_decoder")
+def train(cfg):
+    OmegaConf.resolve(cfg)
+    torch.distributed.init_process_group(backend="nccl")
+    torch.set_float32_matmul_precision("medium")
+    torch.backends.cudnn.allow_tf32 = True
+    trainer = Trainer(**resolve_trainer_cfg(cfg.trainer))
+    log_dir = exp_manager(trainer, cfg.get("exp_manager", None))
+    OmegaConf.save(cfg, log_dir / "exp_config.yaml")
+
+    with trainer.init_module():
+        model = DuplexS2SSpeechDecoderModel(OmegaConf.to_container(cfg.model, resolve=True))
+
+    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)
+
+    trainer.fit(model, datamodule)
+
+
+if __name__ == "__main__":
+    train()

examples/speechlm2/s2s_duplex_train.py

@@ -0,0 +1,55 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+
+import torch
+from lightning.pytorch import Trainer
+from omegaconf import OmegaConf
+
+from nemo.collections.speechlm2 import DataModule, DuplexS2SDataset, DuplexS2SModel
+from nemo.core.config import hydra_runner
+from nemo.utils.exp_manager import exp_manager
+from nemo.utils.trainer_utils import resolve_trainer_cfg
+
+torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
+
+
+@hydra_runner(config_path="conf", config_name="s2s_duplex")
+def train(cfg):
+    OmegaConf.resolve(cfg)
+    torch.distributed.init_process_group(backend="nccl")
+    torch.set_float32_matmul_precision("medium")
+    torch.backends.cudnn.allow_tf32 = True
+    trainer = Trainer(**resolve_trainer_cfg(cfg.trainer))
+    log_dir = exp_manager(trainer, cfg.get("exp_manager", None))
+    OmegaConf.save(cfg, log_dir / "exp_config.yaml")
+
+    with trainer.init_module():
+        model = DuplexS2SModel(OmegaConf.to_container(cfg.model, resolve=True))
+
+    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)
+
+    trainer.fit(model, datamodule)
+
+
+if __name__ == "__main__":
+    train()

examples/speechlm2/salm_asr_decoder_train.py

@@ -0,0 +1,48 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+
+import torch
+from lightning.pytorch import Trainer
+from omegaconf import OmegaConf
+
+from nemo.collections.speechlm2 import DataModule, SALMDataset
+from nemo.collections.speechlm2.models.salm_asr_decoder import SALMWithAsrDecoder
+from nemo.core.config import hydra_runner
+from nemo.utils.exp_manager import exp_manager
+from nemo.utils.trainer_utils import resolve_trainer_cfg
+
+torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
+
+
+@hydra_runner(config_path="conf", config_name="salm")
+def train(cfg):
+    OmegaConf.resolve(cfg)
+    torch.distributed.init_process_group(backend="nccl")
+    torch.set_float32_matmul_precision("medium")
+    trainer = Trainer(**resolve_trainer_cfg(cfg.trainer))
+    log_dir = exp_manager(trainer, cfg.get("exp_manager", None))
+    OmegaConf.save(cfg, log_dir / "exp_config.yaml")
+
+    with trainer.init_module():
+        model = SALMWithAsrDecoder(OmegaConf.to_container(cfg.model, resolve=True))
+
+    dataset = SALMDataset(tokenizer=model.tokenizer)
+    datamodule = DataModule(cfg.data, tokenizer=model.tokenizer, dataset=dataset)
+
+    trainer.fit(model, datamodule)
+
+
+if __name__ == "__main__":
+    train()

examples/speechlm2/salm_eval.py

@@ -0,0 +1,234 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import json
+from dataclasses import dataclass
+from time import perf_counter
+from typing import Optional
+
+import lhotse.dataset
+import torch
+from lhotse import CutSet
+from lhotse.serialization import SequentialJsonlWriter
+from omegaconf import OmegaConf
+from transformers import GenerationConfig
+from whisper_normalizer.basic import BasicTextNormalizer
+from whisper_normalizer.english import EnglishTextNormalizer
+
+from nemo.collections.asr.metrics.wer import word_error_rate_detail
+from nemo.collections.common.data.lhotse.cutset import guess_parse_cutset
+from nemo.collections.speechlm2.models import SALM, SALMWithAsrDecoder
+from nemo.core.config import hydra_runner
+from nemo.utils import logging
+from nemo.utils.get_rank import is_global_rank_zero
+
+
+class ToAudio(torch.utils.data.Dataset):
+    def __getitem__(self, cuts: CutSet):
+        audios, audio_lens = cuts.load_audio(collate=True)
+        return {"cuts": cuts, "audios": audios, "audio_lens": audio_lens}
+
+
+def _resolve_model_cls(pretrained_name: str, use_asr_decoder: bool, use_nemo_automodel: bool | None):
+    """Pick model class. Auto-detects from config.json when use_nemo_automodel is None."""
+    if use_asr_decoder:
+        return SALMWithAsrDecoder
+    if use_nemo_automodel is None:
+        # Auto-detect: peek at config.json
+        from transformers.utils import cached_file
+
+        config_path = cached_file(
+            pretrained_name,
+            "config.json",
+            _raise_exceptions_for_missing_entries=False,
+            _raise_exceptions_for_connection_errors=False,
+        )
+        if config_path is not None:
+            with open(config_path) as f:
+                use_nemo_automodel = json.load(f).get("use_nemo_automodel", False)
+        else:
+            use_nemo_automodel = False
+    if use_nemo_automodel:
+        from nemo.collections.speechlm2.models import SALMAutomodel
+
+        return SALMAutomodel
+    return SALM
+
+
+@dataclass
+class SalmEvalConfig:
+    pretrained_name: str
+    inputs: str
+    batch_size: int = 64
+    max_new_tokens: int = 128
+    output_manifest: Optional[str] = "generations.jsonl"
+    verbose: bool = True
+    use_normalizer: Optional[str] = "english"  # "english", "basic", or "none" / "None"
+    device: str = "cuda"
+    dtype: str = "bfloat16"
+    extra_eos_tokens: Optional[list[str]] = None
+    system_prompt: Optional[str] = None
+    user_prompt: Optional[str] = None
+    enable_thinking: Optional[bool] = None
+    use_asr_decoder: bool = False  # set this to True if using SALMWithAsrDecoder
+    use_nemo_automodel: Optional[bool] = None  # None = auto-detect from config.json
+    # Parallelism sizes for distributed inference (launch with torchrun)
+    tp_size: int = 1
+    ep_size: int = 1
+    pp_size: int = 1
+    cp_size: int = 1
+
+
+@hydra_runner(config_name="SalmEvalConfig", schema=SalmEvalConfig)
+def main(cfg: SalmEvalConfig):
+    logging.info(f'Hydra config:\n{OmegaConf.to_yaml(cfg)}')
+
+    is_distributed = any(s > 1 for s in [cfg.tp_size, cfg.ep_size, cfg.pp_size, cfg.cp_size])
+    model_cls = _resolve_model_cls(cfg.pretrained_name, cfg.use_asr_decoder, cfg.use_nemo_automodel)
+
+    if is_distributed and model_cls is SALM:
+        raise RuntimeError(
+            "Distributed inference requires SALMAutomodel. Set use_nemo_automodel=true or use a checkpoint "
+            "exported from SALMAutomodel."
+        )
+
+    if is_distributed:
+        from nemo.collections.speechlm2.parts.parallel import setup_distributed
+
+        strategy = setup_distributed(
+            tp_size=cfg.tp_size, ep_size=cfg.ep_size, pp_size=cfg.pp_size, cp_size=cfg.cp_size
+        )
+        model = model_cls.from_pretrained(
+            cfg.pretrained_name,
+            device_mesh=strategy.device_mesh,
+            distributed_config=strategy.distributed_config,
+            moe_config=strategy.moe_config,
+            moe_mesh=strategy.moe_mesh,
+            torch_dtype=cfg.dtype,
+        )
+    else:
+        model = model_cls.from_pretrained(cfg.pretrained_name)
+        model = model.to(getattr(torch, cfg.dtype)).to(cfg.device)
+    model = model.eval()
+
+    cuts = guess_parse_cutset(cfg.inputs).sort_by_duration()
+    dloader = torch.utils.data.DataLoader(
+        dataset=ToAudio(),
+        # rank=0 world_size=1 hardcoded so lhotse doesn't accidentally auto-split batches in model parallel settings
+        sampler=lhotse.dataset.DynamicCutSampler(cuts, max_cuts=cfg.batch_size, rank=0, world_size=1),
+        num_workers=1,
+        batch_size=None,
+    )
+
+    normalizer = {"english": EnglishTextNormalizer(), "basic": BasicTextNormalizer()}.get(
+        cfg.use_normalizer, lambda x: x
+    )
+
+    eos_tokens = [model.text_eos_id]
+    if cfg.extra_eos_tokens is not None:
+        for t in cfg.extra_eos_tokens:
+            tid = model.tokenizer.token_to_id(t)
+            assert tid is not None, f"Token '{t}' is not in the model's vocabulary."
+            eos_tokens.append(tid)
+
+    # Construct the prompt from ASR data of the form.
+    # Optional system prompt goes first.
+    prompt = []
+    if cfg.system_prompt is not None:
+        prompt.append({"role": "system", "content": cfg.system_prompt})
+    # If no user prompt is provided, just use the audio placeholder.
+    content = model.audio_locator_tag
+    # Otherwise:
+    # * if user prompt already has audio placeholder, add it as-is,
+    # * if not, append audio placeholder at the end of user prompt
+    if cfg.user_prompt is not None:
+        content = cfg.user_prompt
+        if model.audio_locator_tag not in content:
+            content = f"{content} {model.audio_locator_tag}"
+    prompt.append({"role": "user", "content": content})
+
+    refs = []
+    hyps = []
+    input_durations = []
+    infer_durations = []
+    for batch_idx, batch in enumerate(dloader):
+        ts = perf_counter()
+        answer_ids = model.generate(
+            prompts=[prompt] * len(batch["cuts"]),  # identical prompt for each example
+            audios=batch["audios"].to(model.device, non_blocking=True),
+            audio_lens=batch["audio_lens"].to(model.device, non_blocking=True),
+            generation_config=GenerationConfig(
+                max_new_tokens=cfg.max_new_tokens,
+                bos_token_id=model.text_bos_id,
+                eos_token_id=eos_tokens,
+                pad_token_id=model.text_pad_id,
+            ),
+            enable_thinking=cfg.enable_thinking,
+        )
+        answer_ids = answer_ids.cpu()
+        batch_infer_duration = perf_counter() - ts
+
+        batch_duration = sum(c.duration for c in batch["cuts"])
+        batch_refs = [normalizer(cut.supervisions[0].text) for cut in batch["cuts"]]
+        batch_hyps = [
+            normalizer(model.tokenizer.ids_to_text(parse_hyp(ans, eos_tokens)).strip()) for ans in answer_ids
+        ]
+        if cfg.verbose:
+            batch_wer, _, nins, ndel, nsub = word_error_rate_detail(batch_hyps, batch_refs)
+            batch_rtfx = batch_duration / batch_infer_duration
+            logging.info(
+                f"Batch {batch_idx}: WER={batch_wer:.2%} [ins={nins:.2%} del={ndel:.2%} sub={nsub:.2%}] RTFx={batch_rtfx:.1f}"
+            )
+
+        refs.extend(batch_refs)
+        hyps.extend(batch_hyps)
+        input_durations.append(batch_duration)
+        infer_durations.append(batch_infer_duration)
+
+    wer, _, nins, ndel, nsub = word_error_rate_detail(hypotheses=hyps, references=refs, use_cer=False)
+    rtfx = sum(input_durations) / sum(infer_durations)
+    logging.info(f"WER: {wer:.2%} [ins={nins:.2%} del={ndel:.2%} sub={nsub:.2%}]")
+    logging.info(f"RTFx: {rtfx:.1f}")
+
+    with _create_output_writer(cfg.output_manifest) as writer:
+        for cut, ref, hyp in zip(cuts, refs, hyps):
+            writer.write({"id": cut.id, "duration": cut.duration, "text": ref, "pred_text": hyp})
+
+
+def parse_hyp(answer: torch.Tensor, eos_tokens: list[int]):
+    end = torch.isin(answer, torch.tensor(eos_tokens)).nonzero(as_tuple=True)[0]
+    if end.numel() == 0:
+        return answer
+    end = end[0]
+    return answer[:end]
+
+
+class _NullWriter:
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        return False
+
+    def write(self, data):
+        pass
+
+
+def _create_output_writer(output_manifest: Optional[str]):
+    if output_manifest is None or not is_global_rank_zero():
+        return _NullWriter()
+    return SequentialJsonlWriter(output_manifest)
+
+
+if __name__ == '__main__':
+    main()

examples/speechlm2/salm_generate.py

@@ -0,0 +1,280 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import json
+from dataclasses import dataclass
+from functools import partial
+from time import perf_counter
+from typing import Optional
+
+import lhotse.dataset
+import torch
+from lhotse import CutSet, fastcopy
+from lhotse.dataset import IterableDatasetWrapper
+from lhotse.serialization import SequentialJsonlWriter
+from omegaconf import OmegaConf
+from transformers import GenerationConfig
+
+from nemo.collections.common.data.lhotse import NeMoMultimodalConversation
+from nemo.collections.common.data.lhotse.cutset import cut_to_conversation, guess_parse_cutset
+from nemo.collections.common.data.lhotse.dataloader import tokenize_with_prompt
+from nemo.collections.common.data.lhotse.text_adapters import AudioTurn, TextTurn
+from nemo.collections.speechlm2 import SALM, SALMDataset
+from nemo.collections.speechlm2.models.salm_asr_decoder import SALMWithAsrDecoder
+from nemo.core.config import hydra_runner
+from nemo.utils import logging
+from nemo.utils.get_rank import is_global_rank_zero
+
+
+def _resolve_model_cls(pretrained_name: str, use_asr_decoder: bool, use_nemo_automodel: bool | None):
+    """Pick model class. Auto-detects from config.json when use_nemo_automodel is None."""
+    if use_asr_decoder:
+        return SALMWithAsrDecoder
+    if use_nemo_automodel is None:
+        # Auto-detect: peek at config.json
+        from transformers.utils import cached_file
+
+        config_path = cached_file(
+            pretrained_name,
+            "config.json",
+            _raise_exceptions_for_missing_entries=False,
+            _raise_exceptions_for_connection_errors=False,
+        )
+        if config_path is not None:
+            with open(config_path) as f:
+                use_nemo_automodel = json.load(f).get("use_nemo_automodel", False)
+        else:
+            use_nemo_automodel = False
+    if use_nemo_automodel:
+        from nemo.collections.speechlm2.models import SALMAutomodel
+
+        return SALMAutomodel
+    return SALM
+
+
+@dataclass
+class SalmEvalConfig:
+    pretrained_name: str
+    inputs: str
+    batch_size: int = 64
+    max_new_tokens: int = 128
+    output_manifest: str = "generations.jsonl"
+    verbose: bool = True
+    device: str = "cuda"
+    dtype: str = "bfloat16"
+    extra_eos_tokens: Optional[list[str]] = None
+    system_prompt: Optional[str] = None
+    user_prompt: Optional[str] = None
+    enable_thinking: Optional[bool] = None
+    use_asr_decoder: bool = False  # set this to True if using SALMWithAsrDecoder
+    use_nemo_automodel: Optional[bool] = None  # None = auto-detect from config.json
+    # Parallelism sizes for distributed inference (launch with torchrun)
+    tp_size: int = 1
+    ep_size: int = 1
+    pp_size: int = 1
+    cp_size: int = 1
+
+
+@hydra_runner(config_name="SalmEvalConfig", schema=SalmEvalConfig)
+def main(cfg: SalmEvalConfig):
+    logging.info(f"Hydra config:\n{OmegaConf.to_yaml(cfg)}")
+
+    is_distributed = any(s > 1 for s in [cfg.tp_size, cfg.ep_size, cfg.pp_size, cfg.cp_size])
+    model_cls = _resolve_model_cls(cfg.pretrained_name, cfg.use_asr_decoder, cfg.use_nemo_automodel)
+
+    if is_distributed and model_cls is SALM:
+        raise RuntimeError(
+            "Distributed inference requires SALMAutomodel. Set use_nemo_automodel=true or use a checkpoint "
+            "exported from SALMAutomodel."
+        )
+
+    if is_distributed:
+        from nemo.collections.speechlm2.parts.parallel import setup_distributed
+
+        strategy = setup_distributed(
+            tp_size=cfg.tp_size, ep_size=cfg.ep_size, pp_size=cfg.pp_size, cp_size=cfg.cp_size
+        )
+        model = model_cls.from_pretrained(
+            cfg.pretrained_name,
+            device_mesh=strategy.device_mesh,
+            distributed_config=strategy.distributed_config,
+            moe_config=strategy.moe_config,
+            moe_mesh=strategy.moe_mesh,
+            torch_dtype=cfg.dtype,
+        )
+    else:
+        model = model_cls.from_pretrained(cfg.pretrained_name)
+        model = model.to(getattr(torch, cfg.dtype)).to(cfg.device)
+    model = model.eval()
+
+    conversations = (
+        guess_parse_cutset(cfg.inputs)
+        .map(
+            partial(
+                cut_to_conversation,
+                audio_locator_tag=model.audio_locator_tag,
+                token_equivalent_duration=model.token_equivalent_duration,
+            )
+        )
+        .map(
+            partial(replace_audio_locator_tag, audio_locator_tag=model.audio_locator_tag),
+            apply_fn=None,
+        )
+        .map(
+            partial(set_token_equivalent_duration, token_equivalent_duration=model.token_equivalent_duration),
+            apply_fn=None,
+        )
+        .map(
+            partial(attach_system_and_user_turns, system_prompt=cfg.system_prompt, user_prompt=cfg.user_prompt),
+            apply_fn=None,
+        )
+        .map(strip_response_if_any, apply_fn=None)
+        .map(
+            partial(
+                tokenize_with_prompt,
+                tokenizer=model.tokenizer,
+                prompt_format=model.cfg.prompt_format,
+                enable_thinking=cfg.enable_thinking,
+            ),
+            apply_fn=None,
+        )
+    )
+    conversations = sort_by_length(conversations)
+    dloader = torch.utils.data.DataLoader(
+        dataset=IterableDatasetWrapper(
+            dataset=SALMDataset(model.tokenizer),
+            # rank=0 world_size=1 hardcoded so lhotse doesn't accidentally auto-split batches in model parallel settings
+            sampler=lhotse.dataset.DynamicCutSampler(conversations, max_cuts=cfg.batch_size, rank=0, world_size=1),
+        ),
+        num_workers=1,
+        batch_size=None,
+    )
+
+    eos_tokens = [model.text_eos_id]
+    if cfg.extra_eos_tokens is not None:
+        for t in cfg.extra_eos_tokens:
+            tid = model.tokenizer.token_to_id(t)
+            assert tid is not None, f"Token '{t}' is not in the model's vocabulary."
+            eos_tokens.append(tid)
+
+    num_answer_tokens = []
+    infer_durations = []
+    with _create_output_writer(cfg.output_manifest) as writer:
+        for batch_idx, batch in enumerate(dloader):
+            ts = perf_counter()
+            answer_ids = model.generate(
+                prompts=batch["input_ids"].to(model.device, non_blocking=True),
+                audios=batch["audios"].to(model.device, non_blocking=True),
+                audio_lens=batch["audio_lens"].to(model.device, non_blocking=True),
+                generation_config=GenerationConfig(
+                    max_new_tokens=cfg.max_new_tokens,
+                    bos_token_id=model.text_bos_id,
+                    eos_token_id=eos_tokens,
+                    pad_token_id=model.text_pad_id,
+                ),
+            )
+            answer_ids = answer_ids.cpu()
+            batch_infer_duration = perf_counter() - ts
+
+            batch_contexts = [model.tokenizer.ids_to_text(example) for example in batch["input_ids"]]
+            answer_ids = [parse_hyp(ans, eos_tokens) for ans in answer_ids]
+            batch_num_answer_tokens = [len(ans) for ans in answer_ids]
+            batch_answers = [model.tokenizer.ids_to_text(ans) for ans in answer_ids]
+            for conv, ctx, ans in zip(batch["conversations"], batch_contexts, batch_answers):
+                conv.turns.append(TextTurn(role="assistant", value=ans))
+                for k, v in list(conv.custom.items()):
+                    if isinstance(v, torch.Tensor):
+                        del conv.custom[k]
+                writer.write(conv.to_dict())
+
+            num_answer_tokens.extend(batch_num_answer_tokens)
+            infer_durations.append(batch_infer_duration)
+            if cfg.verbose:
+                batch_token_per_second = sum(batch_num_answer_tokens) / batch_infer_duration
+                logging.info(f"Batch {batch_idx}: TPS={batch_token_per_second:.2f}")
+
+    rtfx = sum(num_answer_tokens) / sum(infer_durations)
+    logging.info(f"TPS: {rtfx:.2f}")
+
+
+def replace_audio_locator_tag(
+    conversation: NeMoMultimodalConversation, audio_locator_tag: str
+) -> NeMoMultimodalConversation:
+    for turn in conversation.turns:
+        if isinstance(turn, AudioTurn):
+            turn.audio_locator_tag = audio_locator_tag
+    return conversation
+
+
+def set_token_equivalent_duration(
+    conversation: NeMoMultimodalConversation, token_equivalent_duration: float
+) -> NeMoMultimodalConversation:
+    conversation.token_equivalent_duration = token_equivalent_duration
+    return conversation
+
+
+def attach_system_and_user_turns(
+    conversation: NeMoMultimodalConversation, system_prompt: str | None = None, user_prompt: str | None = None
+) -> NeMoMultimodalConversation:
+    if system_prompt is None and user_prompt is None:
+        return conversation
+    turns = conversation.turns
+    # Attach user prompt only when no user turn with a text prompt exists.
+    if user_prompt is not None and not any(isinstance(t, TextTurn) and t.role == "user" for t in turns):
+        turns = [TextTurn(role="user", value=user_prompt)] + turns
+    # Attach system prompt only when no system prompt already exists.
+    if system_prompt is not None and not any(t.role == "system" for t in turns):
+        turns = [TextTurn(role="system", value=system_prompt)] + turns
+    return fastcopy(conversation, turns=turns)
+
+
+def strip_response_if_any(
+    conversation: NeMoMultimodalConversation,
+) -> NeMoMultimodalConversation:
+    turns = conversation.turns
+    while turns[-1].role == "assistant":
+        turns = turns[:-1]
+    return fastcopy(conversation, turns=turns)
+
+
+def sort_by_length(conversations: CutSet) -> CutSet:
+    return CutSet(sorted(conversations, key=lambda c: c.total_length, reverse=True))
+
+
+def parse_hyp(answer: torch.Tensor, eos_tokens: list[int]):
+    end = torch.isin(answer, torch.tensor(eos_tokens)).nonzero(as_tuple=True)[0]
+    if end.numel() == 0:
+        return answer
+    end = end[0]
+    return answer[:end]
+
+
+class _NullWriter:
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        return False
+
+    def write(self, data):
+        pass
+
+
+def _create_output_writer(output_manifest: Optional[str]):
+    if output_manifest is None or not is_global_rank_zero():
+        return _NullWriter()
+    return SequentialJsonlWriter(output_manifest)
+
+
+if __name__ == '__main__':
+    main()

examples/speechlm2/salm_train.py

@@ -0,0 +1,56 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+
+import torch
+from lightning.pytorch import Trainer, seed_everything
+from omegaconf import OmegaConf
+
+from nemo.collections.speechlm2 import SALM, DataModule, SALMDataset
+from nemo.core.config import hydra_runner
+from nemo.utils.exp_manager import exp_manager
+from nemo.utils.trainer_utils import resolve_trainer_cfg
+
+if torch.cuda.is_available():
+    torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
+
+
+@hydra_runner(config_path="conf", config_name="salm")
+def train(cfg):
+    OmegaConf.resolve(cfg)
+    if torch.cuda.is_available():
+        torch.distributed.init_process_group(backend="nccl")
+    seed_everything(cfg.data.train_ds.seed)
+    torch.set_float32_matmul_precision("medium")
+    trainer = Trainer(**resolve_trainer_cfg(cfg.trainer))
+    log_dir = exp_manager(trainer, cfg.get("exp_manager", None))
+    OmegaConf.save(cfg, log_dir / "exp_config.yaml")
+
+    model_cls = SALM
+    if cfg.model.get("use_nemo_automodel", False):
+        from nemo.collections.speechlm2 import SALMAutomodel
+
+        model_cls = SALMAutomodel
+
+    with trainer.init_module():
+        model = model_cls(OmegaConf.to_container(cfg.model, resolve=True))
+
+    dataset = SALMDataset(tokenizer=model.tokenizer, multispeaker_cfg=cfg.data.get("multispeaker_cfg", None))
+    datamodule = DataModule(cfg.data, tokenizer=model.tokenizer, dataset=dataset)
+
+    trainer.fit(model, datamodule)
+
+
+if __name__ == "__main__":
+    train()

examples/speechlm2/stt_duplex_infer.py

@@ -0,0 +1,55 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+
+import torch
+from lightning.pytorch import Trainer
+from omegaconf import OmegaConf, open_dict
+
+from nemo.collections.speechlm2 import DataModule, DuplexSTTDataset, DuplexSTTModel
+from nemo.core.config import hydra_runner
+from nemo.utils.exp_manager import exp_manager
+from nemo.utils.trainer_utils import resolve_trainer_cfg
+
+torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
+
+
+@hydra_runner(config_path="conf", config_name="duplex_stt")
+def inference(cfg):
+    OmegaConf.resolve(cfg)
+    torch.distributed.init_process_group(backend="nccl")
+    torch.set_float32_matmul_precision("medium")
+    torch.backends.cudnn.allow_tf32 = True
+    trainer = Trainer(**resolve_trainer_cfg(cfg.trainer))
+    log_dir = exp_manager(trainer, cfg.get("exp_manager", None))
+    OmegaConf.save(cfg, log_dir / "exp_config.yaml")
+
+    with trainer.init_module():
+        model = DuplexSTTModel(OmegaConf.to_container(cfg.model, resolve=True))
+
+    dataset = DuplexSTTDataset(
+        tokenizer=model.tokenizer,
+        frame_length=cfg.data.frame_length,
+        source_sample_rate=cfg.data.source_sample_rate,
+        input_roles=cfg.data.input_roles,
+        output_roles=cfg.data.output_roles,
+        cfg=OmegaConf.to_container(cfg.data, resolve=True),
+    )
+    datamodule = DataModule(cfg.data, tokenizer=model.tokenizer, dataset=dataset)
+
+    trainer.validate(model, datamodule)
+
+
+if __name__ == "__main__":
+    inference()

examples/speechlm2/stt_duplex_train.py

@@ -0,0 +1,71 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import multiprocessing as mp
+import os
+
+import torch
+from lightning.pytorch import Trainer
+from lightning.pytorch.callbacks import ModelCheckpoint
+from omegaconf import OmegaConf, open_dict
+
+from nemo.collections.speechlm2 import DataModule, DuplexSTTDataset, DuplexSTTModel
+from nemo.core.config import hydra_runner
+from nemo.utils.exp_manager import exp_manager
+from nemo.utils.trainer_utils import resolve_trainer_cfg
+
+# Set multiprocessing start method to 'spawn' for CUDA compatibility with DataLoader workers
+# This prevents "Cannot re-initialize CUDA in forked subprocess" errors
+try:
+    mp.set_start_method('spawn', force=True)
+except RuntimeError:
+    pass  # Start method already set
+
+torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
+
+
+@hydra_runner(config_path="conf", config_name="duplex_stt")
+def train(cfg):
+    OmegaConf.resolve(cfg)
+    torch.distributed.init_process_group(backend="nccl")
+    torch.set_float32_matmul_precision("medium")
+    torch.backends.cudnn.allow_tf32 = True
+    trainer = Trainer(**resolve_trainer_cfg(cfg.trainer))
+    log_dir = exp_manager(trainer, cfg.get("exp_manager", None))
+    OmegaConf.save(cfg, log_dir / "exp_config.yaml")
+
+    # avoid using `=` in the checkpoint name
+    for callback in trainer.callbacks:
+        if isinstance(callback, ModelCheckpoint):
+            callback.CHECKPOINT_EQUALS_CHAR = "-"
+
+    with trainer.init_module():
+        model = DuplexSTTModel(OmegaConf.to_container(cfg.model, resolve=True))
+
+    dataset = DuplexSTTDataset(
+        tokenizer=model.tokenizer,
+        frame_length=cfg.data.frame_length,
+        source_sample_rate=cfg.data.source_sample_rate,
+        input_roles=cfg.data.input_roles,
+        output_roles=cfg.data.output_roles,
+        aug_by_swap_role=cfg.data.get("aug_by_swap_role", False),
+        cfg=cfg.data,
+        model_cfg=cfg.model,
+    )
+    datamodule = DataModule(cfg.data, tokenizer=model.tokenizer, dataset=dataset)
+
+    trainer.fit(model, datamodule)
+
+
+if __name__ == "__main__":
+    train()

examples/speechlm2/to_hf.py

@@ -0,0 +1,360 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import json
+import os
+from copy import deepcopy
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import torch
+import torch.distributed as dist
+from omegaconf import DictConfig, OmegaConf
+from safetensors.torch import save_file
+
+from nemo.collections.speechlm2.parts.hf_hub import LLM_BACKBONE_DIR
+from nemo.core.classes.common import safe_instantiate
+from nemo.core.config import hydra_runner
+from nemo.utils.dtype import str_to_dtype
+from nemo.utils.model_utils import import_class_by_path
+
+
+@dataclass
+class HfExportConfig:
+    # Name of the model class to be imported, e.g. nemo.collections.speechlm2.models.SALM
+    class_path: str
+
+    # Path to PyTorch Lightning checkpoint file (normal ckpt) or directory (distributed ckpt)
+    ckpt_path: str
+
+    # Path to the experiment's config, used to instantiate the model class.
+    ckpt_config: str
+
+    # Path where we should save the HuggingFace Hub compatible checkpoint
+    output_dir: str
+
+    # Dtype used for stored parameters
+    dtype: str = "bfloat16"
+
+
+def load_checkpoint(model: torch.nn.Module, checkpoint_path: str) -> None:
+    if Path(checkpoint_path).is_dir():
+        from torch.distributed.checkpoint import load
+
+        state_dict = {"state_dict": model.state_dict()}
+        load(state_dict, checkpoint_id=checkpoint_path)
+        model.load_state_dict(state_dict["state_dict"])
+    else:
+        ckpt_data = torch.load(checkpoint_path, map_location="cpu")
+        model.load_state_dict(ckpt_data["state_dict"])
+
+
+def setup_distributed_from_config(strategy_cfg: dict) -> Any:
+    """Initialize torch.distributed and create a device mesh from a Hydra strategy config.
+
+    Instantiates the strategy from the trainer config dict (as found in the
+    experiment YAML), initializes the process group, resolves automodel
+    configs, and calls :meth:`strategy.create_device_mesh`.
+
+    Returns:
+        An :class:`AutomodelParallelStrategy` with device_mesh ready.
+    """
+    from nemo.utils.trainer_utils import _resolve_automodel_configs
+
+    local_rank = int(os.environ.get("LOCAL_RANK", 0))
+    torch.cuda.set_device(local_rank)
+
+    strategy = safe_instantiate(strategy_cfg)
+    _resolve_automodel_configs(strategy)
+    strategy.create_device_mesh()
+    return strategy
+
+
+def consolidate_state_dict(model: torch.nn.Module) -> dict[str, torch.Tensor]:
+    """Gather a full (non-sharded) state dict from a model with DTensor parameters."""
+    from torch.distributed.tensor import DTensor
+
+    consolidated = {}
+    for key, value in model.state_dict().items():
+        if isinstance(value, DTensor):
+            consolidated[key] = value.full_tensor().cpu()
+        else:
+            consolidated[key] = value.cpu()
+    return consolidated
+
+
+def _canonical_torch_dtype_name(dtype: str | torch.dtype) -> str:
+    """Return the PyTorch dtype name accepted by Transformers configs."""
+    return str(str_to_dtype(dtype)).replace("torch.", "")
+
+
+def _hf_export_config(model: torch.nn.Module, dtype: str | torch.dtype) -> dict[str, Any]:
+    """Build the exported root config without mutating the training config."""
+    config = OmegaConf.to_container(model.cfg) if isinstance(model.cfg, DictConfig) else deepcopy(model.cfg)
+    dtype_name = _canonical_torch_dtype_name(dtype)
+    config["dtype"] = dtype_name
+    config["torch_dtype"] = dtype_name
+    return config
+
+
+def save_hf_checkpoint(model: torch.nn.Module, state_dict: dict, cfg: HfExportConfig) -> None:
+    """Save a consolidated state dict and model config in HuggingFace Hub format."""
+    output_dir = Path(cfg.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    target_dtype = str_to_dtype(cfg.dtype)
+    state_dict = {k: v.to(target_dtype) for k, v in state_dict.items()}
+
+    save_file(state_dict, output_dir / "model.safetensors")
+
+    config = _hf_export_config(model, cfg.dtype)
+    with open(output_dir / "config.json", "w") as f:
+        json.dump(config, f, indent=2)
+    save_llm_backbone_config(model, output_dir)
+
+
+def save_llm_backbone_config(model: torch.nn.Module, output_dir: str | Path) -> None:
+    """Save the original LLM config separately from the NeMo wrapper config."""
+    llm_config = getattr(getattr(model, "llm", None), "config", None)
+    if llm_config is None:
+        return
+
+    llm_backbone_dir = Path(output_dir) / LLM_BACKBONE_DIR
+    llm_backbone_dir.mkdir(parents=True, exist_ok=True)
+    llm_config.save_pretrained(str(llm_backbone_dir))
+
+
+def _detect_vllm_architecture(model_cfg: dict) -> str:
+    """Determine the vLLM plugin model class for the checkpoint.
+
+    The SALM plugin registers a single architecture name and selects between
+    transformer and hybrid backends at instantiation time, so this function
+    just verifies the backbone config is reachable and returns the unified
+    name; the hybrid-vs-transformer split is handled inside the plugin.
+
+    Raises:
+        ValueError: if the HF config can't be loaded or has no 'architectures'.
+    """
+    pretrained_llm = model_cfg.get("pretrained_llm", "")
+    try:
+        from transformers import AutoConfig
+
+        llm_cfg = AutoConfig.from_pretrained(pretrained_llm, trust_remote_code=True)
+    except Exception as e:
+        raise ValueError(
+            f"Could not load HF config for pretrained_llm={pretrained_llm!r}: {e}. "
+            f"Fix the 'pretrained_llm' field or ensure HF access during conversion."
+        ) from e
+
+    archs = getattr(llm_cfg, "architectures", [])
+    if not archs:
+        raise ValueError(f"HF config for {pretrained_llm!r} has empty 'architectures'.")
+
+    return "NeMoSpeechLMForConditionalGeneration"
+
+
+def prepare_for_vllm(output_dir: str, model_cfg: dict) -> None:
+    """Patch a saved checkpoint to be vLLM-ready.
+
+    Adds tokenizer (with audio token and chat template), patches config.json
+    with model_type/architectures, and writes generation_config.json.
+
+    Args:
+        output_dir: Path to the HuggingFace checkpoint directory.
+        model_cfg: Model config dict (from experiment YAML).
+
+    Raises:
+        ValueError: If ``pretrained_llm`` or ``audio_locator_tag`` is missing.
+    """
+    from transformers import AutoTokenizer
+
+    from nemo.utils import logging as LOG
+
+    output_dir = Path(output_dir)
+    pretrained_llm = model_cfg.get("pretrained_llm", "")
+    if not pretrained_llm:
+        raise ValueError("model config has no 'pretrained_llm'; cannot load tokenizer for vLLM")
+
+    # ``model.audio_locator_tag`` is the SoT for the audio placeholder;
+    # fail loud rather than default, since a mismatch is silent at inference.
+    audio_token = model_cfg.get("audio_locator_tag")
+    if not audio_token:
+        raise ValueError("model config has no 'audio_locator_tag' (set it in the training YAML).")
+
+    # 1. Patch config.json (arch, model_type, audio_locator_tag for vLLM plugin).
+    arch_model_cfg = dict(model_cfg)
+    llm_backbone_dir = output_dir / LLM_BACKBONE_DIR
+    if (llm_backbone_dir / "config.json").exists():
+        arch_model_cfg["pretrained_llm"] = str(llm_backbone_dir)
+    arch = _detect_vllm_architecture(arch_model_cfg)
+    config_path = output_dir / "config.json"
+    config = json.loads(config_path.read_text())
+    config["model_type"] = "nemo_speechlm"
+    config["architectures"] = [arch]
+    config["audio_locator_tag"] = audio_token
+    config_path.write_text(json.dumps(config, indent=2) + "\n")
+
+    # 2. Save tokenizer (backbone chat_template carries over via save_pretrained)
+    existing = [
+        f.name
+        for f in output_dir.iterdir()
+        if f.name in ("tokenizer_config.json", "tokenizer.json", "generation_config.json")
+    ]
+    if existing:
+        LOG.info("Overwriting existing files in %s: %s", output_dir, existing)
+    tok = AutoTokenizer.from_pretrained(pretrained_llm, trust_remote_code=True)
+    if audio_token not in tok.get_vocab():
+        tok.add_special_tokens({"additional_special_tokens": [audio_token]})
+    tok.save_pretrained(str(output_dir))
+    # Newer transformers splits long chat_template into a separate
+    # ``chat_template.jinja`` file; inline it back and drop the file.
+    tok_cfg_path = output_dir / "tokenizer_config.json"
+    tok_cfg = json.loads(tok_cfg_path.read_text())
+    jinja_file = output_dir / "chat_template.jinja"
+    if jinja_file.exists():
+        jinja_from_file = jinja_file.read_text()
+        if jinja_from_file.strip():
+            tok_cfg["chat_template"] = jinja_from_file
+        jinja_file.unlink()
+    # Normalize to dict form; transformers writes a list which HF loaders reject.
+    tok_cfg["extra_special_tokens"] = {"audio_token": audio_token}
+    # Some NeMo containers save a proprietary ``TokenizersBackend`` class
+    # unknown to HF; the underlying tokenizer.json is standard, so force
+    # the universal base class.
+    tok_cfg["tokenizer_class"] = "PreTrainedTokenizerFast"
+    tok_cfg_path.write_text(json.dumps(tok_cfg, indent=2) + "\n")
+
+    # 4. Minimal generation_config.json (EOS only; sampling params belong on
+    #    the server, not baked into the checkpoint).
+    gen_cfg = {"eos_token_id": [tok.eos_token_id]}
+    (output_dir / "generation_config.json").write_text(json.dumps(gen_cfg, indent=2) + "\n")
+
+
+def _try_prepare_for_vllm(output_dir: str, model_cfg: dict) -> None:
+    """Run vLLM prep; on ``ValueError``, warn and keep the HF-only output.
+
+    Backward compat for callers that never needed vLLM (e.g., NeMo SALM).
+    """
+    from nemo.utils import logging as LOG
+
+    try:
+        prepare_for_vllm(output_dir, model_cfg)
+    except ValueError as e:
+        LOG.warning(
+            "Checkpoint saved as HF-only; vLLM prep skipped: %s. "
+            "The checkpoint is still loadable by NeMo SALM and plain HF, but "
+            "is NOT vLLM-ready until prep succeeds.",
+            e,
+        )
+
+
+def _uses_automodel_parallel(strategy_cfg: dict) -> bool:
+    """Check if the strategy config targets AutomodelParallelStrategy."""
+    target = strategy_cfg.get("_target_", "")
+    return "AutomodelParallelStrategy" in target
+
+
+@hydra_runner(config_name="HfExportConfig", schema=HfExportConfig)
+def main(cfg: HfExportConfig) -> None:
+    """
+    Read PyTorch Lightning checkpoint and export the model to HuggingFace Hub format.
+    The resulting model can be then initialized via ModelClass.from_pretrained(path).
+
+    Also supports distributed checkpoints for models trained with FSDP2/TP
+    via AutomodelParallelStrategy.  Parallelism sizes (tp_size, pp_size, etc.)
+    are read automatically from the ``trainer.strategy`` section of the
+    experiment config (``ckpt_config``).
+
+    When the checkpoint is a distributed checkpoint (a directory), launch this
+    script via ``torchrun`` with the same number of GPUs used for training.
+
+    Examples:
+        # Single-file checkpoint — original SALM (HF Transformers backend):
+        python to_hf.py \\
+            class_path=nemo.collections.speechlm2.models.SALM \\
+            ckpt_path=/path/to/checkpoint.ckpt \\
+            ckpt_config=/path/to/config.yaml \\
+            output_dir=/path/to/hf_output
+
+        # Single-file checkpoint — SALMAutomodel (NeMo Automodel backend):
+        python to_hf.py \\
+            class_path=nemo.collections.speechlm2.models.SALMAutomodel \\
+            ckpt_path=/path/to/checkpoint.ckpt \\
+            ckpt_config=/path/to/config.yaml \\
+            output_dir=/path/to/hf_output
+
+        # Distributed checkpoint (parallelism read from config automatically):
+        torchrun --nproc-per-node=8 to_hf.py \\
+            class_path=nemo.collections.speechlm2.models.SALMAutomodel \\
+            ckpt_path=/path/to/distributed_ckpt_dir \\
+            ckpt_config=/path/to/config.yaml \\
+            output_dir=/path/to/hf_output
+    """
+    if not Path(cfg.ckpt_path).exists():
+        raise RuntimeError(f"No such file or directory: {cfg.ckpt_path}")
+
+    full_cfg = OmegaConf.to_container(OmegaConf.load(cfg.ckpt_config), resolve=True)
+    model_cfg = full_cfg["model"]
+    model_cfg["torch_dtype"] = _canonical_torch_dtype_name(cfg.dtype)
+    cls = import_class_by_path(cfg.class_path)
+
+    strategy_cfg = full_cfg.get("trainer", {}).get("strategy", {})
+
+    _is_torchrun = "RANK" in os.environ
+    if _is_torchrun and dist.is_available() and not dist.is_initialized():
+        dist.init_process_group(backend="nccl")
+    is_distributed = (
+        _is_torchrun
+        and Path(cfg.ckpt_path).is_dir()
+        and _uses_automodel_parallel(strategy_cfg)
+        and dist.get_world_size() > 1
+    )
+
+    if is_distributed:
+        strategy = setup_distributed_from_config(strategy_cfg)
+
+        # Don't call configure_model() inside __init__ — we set device_mesh first.
+        model_cfg["init_configure_model"] = False
+        model = cls(model_cfg)
+        model.configure_model(
+            device_mesh=strategy.device_mesh,
+            distributed_config=strategy.distributed_config,
+            moe_config=strategy.moe_config,
+            moe_mesh=strategy.moe_mesh,
+        )
+        model_cfg["pretrained_weights"] = False
+
+        load_checkpoint(model, cfg.ckpt_path)
+
+        # Consolidate DTensors to regular tensors and save on rank 0.
+        consolidated = consolidate_state_dict(model)
+        if dist.get_rank() == 0:
+            save_hf_checkpoint(model, consolidated, cfg)
+            _try_prepare_for_vllm(cfg.output_dir, model_cfg)
+
+        dist.barrier()
+        dist.destroy_process_group()
+    else:
+        model_cfg["init_configure_model"] = True
+        model = cls(model_cfg)
+        load_checkpoint(model, cfg.ckpt_path)
+        model = model.to(str_to_dtype(cfg.dtype))
+        model_cfg["pretrained_weights"] = False
+        model.save_pretrained(cfg.output_dir, config=_hf_export_config(model, cfg.dtype))
+        save_llm_backbone_config(model, cfg.output_dir)
+        _try_prepare_for_vllm(cfg.output_dir, model_cfg)
+
+
+if __name__ == "__main__":
+    main()