feat(speechlm2): add padding-free packed speech training

docs/source/speechlm2/configs.rst

@@ -162,14 +162,34 @@ Note the differences from the SALM configuration:
 * ``encoder_chunk_size_seconds`` controls long-audio chunking for the speech encoder.
   Audio rows longer than this value are split on the time axis, encoded as a chunk
   batch, and concatenated back into one embedding sequence before the LLM forward.
-  Set it to ``null`` to disable chunking.
+  Set it to ``null`` to disable chunking. With a ``ParallelExpertEncoder`` and
+  ``packed_encoder_sequences: true``, this same value instead chunks both the ASR
+  and diarization branches after feature stacking; set the data audio-token
+  estimator's ``chunk_size_seconds`` to ``null`` because this does not change the
+  total encoder-token count.
 
 SALMAutomodel-Specific Options
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 The SALMAutomodel config exposes a few extra knobs that pass through to NeMo
 Automodel. All are optional — defaults preserve standard behavior.
 
+**Garbage collection:**
+
+.. code-block:: yaml
+
+    model:
+      # Optional positive optimizer-step interval; null keeps automatic GC.
+      gc_every_steps: null
+
+Setting ``gc_every_steps`` to a positive integer disables Python's automatic
+garbage collector at fit start and uses NeMo Automodel's generation-1 collector
+at that optimizer-step cadence. This avoids an occasional generation-2 scan on
+one distributed rank delaying all peers at the next collective. The cadence is
+counted in optimizer steps, so gradient accumulation does not increase the
+collection frequency. Leave it ``null`` unless profiling shows GC-related rank
+stragglers.
+
 **MoE training:**
 
 .. code-block:: yaml

examples/speechlm2/conf/salm_automodel.yaml

@@ -14,6 +14,11 @@ model:
   # Set to true to use SALMAutomodel (NeMo Automodel backend) instead of SALM (HF Transformers backend).
   use_nemo_automodel: true
 
+  # Optional: disable asynchronous full-heap Python GC and run deterministic
+  # generation-1 collections every N optimizer steps. null preserves Python's
+  # default automatic GC behavior.
+  gc_every_steps: null
+
   # Regexp (re.compile) patterns matching parameters to be frozen.
   freeze_params:
     # Frozen LLM (embed_tokens stays inside llm, so this pattern covers it too)
@@ -50,6 +55,13 @@ model:
   # Set to null to disable encoder chunking and encode each audio row directly.
   encoder_chunk_size_seconds: 30.0
 
+  # Opt in to native token-major (THD) Transformer/MoE activations during training/validation.
+  # The training dataset also keeps raw waveforms packed through preprocessing when enabled.
+  # Existing checkpoints and generation behavior stay unchanged; unsupported adapters fail early.
+  packed_encoder_sequences: false
+  # Separately opt in to the flattened differentiable CP gather.
+  packed_encoder_cp: false
+
   # Uncomment the block below to enable LoRA on the LLM via Automodel.
   # LoRA parameters are kept trainable even when the LLM is frozen.
   # lora:
@@ -89,6 +101,7 @@ model:
   #   fullgraph: false       # Compile the full computation graph
   #   dynamic: true          # Enable dynamic shapes (recommended for variable-length audio)
   #   backend: null           # Compilation backend (null = inductor)
+  #   options: null           # Optional torch.compile backend options dictionary
   #   dynamo_cache_size_limit: 256  # Triton compilation cache limit
 
   # Automodel backend dispatch. Selects the kernel/backend for each major module
@@ -99,18 +112,24 @@ model:
   automodel_backend:
     dispatcher: torch           # Set to "deepep" only if your GPUs have NVLINK/NVSHMEM
   # Optional overrides for other backends (e.g. attn=sdpa to bypass TE):
-  #   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_num_sms: 20    # SM count for DeepEP/UCCL-EP kernels
+  #   attn: te                  # "te" | "sdpa" | "flex" | "eager" | "tilelang"
+  #   linear: te                # "torch" | "te" | "quack"
+  #   rms_norm: torch_fp32      # "torch" | "torch_fp32" | "te" | "quack"
+  #   rope: torch               # "torch" | "quack"
+  #   rope_fusion: false        # Fused RoPE (currently force-disabled in Automodel)
+  #   experts: torch_mm         # "torch" | "te" | "gmm" | "torch_mm" | "torch_mm_mxfp8"
+  #   dispatcher_num_sms: 32    # SM count for DeepEP/UCCL-EP kernels
+  #   dispatcher_share_token_dispatcher: true  # Share one flex dispatcher across MoE layers
+  #   dispatcher_async_dispatch: false         # Return DeepEP/UCCL-EP dispatch asynchronously
   #   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
