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