+  #   compile_attn: false       # Compile attention; requires sdpa/torch/torch and no fused RoPE
+  #   cuda_graph:               # Scoped partial CUDA graphs; empty modules disable them
+  #     modules: []             # attn | te_dpa | moe_router | moe_preprocess
+  #   te_fp8: null              # {recipe: "current" | "block" | "mxfp8"} enables TE FP8
   #                             # (requires linear=te or experts=te)
 
   # Pin the SDPA kernel list used when automodel_backend.attn=sdpa. Accepts
@@ -196,7 +215,19 @@ trainer:
 
     # --- FSDP2 distributed config (plain dict, resolved to FSDP2Config automatically) ---
     # distributed_config:
-    #   sequence_parallel: false    # Enable sequence parallelism (requires tp_size > 1)
+    #   sequence_parallel: false       # Enable sequence parallelism (requires tp_size > 1)
+    #   tp_plan: null                  # Optional custom Transformers TP plan
+    #   patch_is_packed_sequence: false # True is safe only for non-packed training
+    #   mp_policy: null                # null uses Automodel's bf16-compute/fp32-reduce policy
+    #   autocast_dtype: null           # e.g. bfloat16; null disables explicit autocast
+    #   activation_checkpointing_scope: all # all | language | vision | audio | multimodal
+    #   defer_fsdp_grad_sync: true      # Sync gradients only on the final microbatch
+    #   reshard_after_forward: null     # null keeps Automodel's FSDP2 heuristic
+    #   enable_async_tensor_parallel: false
+    #   enable_compile: false           # Per-transformer-layer torch.compile
+    #   enable_fsdp2_prefetch: false
+    #   fsdp2_backward_prefetch_depth: 2
+    #   fsdp2_forward_prefetch_depth: 1
     #   # offload_policy:           # Uncomment to enable CPU offloading
     #   #   _target_: torch.distributed.fsdp.CPUOffloadPolicy
 
@@ -206,12 +237,28 @@ trainer:
     #   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
+    #   mp_policy: null                     # Optional MixedPrecisionPolicy override
 
 data:
   train_ds:
     sample_rate: 16000
     prompt_format: ${model.prompt_format}
     token_equivalent_duration: 0.08
+    # Sample-exact Canary-v2 audio frame count, including per-chunk rounding.
+    # Required when use_packed_sequence_sampling enforces a hard model-token cap.
+    audio_token_estimator: &canary_v2_audio_token_estimator
+      preprocessor:
+        n_fft: 512
+        hop_length: 160
+        stft_pad_amount: 256
+      subsampling:
+        type: conv
+        kernel_size: 3
+        stride: 2
+        padding: 1
+        repeat: 3
+        ceil_mode: false
+      chunk_size_seconds: ${model.encoder_chunk_size_seconds}
     input_cfg:
       - type: lhotse_as_conversation
         cuts_path: ???  # needs to be set
@@ -229,10 +276,13 @@ data:
     # batch_size: null
     # use_bucketing: true
     # use_multimodal_sampling: true
+    # use_packed_sequence_sampling: false  # Sum lengths; requires a padding-free model path
+    # packing_buffer_size: 128  # Packed best-fit lookahead, with or without bucketing
+    # shuffle_buffer_size: 10000  # Reservoir size for ordinary non-packed samplers; legacy packed configs are aliased
     # 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`.
+    #  the sum of input audio frames and output text tokens. `audio_token_estimator` gives
+    #  exact frame counts; `token_equivalent_duration` remains the legacy fallback.
     # batch_tokens: 4000
     # max_tokens: 2048
     # bucket_duration_bins: [64, 128, 256, 384, 512, 768, 1024, 1280, 1536, 2048]
@@ -245,6 +295,7 @@ data:
     # They inherit all settings from validation_ds, but can individually override them.
     prompt_format: ${model.prompt_format}
     token_equivalent_duration: 0.08
+    audio_token_estimator: *canary_v2_audio_token_estimator
     datasets:
       val_set_0:  # rename to your dataset name, add more as needed
         input_cfg:

examples/speechlm2/conf/salm_automodel_pee.yaml

@@ -14,6 +14,11 @@ model:
   # Set to true to use SALMAutomodel (NeMo Automodel backend) instead of SALM (HF Transformers backend).
   use_nemo_automodel: true
 
+  # Optional: disable asynchronous full-heap Python GC and run deterministic
+  # generation-1 collections every N optimizer steps. null preserves Python's
+  # default automatic GC behavior.
+  gc_every_steps: null
+
   # 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
@@ -35,6 +40,13 @@ model:
   # Set to null to disable encoder chunking and encode each audio row directly.
   encoder_chunk_size_seconds: 60.0
 
+  # Opt in to native token-major (THD) Transformer/MoE/PEE activations during training/validation.
+  # The training dataset also keeps raw waveforms packed through preprocessing when enabled.
+  # Default false preserves historical behavior; generation always keeps its streaming-safe path.
+  packed_encoder_sequences: false
+  # Separately opt in to the flattened differentiable CP gather.
+  packed_encoder_cp: false
+
   # ─── 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
@@ -103,26 +115,33 @@ model:
   #   fullgraph: false       # Compile the full computation graph
   #   dynamic: true          # Enable dynamic shapes (recommended for variable-length audio)
   #   backend: null           # Compilation backend (null = inductor)
+  #   options: null           # Optional torch.compile backend options dictionary
   #   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). BackendConfig
   # selects installed kernels by default; 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"
+  #   attn: te                  # "te" | "sdpa" | "flex" | "eager" | "tilelang"
+  #   linear: te                # "torch" | "te" | "quack"
+  #   rms_norm: torch_fp32      # "torch" | "torch_fp32" | "te" | "quack"
+  #   rope: torch               # "torch" | "quack"
+  #   rope_fusion: false        # Fused RoPE (currently force-disabled in Automodel)
+  #   experts: torch_mm         # "torch" | "te" | "gmm" | "torch_mm" | "torch_mm_mxfp8"
   #   dispatcher: deepep        # MoE token dispatcher: "torch" | "deepep" | "hybridep" | "uccl_ep"
-  #   dispatcher_num_sms: 20    # SM count for DeepEP/UCCL-EP kernels
+  #   dispatcher_num_sms: 32    # SM count for DeepEP/UCCL-EP kernels
+  #   dispatcher_share_token_dispatcher: true  # Share one flex dispatcher across MoE layers
+  #   dispatcher_async_dispatch: false         # Return DeepEP/UCCL-EP dispatch asynchronously
   #   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
+  #   compile_attn: false       # Compile attention; requires sdpa/torch/torch and no fused RoPE
+  #   cuda_graph:               # Scoped partial CUDA graphs; empty modules disable them
+  #     modules: []             # attn | te_dpa | moe_router | moe_preprocess
+  #   te_fp8: null              # {recipe: "current" | "block" | "mxfp8"} enables TE FP8
   #                             # (requires linear=te or experts=te)
 
   # Pin the SDPA kernel list used when automodel_backend.attn=sdpa. Accepts
@@ -199,7 +218,19 @@ trainer:
 
     # --- FSDP2 distributed config (plain dict, resolved to FSDP2Config automatically) ---
     # distributed_config:
-    #   sequence_parallel: false    # Enable sequence parallelism (requires tp_size > 1)
+    #   sequence_parallel: false       # Enable sequence parallelism (requires tp_size > 1)
+    #   tp_plan: null                  # Optional custom Transformers TP plan
+    #   patch_is_packed_sequence: false # True is safe only for non-packed training
+    #   mp_policy: null                # null uses Automodel's bf16-compute/fp32-reduce policy
+    #   autocast_dtype: null           # e.g. bfloat16; null disables explicit autocast
+    #   activation_checkpointing_scope: all # all | language | vision | audio | multimodal
+    #   defer_fsdp_grad_sync: true      # Sync gradients only on the final microbatch
+    #   reshard_after_forward: null     # null keeps Automodel's FSDP2 heuristic
+    #   enable_async_tensor_parallel: false
+    #   enable_compile: false           # Per-transformer-layer torch.compile
+    #   enable_fsdp2_prefetch: false
+    #   fsdp2_backward_prefetch_depth: 2
+    #   fsdp2_forward_prefetch_depth: 1
     #   # offload_policy:           # Uncomment to enable CPU offloading
     #   #   _target_: torch.distributed.fsdp.CPUOffloadPolicy
 
@@ -209,6 +240,7 @@ trainer:
     #   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
+    #   mp_policy: null                     # Optional MixedPrecisionPolicy override
 
 data:
   # RTTM/SOT speaker-activity targets for ParallelExpertEncoder training. Active for
@@ -224,6 +256,19 @@ data:
     sample_rate: 16000
     prompt_format: ${model.prompt_format}
     token_equivalent_duration: 0.08
+    # Sample-exact PEE audio frame count, including per-chunk rounding.
+    # Required when use_packed_sequence_sampling enforces a hard model-token cap.
+    audio_token_estimator: &pee_audio_token_estimator
+      preprocessor:
+        n_fft: 512
+        hop_length: 160
+        stft_pad_amount: 256
+      subsampling:
+        type: feature_stacking
+        factor: 8
+      # Dense PEE uses the outer waveform chunker. Set this to null when enabling
+      # packed_encoder_sequences, where PEE chunks internally after feature stacking.
+      chunk_size_seconds: ${model.encoder_chunk_size_seconds}
     input_cfg:
       - type: lhotse_as_conversation
         cuts_path: ???  # needs to be set
@@ -241,10 +286,13 @@ data:
     # batch_size: null
     # use_bucketing: true
     # use_multimodal_sampling: true
+    # use_packed_sequence_sampling: false  # Sum lengths; requires a padding-free model path
+    # packing_buffer_size: 128  # Packed best-fit lookahead, with or without bucketing
+    # shuffle_buffer_size: 10000  # Reservoir size for ordinary non-packed samplers; legacy packed configs are aliased
     # 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`.
+    #  the sum of input audio frames and output text tokens. `audio_token_estimator` gives
+    #  exact frame counts; `token_equivalent_duration` remains the legacy fallback.
     # batch_tokens: 4000
     # max_tokens: 2048
     # bucket_duration_bins: [64, 128, 256, 384, 512, 768, 1024, 1280, 1536, 2048]
@@ -257,6 +305,7 @@ data:
     # They inherit all settings from validation_ds, but can individually override them.
     prompt_format: ${model.prompt_format}
     token_equivalent_duration: 0.08
+    audio_token_estimator: *pee_audio_token_estimator
     datasets:
       val_set_0:  # rename to your dataset name, add more as needed
         input_cfg:

examples/speechlm2/salm_train.py

@@ -48,13 +48,29 @@ def _process_group_timeout(cfg):
     return None
 
 
-def _create_salm_dataset(tokenizer, data_cfg: DictConfig | dict) -> SALMDataset:
+def _create_salm_dataset(
+    tokenizer,
+    data_cfg: DictConfig | dict,
+    *,
+    pack_audio: bool = False,
+    pack_sequences: bool = False,
+) -> SALMDataset:
     """Build SALMDataset without forwarding unset options to legacy NeMo packages."""
     multispeaker_cfg = data_cfg.get("multispeaker_cfg", None)
+    batch_tokens = data_cfg.get("train_ds", {}).get("batch_tokens", None)
     # TODO(Dongji): Remove after all release images ship SALMDataset with multispeaker_cfg support.
-    if multispeaker_cfg is None:
+    if multispeaker_cfg is None and not pack_audio and not pack_sequences and batch_tokens is None:
         return SALMDataset(tokenizer=tokenizer)
-    return SALMDataset(tokenizer=tokenizer, multispeaker_cfg=multispeaker_cfg)
+    kwargs = {"tokenizer": tokenizer}
+    if multispeaker_cfg is not None:
+        kwargs["multispeaker_cfg"] = multispeaker_cfg
+    if pack_audio:
+        kwargs["pack_audio"] = True
+    if pack_sequences:
+        kwargs["pack_sequences"] = True
+    if batch_tokens is not None:
+        kwargs["batch_tokens"] = batch_tokens
+    return SALMDataset(**kwargs)
 
 
 @hydra_runner(config_path="conf", config_name="salm")
@@ -85,7 +101,14 @@ def train(cfg):
     with trainer.init_module():
         model = model_cls(OmegaConf.to_container(cfg.model, resolve=True))
 
-    dataset = _create_salm_dataset(model.tokenizer, cfg.data)
+    dataset = _create_salm_dataset(
+        model.tokenizer,
+        cfg.data,
+        pack_audio=bool(
+            cfg.model.get("use_nemo_automodel", False) and cfg.model.get("packed_encoder_sequences", False)
+        ),
+        pack_sequences=bool(cfg.model.get("use_nemo_automodel", False) and cfg.model.get("packed_sequences", False)),
+    )
     datamodule = DataModule(cfg.data, tokenizer=model.tokenizer, dataset=dataset)
 
     if cfg.get("run_validate_only", False):