diff --git a/docs/source/speechlm2/configs.rst b/docs/source/speechlm2/configs.rst index 1b5833537160..42f2aae6c16f 100644 --- a/docs/source/speechlm2/configs.rst +++ b/docs/source/speechlm2/configs.rst @@ -162,7 +162,11 @@ 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 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ @@ -170,6 +174,22 @@ 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 diff --git a/examples/speechlm2/conf/salm_automodel.yaml b/examples/speechlm2/conf/salm_automodel.yaml index 57377510f81c..b4010abe8e16 100644 --- a/examples/speechlm2/conf/salm_automodel.yaml +++ b/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: diff --git a/examples/speechlm2/conf/salm_automodel_pee.yaml b/examples/speechlm2/conf/salm_automodel_pee.yaml index 18dfe30732f0..35600cdc6e54 100644 --- a/examples/speechlm2/conf/salm_automodel_pee.yaml +++ b/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 -tagged @@ -103,6 +115,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 @@ -110,19 +123,25 @@ model: # 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: diff --git a/examples/speechlm2/salm_train.py b/examples/speechlm2/salm_train.py index ebecdc7d596f..1df9f1cc9678 100644 --- a/examples/speechlm2/salm_train.py +++ b/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): diff --git a/nemo/collections/asr/modules/audio_preprocessing.py b/nemo/collections/asr/modules/audio_preprocessing.py index 148a3d33cc9b..6c98470dca2a 100644 --- a/nemo/collections/asr/modules/audio_preprocessing.py +++ b/nemo/collections/asr/modules/audio_preprocessing.py @@ -21,6 +21,7 @@ import torch +from nemo.collections.asr.parts.packed_sequence import PackedEncoderActivations from nemo.collections.asr.parts.preprocessing.features import FilterbankFeatures from nemo.collections.asr.parts.submodules.spectr_augment import SpecAugment, SpecCutout from nemo.collections.audio.parts.utils.transforms import MFCC @@ -92,11 +93,39 @@ def forward(self, input_signal, length): processed_signal = processed_signal.to(self.dtype_sentinel_tensor.dtype) return processed_signal, processed_length + @torch.no_grad() + def forward_packed(self, input_signal, length, input_signal_cu_seqlens) -> PackedEncoderActivations: + """Preprocess concatenated waveforms into token-flat features. + + The packed frontend uses one vectorized guarded STFT and returns exactly + the valid feature frames. The historical padded :meth:`forward` contract + and checkpoint state remain unchanged. + + Args: + input_signal: Concatenated waveform samples with shape `(sum(length),)`. + length: Per-waveform sample counts with shape `(B,)`. + input_signal_cu_seqlens: Cumulative sample offsets with shape `(B + 1,)`. + + Returns: + Packed time-major features and their sequence metadata. + """ + if input_signal.dtype != torch.float32: + logging.warning( + f"AudioPreprocessor received an input signal of dtype {input_signal.dtype}, rather than " + "torch.float32. Packed preprocessing runs in float32 for numerical stability.", + mode=logging_mode.ONCE, + ) + processed = self.get_features_packed(input_signal.to(torch.float32), length, input_signal_cu_seqlens) + return processed.with_data(processed.data.to(self.dtype_sentinel_tensor.dtype)) + @abstractmethod def get_features(self, input_signal, length): # Called by forward(). Subclasses should implement this. pass + def get_features_packed(self, input_signal, length, input_signal_cu_seqlens) -> PackedEncoderActivations: + raise NotImplementedError(f"{type(self).__name__} does not implement packed waveform preprocessing.") + class AudioToMelSpectrogramPreprocessor(AudioPreprocessor, Exportable): """Featurizer module that converts wavs to mel spectrograms. @@ -284,6 +313,9 @@ def input_example(self, max_batch: int = 8, max_dim: int = 32000, min_length: in def get_features(self, input_signal, length): return self.featurizer(input_signal, length) + def get_features_packed(self, input_signal, length, input_signal_cu_seqlens) -> PackedEncoderActivations: + return self.featurizer.forward_packed(input_signal, length, input_signal_cu_seqlens) + @property def filter_banks(self): return self.featurizer.filter_banks @@ -528,6 +560,15 @@ def forward(self, input_spec, length): augmented_spec = self.spec_augment(input_spec=augmented_spec, length=length) return augmented_spec + @torch.no_grad() + def forward_packed(self, input_spec: PackedEncoderActivations) -> PackedEncoderActivations: + """Apply sequence-local augmentation without padding token-flat features.""" + if isinstance(self.spec_cutout, SpecCutout): + input_spec = self.spec_cutout.forward_packed(input_spec) + if isinstance(self.spec_augment, SpecAugment): + input_spec = self.spec_augment.forward_packed(input_spec) + return input_spec + class MaskedPatchAugmentation(NeuralModule): """ diff --git a/nemo/collections/asr/modules/parallel_expert_encoder.py b/nemo/collections/asr/modules/parallel_expert_encoder.py index 0f86210dc5a3..4f69ae78e9bb 100644 --- a/nemo/collections/asr/modules/parallel_expert_encoder.py +++ b/nemo/collections/asr/modules/parallel_expert_encoder.py @@ -19,7 +19,8 @@ native Transformer encoder on the same mel input, then fuses their outputs with a sinusoidal speaker kernel. The encoder expects unnormalized mels; the ASR and Sortformer branches independently reapply ``normalize_batch`` internally. I/O -matches :class:`ConformerEncoder`. +matches :class:`ConformerEncoder`, including a compatibility fallback for +packed SALM execution. Only self-contained bundles with inline ``asr_encoder_cfg`` and ``diarization_model_cfg`` sections are supported. @@ -44,7 +45,12 @@ from nemo.collections.asr.modules.conformer_encoder import ConformerEncoder from nemo.collections.asr.modules.transformer_encoder import TransformerEncoder -from nemo.collections.asr.parts.preprocessing.features import normalize_batch +from nemo.collections.asr.parts.packed_sequence import ( + PackedEncoderActivations, + pack_encoder_output, + unpack_encoder_output, +) +from nemo.collections.asr.parts.preprocessing.features import normalize_batch, normalize_packed_batch from nemo.core.classes import ModelPT from nemo.core.classes.common import PretrainedModelInfo, Serialization from nemo.core.classes.module import freeze, unfreeze @@ -67,7 +73,9 @@ _BUNDLE_CONFIG_OVERRIDE_KEYS = frozenset( { "asr_normalize_type", + "chunk_size_seconds", "diar_normalize_type", + "frame_shift_seconds", "missing_rttm_target", "speaker_activity_threshold", "speaker_feature_config_version", @@ -266,6 +274,8 @@ def __init__(self, cfg: DictConfig, trainer: Optional[Trainer] = None): speaker_feature_mode=speaker_feature_mode, speaker_activity_threshold=speaker_activity_threshold, spk_kernel_scale=self._cfg.get("spk_kernel_scale", 1.0), + frame_shift_seconds=self._cfg.get("frame_shift_seconds", 0.01), + chunk_size_seconds=self._cfg.get("chunk_size_seconds", None), sync_max_audio_length=self._cfg.get("sync_max_audio_length", False), ) @@ -304,7 +314,7 @@ def is_pe_nemo(cls, nemo_path: str) -> bool: if not str(cfg.get("target", "")).endswith("ParallelExpertEncoderPT"): return False # Keep the released public probe target-based. Runtime loading - # uses the schema resolver and remains strict. + # validates the canonical bundle schema and remains strict. return True except (tarfile.TarError, OSError) as error: logging.warning("[ParallelExpertEncoder] Could not inspect %s: %s", nemo_path, error) @@ -434,6 +444,8 @@ class ParallelExpertEncoder(nn.Module): :class:`TransformerEncoder` used by Transformer AED ASR checkpoints. """ + supports_sequence_packed_output = True + def __init__( self, asr_encoder_cfg: DictConfig, @@ -453,6 +465,8 @@ def __init__( speaker_feature_mode: Optional[str] = None, speaker_activity_threshold: Optional[float] = None, spk_kernel_scale: float = 1.0, + frame_shift_seconds: float = 0.01, + chunk_size_seconds: Optional[float] = None, sync_max_audio_length: bool = False, ): super().__init__() @@ -499,7 +513,7 @@ def __init__( ) # The ASR and diarization experts are called from data-dependent paths - # in both training and replicated inference. Their positional + # in both packed training and replicated inference. Their positional # buffers are local state, so synchronizing the longest feature length # on the default process group is unnecessary and can deadlock when # ranks process different request shapes. @@ -509,6 +523,11 @@ def __init__( self.freeze_diar = bool(freeze_diar) self.freeze_asr = bool(freeze_asr) + self.frame_shift_seconds = float(frame_shift_seconds) + if self.frame_shift_seconds <= 0: + raise ValueError(f"frame_shift_seconds must be positive, got {frame_shift_seconds}.") + self.chunk_size_seconds = self._validate_chunk_size("chunk_size_seconds", chunk_size_seconds) + self.online_inference_length = int(online_inference_length) self.online_inference_enabled: Optional[bool] = None self.chunk_left_context = max(0, int(chunk_left_context)) @@ -570,8 +589,8 @@ def set_activation_checkpointing(self, enabled: bool) -> None: """Wrap trainable ASR stages before FSDP2 sharding. The frozen Sortformer branch is deliberately excluded. Per-layer wrappers - preserve FSDP2 boundaries, unlike a checkpoint around the entire encoder - call. + preserve FSDP2 boundaries and native packed-layer dispatch, unlike a + checkpoint around the entire encoder call. """ if not enabled or self.freeze_asr: return @@ -591,6 +610,91 @@ def set_activation_checkpointing(self, enabled: bool) -> None: if getattr(layer, "_checkpoint_wrapped_module", None) is None: layers[index] = checkpoint_wrapper(layer) + @staticmethod + def _validate_chunk_size(name: str, value: Optional[float]) -> Optional[float]: + if value is None: + return None + value = float(value) + if value <= 0: + raise ValueError(f"{name} must be positive or None, got {value}.") + return value + + def _chunk_size_tokens(self, chunk_size_seconds: Optional[float]) -> Optional[int]: + if chunk_size_seconds is None: + return None + token_seconds = self.frame_shift_seconds * self.subsampling_factor + return max(1, round(chunk_size_seconds / token_seconds)) + + @staticmethod + def _chunk_metadata(packed: PackedEncoderActivations, max_tokens: int) -> PackedEncoderActivations: + chunk_lengths = [] + for sequence_length in packed.lengths.detach().cpu().tolist(): + chunk_lengths.extend([max_tokens] * (sequence_length // max_tokens)) + if sequence_length % max_tokens: + chunk_lengths.append(sequence_length % max_tokens) + lengths = torch.as_tensor(chunk_lengths, dtype=torch.int64, device=packed.data.device) + cu_seqlens = torch.cat( + [ + torch.zeros(1, dtype=torch.int32, device=packed.data.device), + lengths.cumsum(0, dtype=torch.int32), + ] + ).contiguous() + return PackedEncoderActivations( + data=packed.data, + lengths=lengths, + cu_seqlens=cu_seqlens, + max_seqlen=min(max_tokens, packed.max_seqlen), + padding_value=packed.padding_value, + padded_length=None, + ) + + @staticmethod + def _match_packed_module_io(packed: PackedEncoderActivations, module: nn.Module) -> PackedEncoderActivations: + parameter = next(module.parameters(), None) + if parameter is None: + return packed + if packed.data.device != parameter.device: + raise ValueError( + f"Packed input is on {packed.data.device}, but {type(module).__name__} is on {parameter.device}." + ) + if packed.data.dtype == parameter.dtype: + return packed + return packed.with_data(packed.data.to(dtype=parameter.dtype)) + + def _forward_packed_branch( + self, + encoder: nn.Module, + features: PackedEncoderActivations, + chunk_size_seconds: Optional[float], + ) -> PackedEncoderActivations: + """Run an encoder token-flat, optionally splitting after feature stacking.""" + max_tokens = self._chunk_size_tokens(chunk_size_seconds) + packed_forward = getattr(encoder, "forward_sequence_packed", None) + if not callable(packed_forward): + if max_tokens is not None and features.max_seqlen > max_tokens: + raise TypeError(f"{type(encoder).__name__} does not support packed independent chunking.") + padded = unpack_encoder_output(features, total_length=features.padded_length).transpose(1, 2) + encoded, encoded_lengths = encoder(audio_signal=padded, length=features.lengths) + return pack_encoder_output(encoded.transpose(1, 2), encoded_lengths) + if max_tokens is None or features.max_seqlen <= max_tokens: + return packed_forward(features, features.lengths) + + pre_encode = getattr(encoder, "pre_encode", None) + unwrapped_pre_encode = getattr(pre_encode, "_checkpoint_wrapped_module", pre_encode) + if type(unwrapped_pre_encode).__name__ != "FeatureStacking": + raise TypeError( + "Independent post-stacking chunking requires subsampling='feature_stacking'; " + f"got {type(unwrapped_pre_encode).__name__} for {type(encoder).__name__}." + ) + pre_encoded = pre_encode(features) + chunked = self._chunk_metadata(pre_encoded, max_tokens) + encoded_chunks = packed_forward( + chunked, + chunked.lengths, + bypass_pre_encode=True, + ) + return pre_encoded.with_data(encoded_chunks.data) + def _asr_output_frame_boundary(self, input_frame_boundary: int) -> int: """Map an input-frame boundary to the selected ASR encoder's output grid.""" if getattr(self, "asr_encoder_type", "fastconformer") == "transformer": @@ -740,6 +844,38 @@ def forward(self, audio_signal, length, spk_targets=None): runner = self._forward_online if use_online else self._forward return runner(audio_signal=audio_signal, length=length, spk_targets=spk_targets) + def forward_sequence_packed(self, audio_signal, length=None, spk_targets=None) -> PackedEncoderActivations: + """Run Sortformer first and ASR second while keeping encoder states token-flat.""" + if bool(getattr(self, "online_inference_enabled", False)): + raise RuntimeError("forward_sequence_packed is an offline API and cannot run inside online_inference().") + if isinstance(audio_signal, PackedEncoderActivations): + if length is not None and not torch.equal(length.to(audio_signal.lengths), audio_signal.lengths): + raise ValueError("length must match audio_signal.lengths for packed input.") + features = audio_signal + else: + if length is None: + raise ValueError("length is required for padded input.") + features = pack_encoder_output(audio_signal.transpose(1, 2), length) + + self._check_spk_targets(spk_targets, features.batch_size) + needs_diarization = self._should_run_diarization(spk_targets) + diarization_preds = self._run_diarization_packed(features) if needs_diarization else None + asr_encoded = self._run_asr_packed(features) + if diarization_preds is not None and not ( + torch.equal(diarization_preds.lengths, asr_encoded.lengths) + and torch.equal(diarization_preds.cu_seqlens, asr_encoded.cu_seqlens) + ): + raise RuntimeError( + "Sortformer and ASR output metadata diverged: " + f"diar={diarization_preds.lengths.detach().cpu().tolist()} " + f"asr={asr_encoded.lengths.detach().cpu().tolist()}." + ) + return self._fuse_diar_and_asr_packed( + asr_encoded, + spk_targets if spk_targets is not None else diarization_preds, + diarization_preds=diarization_preds, + ) + def _align_diarization_output_resolution( self, predictions: torch.Tensor, embedding_lengths: torch.Tensor ) -> torch.Tensor: @@ -771,6 +907,81 @@ def _downsample_high_resolution_diarization_for_fusion( ) return predictions + def _run_diarization_packed(self, features: PackedEncoderActivations) -> PackedEncoderActivations: + """Normalize each utterance, then run the frozen streaming-trained Sortformer.""" + if self.diar_normalize_type: + features = normalize_packed_batch(features, self.diar_normalize_type) + features = self._match_packed_module_io(features, self.diarization_model.encoder) + with torch.set_grad_enabled(torch.is_grad_enabled() and not self.freeze_diar): + embeddings = self._forward_packed_branch( + self.diarization_model.encoder, + features, + self.chunk_size_seconds, + ) + modules = self.diarization_model.sortformer_modules + projected = embeddings.data + if modules.encoder_proj is not None: + projected = modules.encoder_proj(projected) + + post_encoder = self.diarization_model.transformer_encoder + has_post_layers = post_encoder is not None and len(post_encoder.layers) > 0 + if has_post_layers or self.diarization_model.high_resolution: + padded = unpack_encoder_output(embeddings) + if modules.encoder_proj is not None: + padded = modules.encoder_proj(padded) + predictions = self.diarization_model.forward_infer(padded, embeddings.lengths) + predictions = self._align_diarization_output_resolution(predictions, embeddings.lengths) + return pack_encoder_output(predictions, embeddings.lengths) + if post_encoder is not None and post_encoder.final_layer_norm is not None: + projected = post_encoder.final_layer_norm(projected) + predictions = modules.forward_speaker_sigmoids(projected) + return embeddings.with_data(predictions) + + def _run_asr_packed(self, features: PackedEncoderActivations) -> PackedEncoderActivations: + """Normalize once per utterance, then run the trainable ASR packed path.""" + if self.asr_normalize_type: + features = normalize_packed_batch(features, self.asr_normalize_type) + features = self._match_packed_module_io(features, self.asr_encoder) + with torch.set_grad_enabled(torch.is_grad_enabled() and not self.freeze_asr): + return self._forward_packed_branch( + self.asr_encoder, + features, + self.chunk_size_seconds, + ) + + def _fuse_diar_and_asr_packed( + self, + asr_encoded: PackedEncoderActivations, + spk_targets: Union[torch.Tensor, PackedEncoderActivations], + *, + diarization_preds: Optional[PackedEncoderActivations] = None, + ) -> PackedEncoderActivations: + if isinstance(spk_targets, PackedEncoderActivations): + packed_targets = spk_targets + else: + use_diarization = self._missing_target_rows(spk_targets) + targets = self._align_diar_frames(spk_targets, asr_encoded.max_seqlen).to( + device=asr_encoded.data.device, dtype=asr_encoded.data.dtype + ) + if bool(use_diarization.any().item()): + if diarization_preds is None: + raise ValueError("diarization_preds are required for missing speaker-target rows.") + padded_preds = unpack_encoder_output(diarization_preds) + targets = torch.where( + use_diarization.to(device=targets.device, dtype=torch.bool).view(-1, 1, 1), + padded_preds.to(device=targets.device, dtype=targets.dtype), + targets, + ) + packed_targets = pack_encoder_output(targets, asr_encoded.lengths) + + if not torch.equal(packed_targets.lengths, asr_encoded.lengths): + raise RuntimeError("Packed speaker attributions must match ASR output lengths.") + speaker_features = self._speaker_features(packed_targets.data, asr_encoded.data.dtype) + normalized_states = self.asr_norm(asr_encoded.data) + normalized_targets = self.diar_norm(speaker_features) + infusion = torch.matmul(normalized_targets, self.diar_kernel.to(normalized_targets.dtype)) + return asr_encoded.with_data(normalized_states + getattr(self, "spk_kernel_scale", 1.0) * infusion) + def _run_diarization(self, audio_signal: torch.Tensor, length: torch.Tensor) -> torch.Tensor: if self.diar_normalize_type: audio_signal, _, _ = normalize_batch(audio_signal, length, normalize_type=self.diar_normalize_type) diff --git a/nemo/collections/asr/modules/transformer_encoder.py b/nemo/collections/asr/modules/transformer_encoder.py index bc24f987c484..affc4513621b 100644 --- a/nemo/collections/asr/modules/transformer_encoder.py +++ b/nemo/collections/asr/modules/transformer_encoder.py @@ -16,14 +16,21 @@ import math import random from dataclasses import dataclass -from typing import Optional +from typing import Optional, Sequence import torch import torch.nn as nn -from torch.nn.attention.flex_attention import and_masks, create_block_mask, flex_attention +import torch.nn.functional as F +from torch.nn.attention.flex_attention import and_masks, create_block_mask from nemo.collections.asr.models.configs import CacheAwareStreamingConfig +from nemo.collections.asr.modules import transformer_encoder_utils as _transformer_utils from nemo.collections.asr.parts.mixins.streaming import StreamingEncoder +from nemo.collections.asr.parts.packed_sequence import ( + PackedEncoderActivations, + pack_encoder_output, + packed_encoder_position_ids, +) from nemo.collections.asr.parts.submodules.multi_head_attention import ( PositionalEncoding, RelPositionalEncoding, @@ -35,8 +42,6 @@ from nemo.utils import logging from nemo.utils.decorators import experimental -flex_attention_compiled = torch.compile(flex_attention, dynamic=True) - @dataclass class TransformerEncoderConfig: @@ -218,6 +223,9 @@ def __init__(self, cfg: TransformerEncoderConfig, pos_enc=None): self.self_attention_model = cfg.self_attention_model self._uses_rel_pos = self.self_attention_model == "rel_pos" self._uses_rope = self.self_attention_model == "rope" + self._flash_attention_varlen_static_eligible = ( + not self._uses_rel_pos and self.head_dim <= 256 and self.head_dim % 8 == 0 + ) if self.self_attention_model not in _SUPPORTED_SELF_ATTENTION_MODELS: raise ValueError( f"self_attention_model='{self.self_attention_model}' is not supported. " @@ -313,6 +321,11 @@ def _build_rel_pos_score_mod(self, q, pos_emb): def score_mod(score, b, h, q_idx, kv_idx): return score + rel_pos_bias[b, h, q_idx, kv_idx] + # The compact CPU-with-grad compatibility path consumes the same tensor + # directly because older supported PyTorch releases cannot differentiate + # FlexAttention on CPU. + score_mod._relative_position_bias = rel_pos_bias + # Matrix c: fold u @ K^T into FlexAttention by rewriting Q as (Q + u). return score_mod, q + bias_u @@ -335,7 +348,7 @@ def forward(self, x, block_mask=None, pos_emb=None): if self._uses_rel_pos: score_mod, q = self._build_rel_pos_score_mod(q, pos_emb) - attn_fn = flex_attention_compiled if q.is_cuda else flex_attention + attn_fn = _transformer_utils._get_flex_attention(q) out = attn_fn(q, k, v, block_mask=block_mask, score_mod=score_mod) out = out.transpose(1, 2).contiguous().view(B, T, self.d_model) return self.out_proj(out) @@ -458,11 +471,176 @@ def forward_streaming(self, kv_normed, num_cur, block_mask, pos_emb): if k_mod is not None: k = k_mod - attn_fn = flex_attention_compiled if q.is_cuda else flex_attention + attn_fn = _transformer_utils._get_flex_attention(q) out = attn_fn(q, k, v, block_mask=block_mask, score_mod=score_mod) out = out.transpose(1, 2).contiguous().view(B, num_cur, self.d_model) return self.out_proj(out) + def forward_sequence_packed( + self, + x: torch.Tensor, + *, + lengths: torch.Tensor, + cu_seqlens: torch.Tensor, + max_seqlen: int, + position_ids: Optional[torch.Tensor] = None, + pos_emb: Optional[torch.Tensor] = None, + padded_length: Optional[int] = None, + causal: bool = False, + sequence_offsets: Optional[Sequence[int]] = None, + fused_qkv: bool = False, + ) -> torch.Tensor: + """Run self-attention on token-flat encoder states. + + CUDA fp16/bf16 RoPE, absolute-position, and no-position inputs use + FlashAttention's native variable-length THD kernel when available. Other + inputs use a compact per-utterance FlexAttention reference without + recreating a padded batch-wide activation. + + Args: + x: Token-flat encoder states with shape ``(total_tokens, d_model)``. + lengths: Number of tokens in each sequence. + cu_seqlens: Exclusive cumulative sequence lengths with shape ``(batch_size + 1,)``. + max_seqlen: Maximum sequence length in the packed batch. + position_ids: Optional token-flat positions required by RoPE attention. + pos_emb: Optional relative-position embeddings shared by the packed sequences. + padded_length: Padded source width used to slice relative-position embeddings. + causal: Whether to prevent queries from attending to future keys. + sequence_offsets: Optional host-side sequence boundaries for the reference backend. + fused_qkv: Whether to project Q/K/V with one fused linear operation. + + Returns: + Token-flat attention output with shape ``(total_tokens, d_model)``. + """ + return self._forward_sequence_packed( + x, + lengths=lengths, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + position_ids=position_ids, + pos_emb=pos_emb, + padded_length=padded_length, + causal=causal, + sequence_offsets=sequence_offsets, + fused_qkv=fused_qkv, + ) + + def _forward_sequence_packed( + self, + x, + *, + lengths, + cu_seqlens, + max_seqlen, + position_ids, + pos_emb, + padded_length, + causal, + sequence_offsets, + fused_qkv, + ): + """Project and attend to one token-flat packed batch.""" + q, k, v = self._project_sequence_packed_qkv(x, position_ids=position_ids, fused_qkv=fused_qkv) + out = self._compute_sequence_packed_attention( + q, + k, + v, + lengths=lengths, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + pos_emb=pos_emb, + padded_length=padded_length, + causal=causal, + sequence_offsets=sequence_offsets, + ) + return self.out_proj(out.reshape(x.shape[0], self.d_model)) + + def _project_sequence_packed_qkv(self, x, *, position_ids, fused_qkv): + """Project packed states to Q/K/V and apply any Q/K preparation.""" + total_tokens = x.shape[0] + if fused_qkv: + qkv = self.w_qkv(x).view(total_tokens, 3, self.n_heads, self.head_dim) + q, k, v = (projection.contiguous() for projection in qkv.unbind(dim=1)) + else: + weights = self.w_qkv.weight.view(3, self.d_model, self.d_model) + biases = self.w_qkv.bias.view(3, self.d_model) if self.w_qkv.bias is not None else None + q, k, v = ( + F.linear(x, weights[idx], None if biases is None else biases[idx]).view( + total_tokens, self.n_heads, self.head_dim + ) + for idx in range(3) + ) + return self._prepare_sequence_packed_qkv(q, k, v, position_ids=position_ids) + + def _prepare_sequence_packed_qkv(self, q, k, v, *, position_ids): + """Apply Q/K normalization and rotary position embeddings when configured.""" + if self.qk_norm: + q = self.q_norm(q).to(v.dtype) + k = self.k_norm(k).to(v.dtype) + + if self._uses_rope: + if position_ids is None: + raise ValueError("Packed RoPE attention requires per-token position_ids.") + q, k = _transformer_utils._apply_packed_rope(self.rope, q, k, position_ids) + return q, k, v + + def _compute_sequence_packed_attention( + self, + q, + k, + v, + *, + lengths, + cu_seqlens, + max_seqlen, + pos_emb, + padded_length, + causal, + sequence_offsets, + ): + """Dispatch packed Q/K/V to the fastest compatible attention backend.""" + flash_attention = _transformer_utils._select_flash_attention_varlen( + q, static_eligible=self._flash_attention_varlen_static_eligible + ) + if flash_attention is not None: + out = flash_attention( + q, + k, + v, + cu_seqlens, + cu_seqlens, + max_seqlen, + max_seqlen, + dropout_p=0.0, + softmax_scale=None, + causal=causal, + ) + self._last_sequence_packed_backend = "flash_attention_varlen" + self._last_sequence_packed_provider = getattr(flash_attention, "_sequence_packed_provider", "external") + else: + use_math_reference = ( + q.device.type == 'cpu' + and torch.is_grad_enabled() + and any(tensor.requires_grad for tensor in (q, k, v)) + ) + out = _transformer_utils._packed_flex_attention_reference( + self, + q, + k, + v, + lengths=lengths, + pos_emb=pos_emb, + padded_length=padded_length, + causal=causal, + sequence_offsets=sequence_offsets, + use_math_reference=use_math_reference, + ) + self._last_sequence_packed_backend = ( + "math_attention_reference" if use_math_reference else "flex_attention_reference" + ) + self._last_sequence_packed_provider = None + return out + class TransformerBlock(nn.Module): def __init__(self, cfg: TransformerEncoderConfig, pos_enc=None): @@ -511,6 +689,37 @@ def forward_streaming(self, x_cur, cache_in, block_mask, pos_emb, cache_size): new_cache = kv_in[:, -cache_size:] return x, new_cache + def _forward_sequence_packed( + self, + x, + *, + lengths, + cu_seqlens, + max_seqlen, + position_ids, + pos_emb, + padded_length, + causal, + sequence_offsets, + fused_qkv, + ): + """Run one Transformer block without materializing batch padding.""" + attn_out = self.attn.forward_sequence_packed( + self.norm1(x), + lengths=lengths, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + position_ids=position_ids, + pos_emb=pos_emb, + padded_length=padded_length, + causal=causal, + sequence_offsets=sequence_offsets, + fused_qkv=fused_qkv, + ) + x = x + self.drop(attn_out) + x = x + self.drop(self.ffn(self.norm2(x))) + return x + @experimental class TransformerEncoder(nn.Module): @@ -593,6 +802,9 @@ class TransformerEncoder(nn.Module): sync_max_audio_length: When true, sync positional encoding allocation length across distributed ranks. """ + supports_sequence_packed_output = True + supports_sequence_packed_fused_qkv = True + def __init__( self, feat_in: int = 128, @@ -654,6 +866,7 @@ def __init__( rotary_fraction=rotary_fraction, ) self.d_model = d_model + self.n_heads = n_heads self.n_layers = n_layers self._feat_in = feat_in self.subsampling = subsampling @@ -867,6 +1080,187 @@ def freeze(self) -> None: def unfreeze(self, partial: bool = False) -> None: unfreeze(self, partial=partial) + def forward_sequence_packed( + self, audio_signal, length, bypass_pre_encode=False, *, fused_qkv: bool = False + ) -> PackedEncoderActivations: + """Encode a batch while keeping all Transformer-layer states token-flat. + + This opt-in method leaves :meth:`forward` and its channels-first padded + return contract unchanged, so existing configs, exports, and checkpoints + retain their historical behavior. + + With nonzero training dropout, removing padding changes random-number + indexing. Packed execution is reproducible within its own path, but it does + not promise same-seed elementwise equality with padded execution. + + ``fused_qkv=True`` trades a small transient projection buffer for a larger, + potentially more efficient projection GEMM. Splitting its interleaved result + requires three compacting copies, so it is an explicit performance option, + not a promise of fewer launches. The default remains the lower-peak + independent projection path. + + Args: + audio_signal: A packed feature batch, a padded mel batch with shape ``(B, C, T)``, + or padded pre-encoded states with shape ``(B, T, D)``. + length: Valid input length for each padded sample, or lengths matching packed input. + bypass_pre_encode: Whether ``audio_signal`` already contains encoder-width states. + fused_qkv: Whether each attention layer should use one fused Q/K/V projection. + + Returns: + Token-flat encoded states and their validated packing metadata. + """ + if self.self_attention_model == "rel_pos" and not getattr(self, "_packed_rel_pos_warned", False): + logging.warning( + "Sequence-packed rel_pos attention uses the compact per-utterance reference backend; " + "use rope, abs_pos, or no_pos for the CUDA varlen fast path." + ) + self._packed_rel_pos_warned = True + if isinstance(audio_signal, PackedEncoderActivations): + if ( + length is not None + and length is not audio_signal.lengths + and not torch.equal(length.to(audio_signal.lengths), audio_signal.lengths) + ): + raise ValueError("length must match audio_signal.lengths for packed input.") + expected_width = self.d_model if bypass_pre_encode else self._feat_in + if audio_signal.data.shape[-1] != expected_width: + raise ValueError( + f"Packed audio_signal must have feature width {expected_width}, " + f"got {audio_signal.data.shape[-1]}." + ) + self.update_max_seq_length(seq_length=audio_signal.max_seqlen, device=audio_signal.data.device) + packed, pos_emb, padded_length = self._prepare_packed_input(audio_signal, bypass_pre_encode) + else: + if not bypass_pre_encode and audio_signal.shape[-2] != self._feat_in: + raise ValueError( + f"If bypass_pre_encode is False, audio_signal should have shape " + f"(batch, {self._feat_in}, n_frame) but got last dimension {audio_signal.shape[-2]}." + ) + if bypass_pre_encode and audio_signal.shape[-1] != self.d_model: + raise ValueError( + f"If bypass_pre_encode is True, audio_signal should have shape " + f"(batch, n_frame, {self.d_model}) but got last dimension {audio_signal.shape[-1]}." + ) + if bypass_pre_encode: + self.update_max_seq_length(seq_length=audio_signal.size(1), device=audio_signal.device) + else: + self.update_max_seq_length(seq_length=audio_signal.size(2), device=audio_signal.device) + x, length, pos_emb = self._prepare_sequence_packed_input(audio_signal, length, bypass_pre_encode) + padded_length = x.shape[1] + packed = pack_encoder_output(x, length) + position_ids = packed_encoder_position_ids(packed) if self.self_attention_model == "rope" else None + x = packed.data + fast_path = ( + self.self_attention_model != "rel_pos" + and _transformer_utils._can_use_flash_attention_varlen_layout( + x, + self.d_model // self.n_heads, + ) + ) + sequence_offsets = None if fast_path else tuple(packed.cu_seqlens.tolist()) + for layer in self.layers: + x = _transformer_utils._forward_sequence_packed_layer( + layer, + x, + lengths=packed.lengths, + cu_seqlens=packed.cu_seqlens, + max_seqlen=packed.max_seqlen, + position_ids=position_ids, + pos_emb=pos_emb if self.self_attention_model == "rel_pos" else None, + padded_length=padded_length, + causal=self.attn_mode == "causal", + sequence_offsets=sequence_offsets, + fused_qkv=fused_qkv, + ) + x = self.final_norm(x) + if self.out_proj is not None: + x = self.out_proj(x) + return packed.with_data(x) + + def _prepare_sequence_packed_input(self, audio_signal, length, bypass_pre_encode): + """Prepare padded input for packing while preserving the ordinary frontend path.""" + if length is None: + length = audio_signal.new_full( + (audio_signal.size(0),), + audio_signal.size(1) if bypass_pre_encode else audio_signal.size(-1), + dtype=torch.int64, + device=audio_signal.device, + ) + + if not bypass_pre_encode: + # Unwrap activation-checkpointing (CheckpointWrapper) and match by name: both + # the wrapper and duplicate module copies defeat isinstance(FeatureStacking). + pre_encode_module = getattr(self.pre_encode, "_checkpoint_wrapped_module", self.pre_encode) + is_feature_stacking = type(pre_encode_module).__name__ == "FeatureStacking" + if is_feature_stacking: + x, length = self.pre_encode(audio_signal, length) + else: + x = torch.transpose(audio_signal, 1, 2) + if isinstance(pre_encode_module, nn.Linear): + x = self.pre_encode(x) + elif not is_feature_stacking: + x, length = self.pre_encode(x=x, lengths=length) + length = length.to(torch.int64) + else: + x = audio_signal + length = length.to(torch.int64) + + if self.self_attention_model == "rope": + if self.xscale: + x = x * self.xscale + x = self.dropout_pre_encoder(x) + pos_emb = None + elif self.pos_enc is not None: + x, pos_emb = self.pos_enc(x=x) + else: + pos_emb = None + return self.embed_norm(x), length, pos_emb + + def _prepare_packed_input(self, audio_signal: PackedEncoderActivations, bypass_pre_encode: bool): + """Apply supported pre-encoding and position handling to packed input.""" + if bypass_pre_encode: + packed = audio_signal + else: + pre_encode_module = getattr(self.pre_encode, "_checkpoint_wrapped_module", self.pre_encode) + if type(pre_encode_module).__name__ != "FeatureStacking": + raise TypeError( + "Packed feature input currently requires subsampling='feature_stacking'; " + f"got {type(pre_encode_module).__name__}." + ) + packed = self.pre_encode(audio_signal) + return self._apply_packed_position(packed) + + def _apply_packed_position(self, packed: PackedEncoderActivations): + """Apply the configured positional encoding directly to token-flat states.""" + x = packed.data + if self.self_attention_model == "rope": + if self.xscale: + x = x * self.xscale + x = self.dropout_pre_encoder(x) + pos_emb = None + elif self.self_attention_model == "abs_pos": + position_ids = packed_encoder_position_ids(packed) + if self.pos_enc.xscale: + x = x * self.pos_enc.xscale + pos_emb = self.pos_enc.pe[:, : packed.max_seqlen] + token_pos_emb = pos_emb[0].index_select(0, position_ids) + if self.pos_enc.dropout_emb: + token_pos_emb = self.pos_enc.dropout_emb(token_pos_emb) + x = self.pos_enc.dropout(x + token_pos_emb) + elif self.self_attention_model == "rel_pos": + if self.pos_enc.xscale: + x = x * self.pos_enc.xscale + x = self.pos_enc.dropout(x) + center_pos = self.pos_enc.pe.size(1) // 2 + 1 + start_pos = center_pos - packed.max_seqlen + end_pos = center_pos + packed.max_seqlen - 1 + pos_emb = self.pos_enc.pe[:, start_pos:end_pos] + if self.pos_enc.dropout_emb: + pos_emb = self.pos_enc.dropout_emb(pos_emb) + else: + pos_emb = None + return packed.with_data(self.embed_norm(x)), pos_emb, packed.max_seqlen + @experimental class StreamingTransformerEncoder(TransformerEncoder, StreamingEncoder): diff --git a/nemo/collections/asr/modules/transformer_encoder_utils.py b/nemo/collections/asr/modules/transformer_encoder_utils.py new file mode 100644 index 000000000000..5a8b35a8530f --- /dev/null +++ b/nemo/collections/asr/modules/transformer_encoder_utils.py @@ -0,0 +1,188 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Private packed-attention helpers shared by Transformer encoder execution paths.""" + +from functools import lru_cache + +import torch +from torch.nn.attention.flex_attention import create_block_mask, flex_attention + + +_flex_attention_compiled = torch.compile(flex_attention, dynamic=True) + + +def _get_flex_attention(x): + """Select compiled CUDA FlexAttention or its eager CPU implementation.""" + return _flex_attention_compiled if x.is_cuda else flex_attention + + +def _causal_mask(b, h, q_idx, kv_idx): + """Return whether a query may attend to a key under causal attention.""" + return q_idx >= kv_idx + + +def _apply_packed_rope(rope, q, k, position_ids): + """Apply rotary embeddings to token-flat query and key tensors.""" + cos = rope.cos.index_select(0, position_ids).unsqueeze(1).to(q.dtype) + sin = rope.sin.index_select(0, position_ids).unsqueeze(1).to(q.dtype) + return rope._apply_rotary(q, cos, sin), rope._apply_rotary(k, cos.to(k.dtype), sin.to(k.dtype)) + + +def _packed_flex_attention_reference( + attn, + q, + k, + v, + *, + lengths, + pos_emb, + padded_length, + causal, + sequence_offsets, + use_math_reference=False, +): + """Evaluate packed attention one sequence at a time without padded activations.""" + if sequence_offsets is None: + sequence_offsets = tuple(torch.cat([lengths.new_zeros(1), lengths.cumsum(0)]).tolist()) + outputs = [] + attn_fn = _get_flex_attention(q) + for offset, end in zip(sequence_offsets[:-1], sequence_offsets[1:]): + length = end - offset + if length == 0: + continue + qi = q[offset:end].transpose(0, 1).unsqueeze(0) + ki = k[offset:end].transpose(0, 1).unsqueeze(0) + vi = v[offset:end].transpose(0, 1).unsqueeze(0) + score_mod = None + if attn._uses_rel_pos: + if pos_emb is None or padded_length is None: + raise ValueError("Packed relative-position attention requires max-length positional metadata.") + pos_i = pos_emb[:, padded_length - length : padded_length + length - 1] + score_mod, qi = attn._build_rel_pos_score_mod(qi, pos_i) + if use_math_reference: + out = _packed_math_attention_reference(qi, ki, vi, causal=causal, score_mod=score_mod) + else: + block_mask = None + if causal: + block_mask = create_block_mask(_causal_mask, B=1, H=1, Q_LEN=length, KV_LEN=length, device=q.device) + out = attn_fn(qi, ki, vi, block_mask=block_mask, score_mod=score_mod) + outputs.append(out.squeeze(0).transpose(0, 1)) + if not outputs: + # Keep every attention branch in the autograd graph even when a rank owns + # no valid tokens. FSDP/DDP otherwise observes missing gradients for q/k + # (and relative-position parameters), which can break collectives when + # another rank in the same step has non-empty input. + anchor = q.sum() + k.sum() + if attn._uses_rel_pos: + anchor = anchor + 0.0 * (attn.pos_bias_u.sum() + attn.pos_bias_v.sum()) + anchor = anchor + sum(0.0 * parameter.sum() for parameter in attn.linear_pos.parameters()) + return v + anchor.to(v.dtype) + return torch.cat(outputs, dim=0) + + +def _packed_math_attention_reference(q, k, v, *, causal, score_mod): + """Compute differentiable CPU reference attention for one packed sequence.""" + scores = torch.matmul(q, k.transpose(-2, -1)) * (q.shape[-1] ** -0.5) + if score_mod is not None: + scores = scores + score_mod._relative_position_bias + if causal: + causal_mask = torch.ones(scores.shape[-2:], dtype=torch.bool, device=scores.device).tril() + scores = scores.masked_fill(~causal_mask, torch.finfo(scores.dtype).min) + return torch.matmul(torch.softmax(scores, dim=-1).to(v.dtype), v) + + +def _select_flash_attention_varlen(x, *, static_eligible): + """Return the varlen provider after cheap per-input checks and cached device probing.""" + if not static_eligible or not x.is_cuda or x.dtype not in (torch.float16, torch.bfloat16) or x.shape[0] == 0: + return None + return _get_flash_attention_varlen_for_device(x.device) + + +def _can_use_flash_attention_varlen(q): + """Return whether packed Q can use an available variable-length FlashAttention provider.""" + static_eligible = q.shape[-1] <= 256 and q.shape[-1] % 8 == 0 + return _select_flash_attention_varlen(q, static_eligible=static_eligible) is not None + + +def _can_use_flash_attention_varlen_layout(x, head_dim): + """Return whether a packed layout can use variable-length FlashAttention.""" + static_eligible = head_dim <= 256 and head_dim % 8 == 0 + return _select_flash_attention_varlen(x, static_eligible=static_eligible) is not None + + +@lru_cache(maxsize=None) +def _get_flash_attention_varlen_for_device(device): + """Resolve and cache a variable-length FlashAttention provider for one CUDA device.""" + if torch.version.cuda is None or torch.cuda.get_device_capability(device)[0] < 8: + return None + return _get_flash_attention_varlen() + + +@lru_cache(maxsize=1) +def _get_flash_attention_varlen(): + """Resolve the external or ATen variable-length FlashAttention implementation.""" + try: + from flash_attn import flash_attn_varlen_func + except (ImportError, ModuleNotFoundError): + flash_forward = getattr(torch.ops.aten, "_flash_attention_forward", None) + if flash_forward is None: + return None + + def torch_flash_attention_varlen( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + *, + dropout_p, + softmax_scale, + causal, + ): + """Adapt the ATen FlashAttention operator to the external provider signature.""" + return flash_forward( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + dropout_p, + causal, + False, + scale=softmax_scale, + )[0] + + torch_flash_attention_varlen._sequence_packed_provider = "aten" + return torch_flash_attention_varlen + return flash_attn_varlen_func + + +def _forward_sequence_packed_layer(layer, x, **kwargs): + """Preserve packed execution through PyTorch's activation-checkpoint wrapper.""" + wrapped = getattr(layer, '_checkpoint_wrapped_module', None) + if wrapped is None: + return layer._forward_sequence_packed(x, **kwargs) + packed_forward = getattr(wrapped, '_forward_sequence_packed', None) + checkpoint_fn = getattr(layer, 'checkpoint_fn', None) + if packed_forward is None or checkpoint_fn is None: + raise TypeError( + f"Activation-checkpoint wrapper around {type(wrapped).__name__} cannot execute sequence-packed layers." + ) + return checkpoint_fn(packed_forward, x, **kwargs) diff --git a/nemo/collections/asr/parts/packed_sequence.py b/nemo/collections/asr/parts/packed_sequence.py new file mode 100644 index 000000000000..15cc1f8daf7a --- /dev/null +++ b/nemo/collections/asr/parts/packed_sequence.py @@ -0,0 +1,324 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Token-flat encoder activations and conversions for sequence-packed ASR execution.""" + +from dataclasses import dataclass + +import torch +from torch import Tensor +from torch.utils._pytree import GetAttrKey, register_pytree_node + +__all__ = [ + "PackedEncoderActivations", + "pack_encoder_output", + "packed_encoder_position_ids", + "split_packed_data", + "split_encoder_output", + "unpack_encoder_output", +] + + +@dataclass(frozen=True) +class PackedEncoderActivations: + """An ASR encoder batch stored without inter-utterance padding. + + Args: + data: Valid encoder states concatenated in batch order, shape ``(T_total, D)``. + lengths: Per-utterance token counts, shape ``(B,)`` and dtype ``int64``. + cu_seqlens: Cumulative token counts, shape ``(B + 1,)`` and dtype ``int32``. + max_seqlen: Largest value in ``lengths`` (zero for an empty batch). + padding_value: Value to use when a downstream frame-stacking operation + needs to complete the final partial group. A scalar applies to the + whole batch; a ``(B, D)`` tensor supports sequence-specific normalized + feature padding. Encoder outputs use zero. + padded_length: Optional dense time dimension associated with the packed + data. Frontends preserve it solely for dense-compatible augmentation + RNG ranges; packed kernels continue to use ``max_seqlen``. + + The representation deliberately has an ASR-specific name rather than reusing + :class:`torch.nn.utils.rnn.PackedSequence`: the latter stores sort/unsort metadata + for recurrent networks and has a different layout contract. + """ + + data: Tensor + lengths: Tensor + cu_seqlens: Tensor + max_seqlen: int + padding_value: float | Tensor = 0.0 + padded_length: int | None = None + + def __post_init__(self) -> None: + _validate_packed_encoder_activations(self) + + @property + def batch_size(self) -> int: + """Number of utterances represented by this output.""" + + return int(self.lengths.numel()) + + @property + def total_tokens(self) -> int: + """Number of valid (unpadded) encoder tokens.""" + + return int(self.data.shape[0]) + + def with_data(self, data: Tensor) -> "PackedEncoderActivations": + """Reuse validated sequence metadata with replacement token data. + + Internal encoder stages use this to avoid revalidating unchanged CUDA + offsets (and synchronizing the host) after every projection or fusion. + """ + if data.ndim != 2 or data.shape[0] != self.total_tokens: + raise ValueError(f"replacement data must have shape ({self.total_tokens}, D), got {tuple(data.shape)}.") + if data.device != self.data.device: + raise ValueError(f"replacement data must be on {self.data.device}, got {data.device}.") + return _new_packed_encoder_activations( + data, + self.lengths, + self.cu_seqlens, + self.max_seqlen, + padding_value=self.padding_value, + padded_length=self.padded_length, + ) + + +def _packed_encoder_activations_flatten(packed: PackedEncoderActivations): + values = [ + packed.data, + packed.lengths, + packed.cu_seqlens, + packed.max_seqlen, + packed.padding_value, + packed.padded_length, + ] + return values, None + + +def _packed_encoder_activations_flatten_with_keys(packed: PackedEncoderActivations): + values, context = _packed_encoder_activations_flatten(packed) + names = ( + "data", + "lengths", + "cu_seqlens", + "max_seqlen", + "padding_value", + "padded_length", + ) + keyed_values = [(GetAttrKey(name), value) for name, value in zip(names, values, strict=True)] + return keyed_values, context + + +def _packed_encoder_activations_unflatten(values, context): + del context + # Pytree transforms may temporarily replace every leaf with a sentinel. In + # particular, torch.utils.checkpoint._infer_device_type maps all leaves to + # None merely to inspect CUDA tensors. Re-running public-constructor + # validation on that transient object would reject the transform itself. + # Real PackedEncoderActivations still enter through the validating public + # constructor; pytree reconstruction preserves already-validated metadata. + return _new_packed_encoder_activations(*values) + + +# FSDP2 discovers tensors that need post-forward/pre-backward hooks by walking +# the output pytree. Register this container so gradients from ``data`` reach +# the optimizer-owned sharded parameters when a custom packed forward is used. +# Use a validation-free unflatten function because generic pytree transforms +# are allowed to replace leaves with temporary non-domain values. +register_pytree_node( + PackedEncoderActivations, + _packed_encoder_activations_flatten, + _packed_encoder_activations_unflatten, + serialized_type_name="nemo.collections.asr.parts.packed_sequence.PackedEncoderActivations", + flatten_with_keys_fn=_packed_encoder_activations_flatten_with_keys, +) + + +def pack_encoder_output(padded: Tensor, lengths: Tensor) -> PackedEncoderActivations: + """Compact valid prefixes from a channels-last encoder batch. + + Args: + padded: Encoder states with shape ``(B, T, D)``. + lengths: Valid prefix lengths with shape ``(B,)``. + + Returns: + A :class:`PackedEncoderActivations` whose data has shape ``(sum(lengths), D)``. + """ + + if padded.ndim != 3: + raise ValueError(f"padded must have shape (B, T, D), got {tuple(padded.shape)}.") + lengths, max_seqlen = _normalize_lengths( + lengths, batch_size=padded.shape[0], max_length=padded.shape[1], device=padded.device + ) + mask = _length_mask(lengths, padded.shape[1]) + data = padded[mask] + cu_seqlens = _lengths_to_cu_seqlens(lengths) + return _new_packed_encoder_activations(data, lengths, cu_seqlens, max_seqlen, padded_length=padded.shape[1]) + + +def unpack_encoder_output(packed: PackedEncoderActivations, *, total_length: int | None = None) -> Tensor: + """Restore a packed encoder output to channels-last ``(B, T, D)`` form. + + ``total_length`` defaults to ``packed.max_seqlen``. A larger value is useful at + legacy boundaries that must preserve an externally chosen padded width. + """ + + if total_length is None: + total_length = packed.max_seqlen + if int(total_length) != total_length or total_length < packed.max_seqlen: + raise ValueError(f"total_length must be an integer >= max_seqlen ({packed.max_seqlen}), got {total_length}.") + total_length = int(total_length) + if isinstance(packed.padding_value, Tensor): + padding = packed.padding_value.to(device=packed.data.device, dtype=packed.data.dtype) + output = padding.unsqueeze(1).expand(-1, total_length, -1).clone() + else: + output = packed.data.new_full((packed.batch_size, total_length, packed.data.shape[-1]), packed.padding_value) + if packed.total_tokens: + output[_length_mask(packed.lengths, total_length)] = packed.data + return output + + +def split_encoder_output(packed: PackedEncoderActivations) -> tuple[Tensor, ...]: + """Return one unpadded ``(T_i, D)`` view per utterance.""" + + offsets = packed.cu_seqlens.tolist() + return tuple(packed.data[offsets[i] : offsets[i + 1]] for i in range(packed.batch_size)) + + +def packed_encoder_position_ids(packed: PackedEncoderActivations) -> Tensor: + """Return zero-based positions for every packed token, resetting per utterance.""" + + if packed.total_tokens == 0: + return packed.lengths.new_empty((0,)) + token_indices = torch.arange(packed.total_tokens, device=packed.data.device, dtype=torch.int64) + sequence_ids = torch.bucketize(token_indices, packed.cu_seqlens[1:], right=True) + return token_indices - packed.cu_seqlens[sequence_ids] + + +def split_packed_data(data: Tensor, lengths: Tensor, cu_seqlens: Tensor) -> tuple[Tensor, ...]: + """Validate packed leading-dimension metadata and return one view per sequence. + + This lower-level representation is useful before encoder features exist, such as + for concatenated waveform samples. Unlike :class:`PackedEncoderActivations`, it allows + any data rank and integer offset dtype. + """ + + if data.ndim == 0: + raise ValueError("packed data must have at least one dimension.") + if lengths.ndim != 1: + raise ValueError(f"lengths must be 1D, got shape {tuple(lengths.shape)}.") + if lengths.dtype == torch.bool or lengths.is_floating_point() or lengths.is_complex(): + raise TypeError(f"lengths must have an integer dtype, got {lengths.dtype}.") + if cu_seqlens.ndim != 1 or cu_seqlens.numel() != lengths.numel() + 1: + raise ValueError(f"cu_seqlens must have shape ({lengths.numel() + 1},), got {tuple(cu_seqlens.shape)}.") + if cu_seqlens.dtype == torch.bool or cu_seqlens.is_floating_point() or cu_seqlens.is_complex(): + raise TypeError(f"cu_seqlens must have an integer dtype, got {cu_seqlens.dtype}.") + if data.device != lengths.device or data.device != cu_seqlens.device: + raise ValueError("data, lengths, and cu_seqlens must be on the same device.") + + offsets = cu_seqlens.to(torch.int64) + if not torch.equal(offsets[1:] - offsets[:-1], lengths.to(torch.int64)): + raise ValueError("Differences in cu_seqlens must equal lengths.") + host_offsets = offsets.detach().cpu().tolist() + if host_offsets[0] != 0: + raise ValueError("cu_seqlens must start at zero.") + if any(end < begin for begin, end in zip(host_offsets, host_offsets[1:])): + raise ValueError("cu_seqlens must be non-decreasing.") + if host_offsets[-1] != data.shape[0]: + raise ValueError(f"data has {data.shape[0]} entries, but cu_seqlens ends at {host_offsets[-1]}.") + return tuple(data[begin:end] for begin, end in zip(host_offsets, host_offsets[1:])) + + +def _length_mask(lengths: Tensor, total_length: int) -> Tensor: + positions = torch.arange(total_length, device=lengths.device) + return positions.unsqueeze(0) < lengths.unsqueeze(1) + + +def _lengths_to_cu_seqlens(lengths: Tensor) -> Tensor: + return torch.cat( + [ + torch.zeros(1, dtype=torch.int32, device=lengths.device), + lengths.cumsum(dim=0, dtype=torch.int32), + ] + ).contiguous() + + +def _normalize_lengths(lengths: Tensor, *, batch_size: int, max_length: int, device: torch.device): + if lengths.ndim != 1 or lengths.numel() != batch_size: + raise ValueError(f"lengths must have shape ({batch_size},), got {tuple(lengths.shape)}.") + if lengths.device != device: + raise ValueError(f"lengths must be on {device}, got {lengths.device}.") + if lengths.dtype == torch.bool or lengths.is_floating_point() or lengths.is_complex(): + raise TypeError(f"lengths must have an integer dtype, got {lengths.dtype}.") + lengths = lengths.to(torch.int64) + # Varlen kernels need a host max length. Validate the same host copy so the + # hot packing path pays one device synchronization rather than one per check. + host_lengths = lengths.detach().to(device="cpu") + if bool(((host_lengths < 0) | (host_lengths > max_length)).any()): + raise ValueError(f"lengths must be between 0 and padded time {max_length}, got {host_lengths.tolist()}.") + max_seqlen = int(host_lengths.max()) if host_lengths.numel() else 0 + return lengths, max_seqlen + + +def _validate_packed_encoder_activations(packed: PackedEncoderActivations) -> None: + if packed.data.ndim != 2: + raise ValueError(f"data must have shape (T_total, D), got {tuple(packed.data.shape)}.") + if packed.lengths.ndim != 1 or packed.lengths.dtype != torch.int64: + raise ValueError(f"lengths must be 1D int64, got shape={tuple(packed.lengths.shape)}, {packed.lengths.dtype}.") + if packed.cu_seqlens.ndim != 1 or packed.cu_seqlens.dtype != torch.int32: + raise ValueError( + "cu_seqlens must be 1D int32, got " f"shape={tuple(packed.cu_seqlens.shape)}, {packed.cu_seqlens.dtype}." + ) + if not packed.cu_seqlens.is_contiguous(): + raise ValueError("cu_seqlens must be contiguous for variable-length attention kernels.") + if packed.data.device != packed.lengths.device or packed.data.device != packed.cu_seqlens.device: + raise ValueError("data, lengths, and cu_seqlens must be on the same device.") + if packed.cu_seqlens.numel() != packed.lengths.numel() + 1: + raise ValueError("cu_seqlens must have exactly batch_size + 1 entries.") + expected = _lengths_to_cu_seqlens(packed.lengths) + if not torch.equal(packed.cu_seqlens, expected): + raise ValueError("cu_seqlens must start at zero and have differences equal to lengths.") + if packed.data.shape[0] != int(expected[-1].item()): + raise ValueError(f"data has {packed.data.shape[0]} tokens, but cu_seqlens ends at {int(expected[-1].item())}.") + expected_max = int(packed.lengths.max().item()) if packed.lengths.numel() else 0 + if int(packed.max_seqlen) != packed.max_seqlen or packed.max_seqlen != expected_max: + raise ValueError(f"max_seqlen must equal max(lengths)={expected_max}, got {packed.max_seqlen}.") + if isinstance(packed.padding_value, Tensor): + expected_shape = (packed.batch_size, packed.data.shape[1]) + if packed.padding_value.shape != expected_shape: + raise ValueError( + f"tensor padding_value must have shape {expected_shape}, got {tuple(packed.padding_value.shape)}." + ) + if packed.padding_value.device != packed.data.device: + raise ValueError( + f"tensor padding_value must be on {packed.data.device}, got {packed.padding_value.device}." + ) + if packed.padded_length is not None and ( + int(packed.padded_length) != packed.padded_length or packed.padded_length < packed.max_seqlen + ): + raise ValueError( + f"padded_length must be an integer >= max_seqlen ({packed.max_seqlen}), got {packed.padded_length}." + ) + + +def _new_packed_encoder_activations(data, lengths, cu_seqlens, max_seqlen, padding_value=0.0, padded_length=None): + """Construct from metadata that a public constructor already validated.""" + packed = object.__new__(PackedEncoderActivations) + object.__setattr__(packed, "data", data) + object.__setattr__(packed, "lengths", lengths) + object.__setattr__(packed, "cu_seqlens", cu_seqlens) + object.__setattr__(packed, "max_seqlen", max_seqlen) + object.__setattr__(packed, "padding_value", padding_value) + object.__setattr__(packed, "padded_length", padded_length) + return packed diff --git a/nemo/collections/asr/parts/preprocessing/features.py b/nemo/collections/asr/parts/preprocessing/features.py index 75ce616ee7f5..b0da015ec2d4 100644 --- a/nemo/collections/asr/parts/preprocessing/features.py +++ b/nemo/collections/asr/parts/preprocessing/features.py @@ -41,6 +41,7 @@ import torch import torch.nn as nn +from nemo.collections.asr.parts.packed_sequence import PackedEncoderActivations, _new_packed_encoder_activations from nemo.collections.asr.parts.preprocessing.perturb import AudioAugmentor from nemo.collections.asr.parts.preprocessing.segment import AudioSegment from nemo.utils import logging @@ -70,9 +71,13 @@ def normalize_batch(x, seq_len, normalize_type): ) time_steps = torch.arange(max_time, device=x.device).unsqueeze(0).expand(batch_size, max_time) valid_mask = time_steps < seq_len.unsqueeze(1) - x_mean_numerator = torch.where(valid_mask.unsqueeze(1), x, 0.0).sum(axis=2) x_mean_denominator = valid_mask.sum(axis=1) - x_mean = x_mean_numerator / x_mean_denominator.unsqueeze(1) + # Reference-centering keeps constant inputs exact across reduction backends + # and matches the packed normalization path. + reference = x[:, :, 0] if max_time else x.new_zeros((batch_size, x.shape[1])) + reference = reference.masked_fill((x_mean_denominator == 0).unsqueeze(1), 0.0) + centered = torch.where(valid_mask.unsqueeze(1), x - reference.unsqueeze(2), 0.0) + x_mean = reference + centered.sum(axis=2) / x_mean_denominator.clamp_min(1).unsqueeze(1) # Subtract 1 in the denominator to correct for the bias. x_std = torch.sqrt( @@ -106,6 +111,28 @@ def normalize_batch(x, seq_len, normalize_type): return x, x_mean, x_std +def normalize_packed_batch(packed: PackedEncoderActivations, normalize_type) -> PackedEncoderActivations: + """Normalize token-flat features independently within each sequence.""" + if not normalize_type or packed.total_tokens == 0: + return packed + sequence_ids = torch.repeat_interleave(torch.arange(packed.batch_size, device=packed.data.device), packed.lengths) + data, padding_value = _normalize_packed_features_and_padding( + packed.data, + packed.lengths, + sequence_ids, + normalize_type, + padding_value=packed.padding_value, + ) + return _new_packed_encoder_activations( + data, + packed.lengths, + packed.cu_seqlens, + packed.max_seqlen, + padding_value, + padded_length=packed.padded_length, + ) + + def clean_spectrogram_batch(spectrogram: torch.Tensor, spectrogram_len: torch.Tensor, fill_value=0.0) -> torch.Tensor: """ Fill spectrogram values outside the length with `fill_value` @@ -355,13 +382,14 @@ def __init__( self.use_grads = use_grads if not use_grads: self.forward = torch.no_grad()(self.forward) + self.forward_packed = torch.no_grad()(self.forward_packed) self._rng = random.Random() if rng is None else rng self.nb_augmentation_prob = nb_augmentation_prob if self.nb_augmentation_prob > 0.0: if nb_max_freq >= sample_rate / 2: self.nb_augmentation_prob = 0.0 else: - self._nb_max_fft_bin = int((nb_max_freq / sample_rate) * n_fft) + self._nb_max_fft_bin = int((nb_max_freq / sample_rate) * self.n_fft) # log_zero_guard_value is the the small we want to use, we support # an actual number, or "tiny", or "eps" @@ -376,13 +404,15 @@ def __init__( logging.debug(f"using grads: {use_grads}") logging.debug(f"nb_augmentation_prob: {nb_augmentation_prob}") - def stft(self, x): + def stft(self, x, *, center=None): + if center is None: + center = not self.exact_pad return torch.stft( x, n_fft=self.n_fft, hop_length=self.hop_length, win_length=self.win_length, - center=False if self.exact_pad else True, + center=center, window=self.window.to(dtype=torch.float, device=x.device), return_complex=True, pad_mode="constant", @@ -413,6 +443,123 @@ def get_seq_len(self, seq_len): def filter_banks(self): return self.fb + def forward_packed(self, x, seq_len, cu_seqlens, linear_spec=False) -> PackedEncoderActivations: + """Compute features from concatenated waveforms with one vectorized STFT. + + Each utterance is placed in a hop-aligned block with the same zero guard + that the dense STFT applies at its boundaries. Only valid frames are + gathered from the resulting single STFT, so both input and output remain + sequence-packed and no ``B x T`` waveform or feature tensor is created. + + ``pad_to`` is intentionally ignored: it is a dense-layout optimization and + packed output contains exactly ``sum(output_lengths)`` frames. + """ + seq_len, cu_seqlens, host_seq_len = _validate_packed_waveforms(x, seq_len, cu_seqlens) + feature_lengths = torch.where(seq_len == 0, 0, self.get_seq_len(seq_len)) + host_feature_lengths = torch.where(host_seq_len == 0, 0, self.get_seq_len(host_seq_len)) + padded_length = _dense_feature_width(host_feature_lengths, self.pad_to, self.max_length) + max_seqlen = int(host_feature_lengths.max()) if host_feature_lengths.numel() else 0 + total_frames = int(host_feature_lengths.sum()) + if bool((host_feature_lengths < 0).any()): + raise ValueError( + "Packed waveform lengths are too short for this STFT configuration; " + f"computed feature lengths {host_feature_lengths.tolist()}." + ) + if seq_len.numel() == 0 or total_frames == 0: + feature_dim = self.n_fft // 2 + 1 if linear_spec else self.nfilt * self.frame_splicing + return _empty_packed_features( + x, + feature_lengths, + feature_dim, + padding_value=self.pad_value, + padded_length=padded_length, + max_seqlen=max_seqlen, + ) + + guard = self.stft_pad_amount if self.stft_pad_amount is not None else self.n_fft // 2 + block_lengths = _round_up(seq_len + 2 * guard, self.hop_length) + block_offsets = torch.cat([seq_len.new_zeros(1), block_lengths.cumsum(0)]) + guarded_size = int(_round_up(host_seq_len + 2 * guard, self.hop_length).sum()) + guarded = x.new_zeros(guarded_size) + + guarded_positions = torch.arange(x.numel(), device=x.device) + sample_sequence_ids = torch.bucketize(guarded_positions, cu_seqlens[1:], right=True) + guarded_positions -= cu_seqlens[sample_sequence_ids] + guarded_positions += block_offsets[sample_sequence_ids] + guard + + if self.stft_pad_amount is None: + samples = _dither_and_preemphasize_packed( + x, seq_len, cu_seqlens, self.preemph, self.dither if self.training else 0.0 + ) + guarded[guarded_positions] = samples + else: + guarded[guarded_positions] = x + guarded = _dither_and_preemphasize_exact_pad_blocks( + guarded, + seq_len, + block_lengths, + block_offsets, + self.preemph, + self.dither if self.training else 0.0, + ) + del guarded_positions, sample_sequence_ids + + with torch.amp.autocast(x.device.type, enabled=False): + spectra = self.stft(guarded.unsqueeze(0), center=False)[0] + + frame_cu_seqlens = torch.cat([feature_lengths.new_zeros(1), feature_lengths.cumsum(0)]) + frame_indices = torch.arange(total_frames, device=x.device) + frame_sequence_ids = torch.bucketize(frame_indices, frame_cu_seqlens[1:], right=True) + local_frames = frame_indices - frame_cu_seqlens[frame_sequence_ids] + global_frames = torch.div(block_offsets[frame_sequence_ids], self.hop_length, rounding_mode="floor") + global_frames = global_frames + local_frames + spectra = spectra.index_select(-1, global_frames).transpose(0, 1) + + guard_value = 0 if not self.use_grads else CONSTANT + spectra = torch.sqrt(torch.view_as_real(spectra).pow(2).sum(-1) + guard_value) + if self.training and self.nb_augmentation_prob > 0.0: + narrowband = torch.tensor( + self._rng.choices( + (True, False), + weights=(self.nb_augmentation_prob, 1.0 - self.nb_augmentation_prob), + k=feature_lengths.numel(), + ), + device=x.device, + ) + keep = ~(narrowband[frame_sequence_ids].unsqueeze(1) & _high_frequency_mask(spectra, self._nb_max_fft_bin)) + spectra = spectra * keep + if self.mag_power != 1.0: + spectra = spectra.pow(self.mag_power) + if linear_spec: + return _make_packed_features( + spectra, + feature_lengths, + padding_value=self.pad_value, + padded_length=padded_length, + max_seqlen=max_seqlen, + ) + + with torch.amp.autocast(x.device.type, enabled=False): + features = torch.matmul(self.fb.to(spectra.dtype), spectra.transpose(0, 1).unsqueeze(0))[0].transpose(0, 1) + if self.log: + if self.log_zero_guard_type == "add": + features = torch.log(features + self.log_zero_guard_value_fn(features)) + elif self.log_zero_guard_type == "clamp": + features = torch.log(torch.clamp(features, min=self.log_zero_guard_value_fn(features))) + else: + raise ValueError("log_zero_guard_type was not understood") + if self.frame_splicing > 1: + features = features.repeat(1, self.frame_splicing) + if self.normalize: + features = _normalize_packed_features(features, feature_lengths, frame_sequence_ids, self.normalize) + return _make_packed_features( + features, + feature_lengths, + padding_value=self.pad_value, + padded_length=padded_length, + max_seqlen=max_seqlen, + ) + def forward(self, x, seq_len, linear_spec=False): seq_len_time = seq_len seq_len_unfixed = self.get_seq_len(seq_len) @@ -493,3 +640,169 @@ def forward(self, x, seq_len, linear_spec=False): if pad_amt != 0: x = nn.functional.pad(x, (0, pad_to - pad_amt), value=self.pad_value) return x, seq_len + + +def _validate_packed_waveforms(x, seq_len, cu_seqlens): + if x.ndim != 1: + raise ValueError(f"packed waveform data must be 1D, got shape {tuple(x.shape)}.") + if seq_len.ndim != 1: + raise ValueError(f"length must be 1D, got shape {tuple(seq_len.shape)}.") + if seq_len.dtype == torch.bool or seq_len.is_floating_point() or seq_len.is_complex(): + raise TypeError(f"length must have an integer dtype, got {seq_len.dtype}.") + if cu_seqlens.ndim != 1 or cu_seqlens.numel() != seq_len.numel() + 1: + raise ValueError(f"cu_seqlens must have shape ({seq_len.numel() + 1},), got {tuple(cu_seqlens.shape)}.") + if cu_seqlens.dtype == torch.bool or cu_seqlens.is_floating_point() or cu_seqlens.is_complex(): + raise TypeError(f"cu_seqlens must have an integer dtype, got {cu_seqlens.dtype}.") + if x.device != seq_len.device or x.device != cu_seqlens.device: + raise ValueError("packed waveform data, length, and cu_seqlens must be on the same device.") + seq_len = seq_len.to(torch.int64) + cu_seqlens = cu_seqlens.to(torch.int64) + host_metadata = torch.cat((seq_len, cu_seqlens)).detach().cpu() + host_seq_len = host_metadata[: seq_len.numel()] + host_cu_seqlens = host_metadata[seq_len.numel() :] + if host_cu_seqlens.numel() and int(host_cu_seqlens[0]) != 0: + raise ValueError("cu_seqlens must start at zero.") + if not torch.equal(host_cu_seqlens[1:] - host_cu_seqlens[:-1], host_seq_len): + raise ValueError("Differences in cu_seqlens must equal length.") + if bool((host_seq_len < 0).any()): + raise ValueError("length must be non-negative.") + if int(host_cu_seqlens[-1]) != x.shape[0]: + raise ValueError( + f"packed waveform data has {x.shape[0]} samples, but cu_seqlens ends at {host_cu_seqlens[-1]}." + ) + return seq_len, cu_seqlens, host_seq_len + + +def _round_up(values, multiple): + return torch.div(values + multiple - 1, multiple, rounding_mode="floor") * multiple + + +def _dense_feature_width(lengths, pad_to, max_length): + width = int(lengths.max().item()) + 1 if lengths.numel() else 0 + if pad_to == "max": + return int(max_length) + if pad_to > 0 and width % pad_to: + width += pad_to - width % pad_to + return width + + +def _dither_and_preemphasize_packed(x, lengths, cu_seqlens, preemph, dither): + samples = x + dither * torch.randn_like(x) if dither > 0 else x + if preemph is None or samples.numel() == 0: + return samples + emphasized = torch.cat([samples[:1], samples[1:] - preemph * samples[:-1]]) + starts = cu_seqlens[:-1][lengths > 0] + return emphasized.scatter(0, starts, samples.index_select(0, starts)) + + +def _dither_and_preemphasize_exact_pad_blocks(guarded, lengths, block_lengths, block_offsets, preemph, dither): + if dither > 0: + guarded = guarded + dither * torch.randn_like(guarded) + if preemph is not None and guarded.numel() > 0: + emphasized = torch.cat([guarded[:1], guarded[1:] - preemph * guarded[:-1]]) + starts = block_offsets[:-1] + guarded = emphasized.scatter(0, starts, guarded.index_select(0, starts)) + block_ids = torch.repeat_interleave(torch.arange(lengths.numel(), device=guarded.device), block_lengths) + local_samples = torch.arange(guarded.numel(), device=guarded.device) - block_offsets[block_ids] + guarded = guarded.masked_fill(local_samples >= lengths[block_ids], 0.0) + return guarded + + +def _high_frequency_mask(spectra, first_masked_bin): + bins = torch.arange(spectra.shape[1], device=spectra.device) + return bins.unsqueeze(0) >= first_masked_bin + + +def _normalize_packed_features(features, lengths, sequence_ids, normalize_type): + normalized, _ = _normalize_packed_features_and_padding(features, lengths, sequence_ids, normalize_type) + return normalized + + +def _normalize_packed_features_and_padding(features, lengths, sequence_ids, normalize_type, *, padding_value=None): + if normalize_type == "per_feature": + input_dtype = features.dtype + statistics_features = _packed_normalization_statistics_features(features) + denominator = lengths.clamp_min(1).unsqueeze(1) + reference = _packed_segment_reference(statistics_features, lengths) + mean = reference + _packed_segment_sum(statistics_features - reference[sequence_ids], lengths) / denominator + centered = statistics_features - mean[sequence_ids] + variance = _packed_segment_sum(centered.square(), lengths) / (denominator - 1) + std = torch.sqrt(variance).masked_fill(variance.isnan(), 0.0) + CONSTANT + normalized = (centered / std[sequence_ids]).to(input_dtype) + normalized_padding = features.new_zeros((lengths.numel(), features.shape[1])) + return normalized, normalized_padding + if normalize_type == "all_features": + input_dtype = features.dtype + statistics_features = _packed_normalization_statistics_features(features) + denominator = lengths * features.shape[1] + mean = _packed_segment_sum(statistics_features.sum(1), lengths) / denominator.clamp_min(1) + centered = statistics_features - mean[sequence_ids].unsqueeze(1) + variance = _packed_segment_sum(centered.square().sum(1), lengths) / (denominator.clamp_min(1) - 1) + std = torch.sqrt(variance).masked_fill(variance.isnan(), 0.0) + CONSTANT + normalized = (centered / std[sequence_ids].unsqueeze(1)).to(input_dtype) + padding = _expand_packed_padding(padding_value, statistics_features, lengths) + normalized_padding = ( + None if padding is None else ((padding - mean.unsqueeze(1)) / std.unsqueeze(1)).to(input_dtype) + ) + return normalized, normalized_padding + if "fixed_mean" in normalize_type and "fixed_std" in normalize_type: + mean = torch.as_tensor(normalize_type["fixed_mean"], device=features.device, dtype=features.dtype) + std = torch.as_tensor(normalize_type["fixed_std"], device=features.device, dtype=features.dtype) + if mean.numel() == features.shape[1]: + normalized = (features - mean) / std + padding = _expand_packed_padding(padding_value, features, lengths) + normalized_padding = None if padding is None else (padding - mean) / std + return normalized, normalized_padding + mean = mean.view(lengths.numel(), features.shape[1]) + std = std.view(lengths.numel(), features.shape[1]) + normalized = (features - mean[sequence_ids]) / std[sequence_ids] + padding = _expand_packed_padding(padding_value, features, lengths) + normalized_padding = None if padding is None else (padding - mean) / std + return normalized, normalized_padding + return features, padding_value + + +def _packed_segment_sum(values, lengths): + # Public packed entry points validate lengths; avoid repeating their synchronizing checks here. + return torch.segment_reduce(values, "sum", lengths=lengths, unsafe=True) + + +def _packed_segment_reference(values, lengths): + """Return one value per segment, with zeros for empty segments.""" + starts = torch.cat([lengths.new_zeros(1), lengths.cumsum(0)[:-1]]).long() + references = values.index_select(0, starts.clamp_max(values.shape[0] - 1)) + return references.masked_fill((lengths == 0).unsqueeze(1), 0.0) + + +def _packed_normalization_statistics_features(features): + """Accumulate packed normalization statistics safely for low-precision inputs.""" + if features.dtype in (torch.float16, torch.bfloat16): + return features.float() + return features + + +def _expand_packed_padding(padding_value, features, lengths): + if padding_value is None: + return None + if isinstance(padding_value, torch.Tensor): + return padding_value.to(features) + return features.new_full((lengths.numel(), features.shape[1]), padding_value) + + +def _make_packed_features(features, lengths, *, padding_value=0.0, padded_length=None, max_seqlen=None): + cu_seqlens = torch.cat([lengths.new_zeros(1, dtype=torch.int32), lengths.cumsum(0, dtype=torch.int32)]) + if max_seqlen is None: + max_seqlen = int(lengths.max().item()) if lengths.numel() else 0 + return _new_packed_encoder_activations( + features, lengths.to(torch.int64), cu_seqlens, max_seqlen, padding_value, padded_length + ) + + +def _empty_packed_features(x, lengths, feature_dim, *, padding_value=0.0, padded_length=None, max_seqlen=None): + return _make_packed_features( + x.new_empty((0, feature_dim)), + lengths, + padding_value=padding_value, + padded_length=padded_length, + max_seqlen=max_seqlen, + ) diff --git a/nemo/collections/asr/parts/submodules/multi_head_attention.py b/nemo/collections/asr/parts/submodules/multi_head_attention.py index 374e198b874a..2519031ae453 100644 --- a/nemo/collections/asr/parts/submodules/multi_head_attention.py +++ b/nemo/collections/asr/parts/submodules/multi_head_attention.py @@ -1232,7 +1232,7 @@ def __init__(self, d_k, rotary_fraction=1.0, rope_base=10000.0, max_len=5000): self.rope_base = rope_base self.max_len = max_len - inv_freq = 1.0 / (rope_base ** (torch.arange(0, d_k_rot, 2, dtype=torch.float32) / d_k_rot)) + inv_freq = 1.0 / (rope_base ** (torch.arange(0, d_k_rot, 2, dtype=torch.float32, device="cpu") / d_k_rot)) self.register_buffer('inv_freq', inv_freq, persistent=False) def _rotate_half(self, x): @@ -1256,11 +1256,18 @@ def create_pe(self, positions, dtype): ``dtype`` for storage. The final cast to Q/K runtime dtype happens in ``forward``. """ - freqs = torch.outer(positions, self.inv_freq.to(device=positions.device, dtype=torch.float32)) + # Build this non-persistent cache with one device-independent fp32 + # reference. CPU and CUDA transcendental kernels can round differently + # before the model-dtype cast, making identical weights produce different + # rotations solely because the model was constructed on another device. + target_device = positions.device + positions_cpu = positions.detach().to(device="cpu", dtype=torch.float32) + inv_freq_cpu = self.inv_freq.detach().to(device="cpu", dtype=torch.float32) + freqs = torch.outer(positions_cpu, inv_freq_cpu) # Duplicate to align with `_rotate_half`: tail half mirrors the head half. emb = torch.cat((freqs, freqs), dim=-1) - cos = emb.cos().to(dtype) - sin = emb.sin().to(dtype) + cos = emb.cos().to(device=target_device, dtype=dtype) + sin = emb.sin().to(device=target_device, dtype=dtype) if hasattr(self, 'cos'): self.cos = cos self.sin = sin diff --git a/nemo/collections/asr/parts/submodules/spectr_augment.py b/nemo/collections/asr/parts/submodules/spectr_augment.py index 388bd5e65472..b4dd1ffca35a 100644 --- a/nemo/collections/asr/parts/submodules/spectr_augment.py +++ b/nemo/collections/asr/parts/submodules/spectr_augment.py @@ -19,6 +19,7 @@ import torch import torch.nn as nn +from nemo.collections.asr.parts.packed_sequence import PackedEncoderActivations, packed_encoder_position_ids from nemo.core.classes import Typing, typecheck from nemo.core.neural_types import LengthsType, NeuralType, SpectrogramType @@ -100,6 +101,31 @@ def forward(self, input_spec, length): else: return self._forward_legacy(input_spec, length) + @torch.no_grad() + def forward_packed(self, input_spec: PackedEncoderActivations) -> PackedEncoderActivations: + """Apply GPU-vectorized sequence-local masks directly to token-flat features. + + ``use_vectorized_code`` only preserves the historical dense backend choice; + the new packed API always avoids the legacy Python loop. + """ + data = _apply_packed_axis_masks( + input_spec.data, + input_spec, + num_masks=self.time_masks, + width=self.time_width, + axis=self.TIME_AXIS, + mask_value=self.mask_value, + ) + data = _apply_packed_axis_masks( + data, + input_spec, + num_masks=self.freq_masks, + width=self.freq_width, + axis=self.FREQ_AXIS, + mask_value=self.mask_value, + ) + return input_spec.with_data(data) + def _forward_legacy(self, input_spec, length): batch_size, num_freq_bins, _ = input_spec.shape # Move lengths to CPU before repeated indexing @@ -262,3 +288,77 @@ def forward(self, input_spec): input_spec[idx, rect_x : rect_x + w_x, rect_y : rect_y + w_y] = 0.0 return input_spec + + @torch.no_grad() + def forward_packed(self, input_spec: PackedEncoderActivations) -> PackedEncoderActivations: + """Apply sequence-local rectangular masks to token-flat features.""" + if input_spec.total_tokens == 0 or self.rect_masks == 0: + return input_spec + batch_size = input_spec.batch_size + data = input_spec.data + freq_start = torch.zeros((batch_size, self.rect_masks), dtype=torch.long, device=data.device) + freq_width = torch.zeros_like(freq_start) + time_start = torch.zeros((batch_size, self.rect_masks), dtype=torch.long, device=data.device) + time_width = torch.zeros_like(time_start) + dense_time = input_spec.padded_length if input_spec.padded_length is not None else input_spec.max_seqlen + for row in range(batch_size): + for mask in range(self.rect_masks): + freq_start[row, mask] = self._rng.randint(0, data.shape[1] - self.rect_freq) + time_start[row, mask] = self._rng.randint(0, dense_time - self.rect_time) + freq_width[row, mask] = self._rng.randint(0, self.rect_freq) + time_width[row, mask] = self._rng.randint(0, self.rect_time) + data = _apply_packed_range_masks( + data, + input_spec, + freq_start=freq_start, + freq_width=freq_width, + time_start=time_start, + time_width=time_width, + mask_value=0.0, + rectangular=True, + ) + return input_spec.with_data(data) + + +def _apply_packed_axis_masks(data, packed, *, num_masks, width, axis, mask_value): + if num_masks == 0 or packed.total_tokens == 0: + return data + batch_size = packed.batch_size + axis_length = packed.max_seqlen if axis == SpecAugment.TIME_AXIS else data.shape[1] + if axis == SpecAugment.TIME_AXIS and isinstance(width, float): + width = torch.clamp(width * packed.lengths, max=axis_length).unsqueeze(1) + mask_width = (torch.rand((batch_size, num_masks), device=data.device, dtype=torch.float32) * width).long() + mask_start = torch.rand((batch_size, num_masks), device=data.device, dtype=torch.float32) + if axis == SpecAugment.TIME_AXIS: + mask_start = (mask_start * (packed.lengths.unsqueeze(1) - mask_width)).long() + token_indices = torch.arange(packed.total_tokens, device=data.device) + sequence_ids = torch.bucketize(token_indices, packed.cu_seqlens[1:], right=True) + positions = (token_indices - packed.cu_seqlens[sequence_ids]).unsqueeze(1) + mask = ((positions >= mask_start[sequence_ids]) & (positions < (mask_start + mask_width)[sequence_ids])).any(1) + return data.masked_fill(mask.unsqueeze(1), mask_value) + mask_start = (mask_start * (axis_length - mask_width)).long() + bins = torch.arange(axis_length, device=data.device) + mask = ((bins >= mask_start.unsqueeze(-1)) & (bins < (mask_start + mask_width).unsqueeze(-1))).any(1) + token_indices = torch.arange(packed.total_tokens, device=data.device) + sequence_ids = torch.bucketize(token_indices, packed.cu_seqlens[1:], right=True) + return data.masked_fill(mask[sequence_ids], mask_value) + + +def _apply_packed_range_masks( + data, packed, *, freq_start, freq_width, time_start, time_width, mask_value, rectangular=False +): + if packed.total_tokens == 0: + return data + sequence_ids = torch.repeat_interleave(torch.arange(packed.batch_size, device=data.device), packed.lengths) + positions = packed_encoder_position_ids(packed) + bins = torch.arange(data.shape[1], device=data.device) + time_mask = (positions[:, None] >= time_start[sequence_ids]) & ( + positions[:, None] < (time_start + time_width)[sequence_ids] + ) + freq_mask = (bins >= freq_start.unsqueeze(-1)) & (bins < (freq_start + freq_width).unsqueeze(-1)) + freq_mask = freq_mask[sequence_ids].transpose(1, 2) + if rectangular: + mask = (time_mask.unsqueeze(1) & freq_mask).any(2) + else: + mask = time_mask.any(1).unsqueeze(1) | freq_mask.any(2) + return data.masked_fill(mask, mask_value) diff --git a/nemo/collections/asr/parts/submodules/subsampling.py b/nemo/collections/asr/parts/submodules/subsampling.py index 3ce6950b981e..a98a8084e34a 100644 --- a/nemo/collections/asr/parts/submodules/subsampling.py +++ b/nemo/collections/asr/parts/submodules/subsampling.py @@ -19,6 +19,7 @@ import torch.nn as nn from torch.nn import LayerNorm +from nemo.collections.asr.parts.packed_sequence import PackedEncoderActivations, _new_packed_encoder_activations from nemo.collections.asr.parts.submodules.causal_convs import CausalConv1D, CausalConv2D from nemo.core.utils.optional_libs import TRITON_AVAILABLE, triton_required from nemo.utils import logging @@ -58,7 +59,7 @@ def get_streaming_cache_size(self): non-overlapping, so a chunk aligned to ``subsampling_factor`` needs none.""" return 0 - def forward(self, x, lengths): + def forward(self, x, lengths=None): """ Args: x: (B, C, T) input features. @@ -67,6 +68,12 @@ def forward(self, x, lengths): x: (B, T', feat_out) stacked and projected features. lengths: (B,) updated lengths after subsampling. """ + if isinstance(x, PackedEncoderActivations): + if lengths is not None: + raise ValueError("lengths must be omitted when x is PackedEncoderActivations.") + return self.forward_packed(x) + if lengths is None: + raise ValueError("lengths are required for padded FeatureStacking input.") x = x.transpose(1, 2) # (B, C, T) -> (B, T, C) b, t, c = x.size() pad_size = (self.subsampling_factor - (t % self.subsampling_factor)) % self.subsampling_factor @@ -78,6 +85,44 @@ def forward(self, x, lengths): lengths = self.compute_num_out_frames(lengths) return x, lengths + def forward_packed(self, packed: PackedEncoderActivations) -> PackedEncoderActivations: + """Stack and project token-flat feature sequences without batch padding.""" + stacked = self.stack_packed(packed) + return stacked.with_data(self.proj(stacked.data)) + + def stack_packed(self, packed: PackedEncoderActivations) -> PackedEncoderActivations: + """Stack token-flat frames, retaining packed metadata for grouped projection.""" + if packed.data.shape[1] * self.subsampling_factor != self.proj.in_features: + raise ValueError( + f"Expected packed feature width {self.proj.in_features // self.subsampling_factor}, " + f"got {packed.data.shape[1]}." + ) + output_lengths = self.compute_num_out_frames(packed.lengths) + output_cu_seqlens = torch.cat( + [output_lengths.new_zeros(1, dtype=torch.int32), output_lengths.cumsum(0, dtype=torch.int32)] + ) + slot_count = int(output_cu_seqlens[-1].item()) * self.subsampling_factor + slot_indices = torch.arange(slot_count, device=packed.data.device) + slot_boundaries = output_cu_seqlens[1:].to(torch.int64) * self.subsampling_factor + slot_sequence_ids = torch.bucketize(slot_indices, slot_boundaries, right=True) + if isinstance(packed.padding_value, torch.Tensor): + slots = packed.padding_value.to(packed.data)[slot_sequence_ids].clone() + else: + slots = packed.data.new_full((slot_count, packed.data.shape[1]), packed.padding_value) + if packed.padded_length is not None and slot_count: + slot_starts = output_cu_seqlens[slot_sequence_ids].to(torch.int64) * self.subsampling_factor + local_slots = slot_indices - slot_starts + slots.masked_fill_(local_slots.unsqueeze(1) >= packed.padded_length, 0.0) + if packed.total_tokens: + token_indices = torch.arange(packed.total_tokens, device=packed.data.device) + sequence_ids = torch.bucketize(token_indices, packed.cu_seqlens[1:], right=True) + local_positions = token_indices - packed.cu_seqlens[sequence_ids] + slot_positions = output_cu_seqlens[sequence_ids].to(torch.int64) * self.subsampling_factor + slots[slot_positions + local_positions] = packed.data + stacked = slots.reshape(-1, self.proj.in_features) + max_seqlen = self.compute_num_out_frames(packed.max_seqlen) + return _new_packed_encoder_activations(stacked, output_lengths, output_cu_seqlens, max_seqlen) + class StackingSubsampling(torch.nn.Module): """Stacking subsampling which simply stacks consecutive frames to reduce the sampling rate diff --git a/nemo/collections/speechlm2/data/salm_dataset.py b/nemo/collections/speechlm2/data/salm_dataset.py index 7a95a35f57a2..e4dfef15d73e 100644 --- a/nemo/collections/speechlm2/data/salm_dataset.py +++ b/nemo/collections/speechlm2/data/salm_dataset.py @@ -37,8 +37,10 @@ from nemo.collections.common.data.lhotse import NeMoMultimodalConversation from nemo.collections.common.data.lhotse.text_adapters import ( AudioTurn, + Formattable, TextTurn, collate_conversation_audio_fault_tolerant, + collate_conversation_audio_packed_fault_tolerant, ) from nemo.collections.common.data.prompt_fn import registered_prompt_format_fn from nemo.collections.common.prompts import Llama2PromptFormatter @@ -67,11 +69,29 @@ class SALMDataset(torch.utils.data.Dataset): ``spk_targets`` / ``spk_target_length``. Rows without an explicit RTTM path contain the reserved value ``-1`` so the perception encoder can replace them with inferred speaker activity. + pack_audio (bool): + Return valid waveform samples contiguously as `packed_audio_samples` + plus `audio_cu_seqlens`, instead of materializing `audios[B, T_max]`. + Defaults to `False` for complete batch-API compatibility. + pack_sequences (bool): + Return every variable-length sequence tensor without batch padding. + Text IDs and masks are concatenated to shape ``[T_total]`` and + described by ``text_cu_seqlens``; audio is returned in the same + packed form as ``pack_audio=True``; optional speaker targets are + concatenated to shape ``[T_spk_total, N_spk]`` and described by + ``spk_target_cu_seqlens``. This option implies ``pack_audio=True``. + batch_tokens (int | None): + Token budget used by the Lhotse sampler. When provided, and the + sampler attached an exact ``num_tokens`` measurement to every + retained conversation, the batch contains a scalar + ``packing_efficiency`` equal to the measured token sum divided by + this budget. strict_audio_loading (bool): - Re-raise audio collation failures and reject batches where the - fault-tolerant collator dropped or reordered conversations/audio. - Defaults to ``False``; the datamodule configures it from - ``fault_tolerant_audio_loading`` for each loader. + Re-raises audio collation errors and + rejects conversations or audio items dropped or reordered by the + fault-tolerant collator. Defaults to ``False`` so audio I/O and + decoder failures remain fault tolerant. The datamodule controls + this through ``fault_tolerant_audio_loading``. [ SOT Example for overlapping speakers ] Speaker-parallel transcription as a timeline: @@ -84,11 +104,17 @@ class SALMDataset(torch.utils.data.Dataset): Returns: A dictionary with the following keys: - - audios: Tensor of audio waveform samples [B_audio, T_samples] + - audios: Tensor of audio waveform samples [B_audio, T_samples] (default mode) + - packed_audio_samples: Tensor of contiguous waveform samples [T_total] (packed mode) + - audio_cu_seqlens: Tensor of cumulative waveform offsets [B_audio + 1] (packed mode) - audio_lens: Tensor of audio lengths [B_audio] - - input_ids: Tensor of text token IDs [B, T_tokens], including audio_locator_tag tokens - - loss_mask: Boolean tensor [B, T_tokens] indicating which tokens are part of the - assistant's responses (True) and should be used for computing loss + - input_ids: Tensor of text token IDs [B, T_tokens] (padded mode) or + [T_total] (packed mode), including audio_locator_tag tokens + - loss_mask: Boolean tensor with the same shape as input_ids indicating which + tokens are part of the assistant's responses (True) and should be used for loss + - text_cu_seqlens: Tensor of cumulative text offsets [B + 1] (packed mode) + - packing_efficiency: Optional scalar measuring sampled tokens divided by + ``batch_tokens`` Notes: - Each audio_locator_tag token in input_ids corresponds to an audio segment in audios @@ -106,11 +132,19 @@ def __init__( self, tokenizer: AutoTokenizer, multispeaker_cfg: dict | None = None, + pack_audio: bool = False, + pack_sequences: bool = False, + batch_tokens: int | None = None, strict_audio_loading: bool = False, ) -> None: self.tokenizer = tokenizer self.pad_id = get_pad_id(tokenizer) + self.pack_sequences = bool(pack_sequences) + self.pack_audio = bool(pack_audio) or self.pack_sequences + self.batch_tokens = int(batch_tokens) if batch_tokens is not None else None self.strict_audio_loading = bool(strict_audio_loading) + if self.batch_tokens is not None and self.batch_tokens <= 0: + raise ValueError(f"batch_tokens must be positive, got {self.batch_tokens}") # Setting USE_AIS_GET_BATCH=true makes the loader issue a single AIStore GetBatch # call per minibatch, paired with URL-backed cuts produced by the multimodal # conversation adapters (NeMoMultimodalConversation{Jsonl,ShareGPTJsonl}Adapter). @@ -126,11 +160,20 @@ def __init__( ) self.multispeaker_cfg = MultiSpeakerConfig.from_dict(multispeaker_cfg) self.multispeaker_processor = ( - SALMMultiSpeakerProcessor(self.multispeaker_cfg) if self.multispeaker_cfg is not None else None + SALMMultiSpeakerProcessor(self.multispeaker_cfg, pack_targets=self.pack_sequences) + if self.multispeaker_cfg is not None + else None ) def with_fault_tolerant_audio_loading(self, enabled: bool) -> "SALMDataset": - """Return a per-loader view with the requested audio I/O policy.""" + """Return a per-loader view with the requested audio I/O failure policy. + + ``DataModule`` shares one dataset factory between train/validation/test, + while each loader may have a different audio-loading policy. A shallow + copy keeps the tokenizer and model-independent processors shared, and a + copied ``AudioSamples`` instance avoids mutating another loader's + strictness state. + """ enabled = bool(enabled) dataset = copy(self) dataset.strict_audio_loading = not enabled @@ -142,15 +185,29 @@ def with_fault_tolerant_audio_loading(self, enabled: bool) -> "SALMDataset": return dataset def __getitem__(self, conversations: CutSet) -> dict | None: + # The collator retains its fault-tolerant 3-tuple API, but strict mode + # verifies exact conversation/audio identity and raises on any drop. + # DataModule gates FallbackDataset behind the same audio-loading policy. if self.strict_audio_loading: requested_conversation_ids = tuple(id(conversation) for conversation in conversations) requested_audio_cut_ids = _audio_cut_ids(conversations) else: requested_conversation_ids = requested_audio_cut_ids = () + try: - audios, audio_lens, conversations = collate_conversation_audio_fault_tolerant( - conversations, self.load_audio - ) + if self.pack_audio: + packed_audio_samples, audio_cu_seqlens, audio_lens, conversations = ( + collate_conversation_audio_packed_fault_tolerant(conversations, self.load_audio) + ) + audio_inputs = { + "packed_audio_samples": packed_audio_samples, + "audio_cu_seqlens": audio_cu_seqlens, + } + else: + audios, audio_lens, conversations = collate_conversation_audio_fault_tolerant( + conversations, self.load_audio + ) + audio_inputs = {"audios": audios} except Exception as e: if self.strict_audio_loading: raise @@ -170,23 +227,55 @@ def __getitem__(self, conversations: CutSet) -> dict | None: ): raise RuntimeError( "Strict SALM validation dropped or reordered audio items: " - f"requested={len(requested_audio_cut_ids)} materialized={len(audio_lens)}" + f"requested={len(requested_audio_cut_ids)} " + f"materialized={len(audio_lens)}" ) if not conversations: if self.strict_audio_loading: - raise RuntimeError("Strict SALM validation dropped or reordered conversations into an empty batch.") + raise RuntimeError("Strict SALM loading produced an empty conversation batch.") return None + input_ids = [c.input_ids for c in conversations] + loss_masks = [getattr(c, "mask", torch.empty(0)) for c in conversations] + if self.pack_sequences: + packed_input_ids, text_cu_seqlens = pack_vectors(input_ids) + packed_loss_mask, loss_mask_cu_seqlens = pack_vectors(loss_masks) + if not torch.equal(text_cu_seqlens, loss_mask_cu_seqlens): + raise ValueError( + "Each SALM loss mask must have the same length as its input IDs; " + f"got offsets {text_cu_seqlens.tolist()} and {loss_mask_cu_seqlens.tolist()}." + ) + text_inputs = { + "input_ids": packed_input_ids, + "loss_mask": packed_loss_mask.to(torch.bool), + "text_cu_seqlens": text_cu_seqlens, + } + else: + text_inputs = { + "input_ids": left_collate_vectors(input_ids, padding_value=self.pad_id), + "loss_mask": left_collate_vectors(loss_masks, padding_value=0).to(torch.bool), + } + batch = { - "audios": audios, + **audio_inputs, + **text_inputs, "audio_lens": audio_lens, - "input_ids": left_collate_vectors([c.input_ids for c in conversations], padding_value=self.pad_id), - "loss_mask": left_collate_vectors( - [getattr(c, "mask", torch.empty(0)) for c in conversations], padding_value=0 - ).to(torch.bool), - "conversations": drop_in_memory_data(conversations), + # Keep decoded in-memory audio available until auxiliary targets + # are materialized. Native ShareGPT WDS cuts intentionally use + # memory-backed recordings; dropping them first replaces their + # sources with unresolved Shar placeholders, and multichannel + # downmixing in SALMMultiSpeakerProcessor then cannot load audio. + "conversations": conversations, } + if self.batch_tokens is not None: + sampled_lengths = [getattr(conversation, "num_tokens", None) for conversation in conversations] + if all(length is not None for length in sampled_lengths): + batch["packing_efficiency"] = torch.tensor( + sum(sampled_lengths) / self.batch_tokens, + dtype=torch.float32, + ) if self.multispeaker_processor is not None: self.multispeaker_processor(batch) + batch["conversations"] = drop_in_memory_data(conversations) return batch @@ -199,8 +288,23 @@ def left_collate_vectors( return pad_sequence(tensors, batch_first=True, padding_value=padding_value, padding_side="left") +def pack_vectors(tensors: Iterable[Union[torch.Tensor, np.ndarray]]) -> tuple[torch.Tensor, torch.Tensor]: + """Concatenate 1-D rows and return their cumulative offsets without padding.""" + tensors = [torch.as_tensor(t) for t in tensors] + if not tensors: + raise ValueError("Cannot pack an empty sequence collection.") + if not all(t.ndim == 1 for t in tensors): + raise ValueError(f"Expected only 1-D input tensors, got shapes {[tuple(t.shape) for t in tensors]}.") + values = torch.cat(tensors, dim=0) + lengths = torch.as_tensor([t.shape[0] for t in tensors], dtype=torch.long, device=values.device) + cu_seqlens = torch.cat([lengths.new_zeros(1), lengths.cumsum(0)]) + return values, cu_seqlens + + def drop_in_memory_data(conversations: CutSet) -> CutSet: - def _drop(conversation: NeMoMultimodalConversation) -> NeMoMultimodalConversation: + def _drop(conversation: Formattable) -> Formattable: + if not isinstance(conversation, NeMoMultimodalConversation): + return conversation turns = [] for t in conversation.turns: if isinstance(t, AudioTurn): @@ -211,8 +315,14 @@ def _drop(conversation: NeMoMultimodalConversation) -> NeMoMultimodalConversatio return conversations.map(_drop, apply_fn=None) -def _audio_cut_ids(conversations: Iterable[NeMoMultimodalConversation]) -> tuple[str, ...]: - return tuple(cut.id for conversation in conversations for cut in conversation.list_cuts()) +def _audio_cut_ids(conversations: Iterable[Formattable]) -> tuple[str, ...]: + cut_ids = [] + for conversation in conversations: + if isinstance(conversation, NeMoMultimodalConversation): + cut_ids.extend(cut.id for cut in conversation.list_cuts()) + elif not isinstance(conversation, Formattable): + raise TypeError("SALMDataset expected a prompt-formatted example, " f"got {type(conversation).__name__}.") + return tuple(cut_ids) @registered_prompt_format_fn(NeMoMultimodalConversation, Llama2PromptFormatter) @@ -275,8 +385,9 @@ def from_dict(cfg: dict | None) -> "MultiSpeakerConfig | None": class SALMMultiSpeakerProcessor: """Add SOT activity targets, using ``-1`` rows to request inferred diarization.""" - def __init__(self, cfg: MultiSpeakerConfig) -> None: + def __init__(self, cfg: MultiSpeakerConfig, *, pack_targets: bool = False) -> None: self.cfg = cfg + self.pack_targets = bool(pack_targets) def __call__(self, batch: dict) -> None: """Attach RTTM targets or missing-RTTM sentinels to ``batch`` in place.""" @@ -293,23 +404,50 @@ def __call__(self, batch: dict) -> None: [bool(torch.all(activity == -1.0)) for activity in speaker_activities], dtype=torch.bool, ) - targets, target_length = collate_speaker_activity_targets( - speaker_activities, - batch["audio_lens"], - num_speakers=cfg.num_speakers, - num_sample_per_mel_frame=cfg.num_sample_per_mel_frame, - num_mel_frame_per_target_frame=cfg.num_mel_frame_per_target_frame, - dtype=batch["audios"].dtype, - ) - targets[missing_rttm_rows] = -1.0 - batch["spk_targets"] = targets - batch["spk_target_length"] = target_length + dtype = (batch["audios"] if "audios" in batch else batch["packed_audio_samples"]).dtype + if self.pack_targets: + normalized = [] + for activity, missing_rttm in zip(speaker_activities, missing_rttm_rows): + n_spk = activity.shape[1] + if n_spk > cfg.num_speakers: + activity = activity[:, : cfg.num_speakers] + elif n_spk < cfg.num_speakers: + activity = torch.nn.functional.pad( + activity, + (0, cfg.num_speakers - n_spk), + mode="constant", + value=0.0, + ) + activity = activity.to(dtype=dtype) + if missing_rttm: + activity = torch.full_like(activity, -1.0) + normalized.append(activity) + target_length = torch.as_tensor([target.shape[0] for target in normalized], dtype=torch.long) + targets = torch.cat(normalized, dim=0) + target_cu_seqlens = torch.cat([target_length.new_zeros(1), target_length.cumsum(0)]) + batch["spk_targets"] = targets + batch["spk_target_length"] = target_length + batch["spk_target_cu_seqlens"] = target_cu_seqlens + else: + targets, target_length = collate_speaker_activity_targets( + speaker_activities, + batch["audio_lens"], + num_speakers=cfg.num_speakers, + num_sample_per_mel_frame=cfg.num_sample_per_mel_frame, + num_mel_frame_per_target_frame=cfg.num_mel_frame_per_target_frame, + dtype=dtype, + ) + targets[missing_rttm_rows] = -1.0 + batch["spk_targets"] = targets + batch["spk_target_length"] = target_length def _build_speaker_activities(self, conversations: CutSet) -> list[torch.Tensor]: cfg = self.cfg speaker_activities = [] for conversation in conversations: - for turn in conversation.turns: + # Generic prompt-formatted text-only examples intentionally do not + # expose multimodal turns and have no speaker targets to materialize. + for turn in getattr(conversation, "turns", ()): if not isinstance(turn, AudioTurn): continue diff --git a/nemo/collections/speechlm2/models/salm_automodel.py b/nemo/collections/speechlm2/models/salm_automodel.py index 81a56d9cd366..a7f97a22ef8f 100644 --- a/nemo/collections/speechlm2/models/salm_automodel.py +++ b/nemo/collections/speechlm2/models/salm_automodel.py @@ -15,7 +15,7 @@ import re import warnings from collections import defaultdict -from contextlib import contextmanager +from contextlib import contextmanager, nullcontext from typing import Any import torch @@ -23,7 +23,7 @@ from lightning import LightningModule from omegaconf import DictConfig, OmegaConf from torch import Tensor -from torch.distributed.fsdp import fully_shard +from torch.distributed.fsdp import fully_shard, register_fsdp_forward_method from torch.distributed.tensor import DTensor from torch.distributed.tensor.parallel import loss_parallel from transformers import GenerationConfig @@ -34,6 +34,7 @@ from nemo.collections.speechlm2.models.salm import _resolve_audios_in_prompt, replace_placeholders_and_build_targets from nemo.collections.speechlm2.parts.automodel_lora import ensure_lora_trainable, make_peft_config, maybe_install_lora from nemo.collections.speechlm2.parts.encoder_chunking import encode_audio_with_optional_chunking +from nemo.collections.speechlm2.parts.gc import GarbageCollectionManager from nemo.collections.speechlm2.parts.hf_hub import HFHubMixin from nemo.collections.speechlm2.parts.mtp import ( build_mtp_loss_fn, @@ -82,6 +83,22 @@ def __init__(self, cfg) -> None: self._use_fsdp = False self._use_tp = False + self._garbage_collection = GarbageCollectionManager(self.cfg.get("gc_every_steps", None)) + self._fused_linear_cross_entropy = None + cross_entropy_backend = str(self.cfg.get("cross_entropy_backend", "eager")) + if cross_entropy_backend not in ("eager", "fused_linear"): + raise ValueError( + "model.cross_entropy_backend must be 'eager' or 'fused_linear', " f"got {cross_entropy_backend!r}." + ) + if cross_entropy_backend == "fused_linear": + if self.lss_loss is not None: + raise ValueError( + "model.cross_entropy_backend='fused_linear' is incompatible with model.lss_loss because " + "the training path deliberately does not materialize full logits." + ) + from nemo_automodel.components.loss.linear_ce import FusedLinearCrossEntropy + + self._fused_linear_cross_entropy = FusedLinearCrossEntropy(ignore_index=-100, reduction="sum") if self.cfg.get("init_configure_model", False): self.configure_model() @@ -218,21 +235,41 @@ def forward( llm_input_ids = torch.zeros((1, seq_len), device=input_embeds.device, dtype=torch.long) mtp_embed_inputs = tuple(llm_kwargs.pop("mtp_embed_inputs", ())) - out = self.llm( - llm_input_ids, - *mtp_embed_inputs, - inputs_embeds=input_embeds, - attention_mask=attention_mask, - past_key_values=cache, - use_cache=cache is not None, - return_dict=True, - **llm_kwargs, + llm_positional_args = (llm_input_ids, *mtp_embed_inputs) if mtp_embed_inputs else () + if not mtp_embed_inputs: + llm_kwargs["input_ids"] = llm_input_ids + use_fused_linear_ce = ( + self.training and getattr(self, "_fused_linear_cross_entropy", None) is not None and cache is None ) + if use_fused_linear_ce: + llm_kwargs["output_hidden_states"] = True + llm_kwargs["compute_logits"] = False + + backend = getattr(self.llm, "backend", None) + te_fp8 = getattr(backend, "te_fp8", None) + fp8_ctx = te_fp8.maybe_te_autocast() if te_fp8 is not None else nullcontext() + with fp8_ctx: + out = self.llm( + *llm_positional_args, + inputs_embeds=input_embeds, + attention_mask=attention_mask, + past_key_values=cache, + use_cache=cache is not None, + return_dict=True, + **llm_kwargs, + ) if not isinstance(out, dict): # NeMo Automodel doesn't respect return_dict=True yet ans = {"logits": out} else: ans = {"logits": out['logits']} # (B, T, text_vocab_size) + if use_fused_linear_ce: + hidden_states = out.get("hidden_states", None) + if hidden_states is None: + raise RuntimeError("Fused linear CE requires the LLM to return final hidden states.") + if isinstance(hidden_states, (list, tuple)): + hidden_states = hidden_states[-1] + ans["hidden_states"] = hidden_states if cache is not None: ans["cache"] = out["past_key_values"] # MTP per-depth hidden states are returned when an MTP head is attached and @@ -307,8 +344,16 @@ def prepare_inputs(self, batch: dict, *, include_mtp_inputs: bool = True): device_mesh = getattr(self, "_device_mesh", None) spk_targets = batch.get("spk_targets", None) spk_target_lengths = batch.get("spk_target_length", None) + spk_target_cu_seqlens = batch.get("spk_target_cu_seqlens", None) cp_mesh, _, _ = get_cp_mesh(device_mesh) fsdp_sync_group = get_perception_fsdp_group(device_mesh) + packed_encoder_sequences = bool(self.cfg.get("packed_encoder_sequences", False)) + packed_encoder_cp = bool(self.cfg.get("packed_encoder_cp", False)) + audio_lens = batch["audio_lens"] + audio_cu_seqlens = batch.get("audio_cu_seqlens") + audios = batch.get("audios") + if audios is None: + audios = batch["packed_audio_samples"] # Source audio encoding. Input audio: (B, T_samples), audio embeddings: (B, T, H). # Routing uses valid targets for RTTM rows, a -1 sentinel for non-RTTM @@ -321,19 +366,32 @@ def prepare_inputs(self, batch: dict, *, include_mtp_inputs: bool = True): uses_parallel_expert_encoder = self._uses_parallel_expert_encoder() audio_embs, dummy_audio_loss = encode_audio_with_cp_distribution( self.perception, - batch["audios"], - batch["audio_lens"], - chunk_size_seconds=self.cfg.get("encoder_chunk_size_seconds", None), - chunk_batch_size=self.cfg.get("encoder_chunk_batch_size", None), + audios, + audio_lens, + audio_cu_seqlens=audio_cu_seqlens, + # A ParallelExpertEncoder applies this shared setting to both packed + # post-stacking branches. Do not split its waveform a second time on + # that path. Dense PEE execution retains the ordinary outer chunker. + chunk_size_seconds=( + None + if uses_parallel_expert_encoder and packed_encoder_sequences + else self.cfg.get("encoder_chunk_size_seconds", None) + ), + chunk_batch_size=( + None + if uses_parallel_expert_encoder and packed_encoder_sequences + else self.cfg.get("encoder_chunk_batch_size", None) + ), sampling_rate=self.sampling_rate, cp_mesh=cp_mesh, spk_targets=spk_targets if uses_parallel_expert_encoder else None, spk_target_lengths=spk_target_lengths if uses_parallel_expert_encoder else None, + spk_target_cu_seqlens=spk_target_cu_seqlens if uses_parallel_expert_encoder else None, fsdp_sync_group=fsdp_sync_group, return_dummy_loss=True, + sequence_packed=packed_encoder_sequences, + packed_cp_gather=packed_encoder_cp, ) - input_ids_to_embed = torch.where(batch["input_ids"] == self.audio_locator_tag_id, 0, batch["input_ids"]) - text_embs = self._embed_tokens(input_ids_to_embed) target_ids_full = batch["input_ids"].where(batch["loss_mask"], -100) # CrossEntropyLoss().ignore_index # Packed-sequence (THD) path — used for both training and validation when enabled. @@ -341,20 +399,27 @@ def prepare_inputs(self, batch: dict, *, include_mtp_inputs: bool = True): if self.cfg.get("packed_sequences", False): from nemo.collections.speechlm2.parts.packed_sequences import prepare_packed_llm_inputs + te_fp8 = getattr(getattr(self.llm, "backend", None), "te_fp8", None) + ans = prepare_packed_llm_inputs( input_ids=batch["input_ids"], - text_embs=text_embs, + text_embs=None, audio_embs=audio_embs, target_ids=target_ids_full, padding_id=self.text_pad_id, placeholder_id=self.audio_locator_tag_id, device_mesh=device_mesh, mtp_num_depths=self._mtp_num_depths if include_mtp_inputs else 0, + embed_tokens=self._embed_tokens, + text_cu_seqlens=batch.get("text_cu_seqlens"), + token_alignment=8 if te_fp8 is not None else 1, ) if dummy_audio_loss is not None: ans["dummy_audio_loss"] = dummy_audio_loss return ans + input_ids_to_embed = torch.where(batch["input_ids"] == self.audio_locator_tag_id, 0, batch["input_ids"]) + text_embs = self._embed_tokens(input_ids_to_embed) input_embs, target_ids, attention_mask = replace_placeholders_and_build_targets( input_ids=batch["input_ids"], embeds=text_embs, @@ -393,6 +458,12 @@ def on_fit_start(self) -> None: averaging (see ``_configure_moe_aux_loss_scaler``).""" self._validate_parallelism_compatibility() self._configure_moe_aux_loss_scaler() + self._garbage_collection.on_fit_start() + + def optimizer_step(self, epoch, batch_idx, optimizer, optimizer_closure=None) -> None: + """Run configured manual GC after each completed optimizer step.""" + super().optimizer_step(epoch, batch_idx, optimizer, optimizer_closure) + self._garbage_collection.on_optimizer_step() def on_validation_start(self) -> None: """Reject unsupported parallel layouts for fit and standalone validation.""" @@ -450,6 +521,46 @@ def training_step(self, dataloader_iter): batch, batch_idx = read_batch(dataloader_iter, self) return self._training_step_batch(batch, batch_idx) + def _compute_training_cross_entropy_sum( + self, + forward_outputs: dict[str, Tensor], + target_ids: Tensor, + dp_group, + *, + lm_weight: Tensor | None = None, + ) -> tuple[Tensor, Tensor | None]: + """Return local summed CE and optional full logits used by auxiliary losses. + + ``lm_weight`` may be a previously materialized regular tensor shared + with the MTP loss. Supplying it avoids a second FSDP DTensor gather. + """ + fused_linear_cross_entropy = getattr(self, "_fused_linear_cross_entropy", None) + if fused_linear_cross_entropy is not None: + hidden_states = forward_outputs.get("hidden_states", None) + if hidden_states is None: + raise RuntimeError("Fused linear CE requires final hidden states from forward().") + if lm_weight is None: + lm_head = self.llm.get_output_embeddings() if hasattr(self.llm, "get_output_embeddings") else None + if lm_head is None: + lm_head = self.llm.lm_head + lm_weight = lm_head.weight + loss_sum = fused_linear_cross_entropy( + hidden_states, + target_ids, + lm_weight, + grad_reduce_group=dp_group, + ) + return loss_sum, None + + logits = forward_outputs["logits"] + loss_sum = torch.nn.functional.cross_entropy( + logits.reshape(-1, logits.size(-1)), + target_ids.reshape(-1), + reduction="sum", + ignore_index=-100, + ) + return loss_sum, logits + def _training_step_batch(self, batch: dict | None, batch_idx: int): self._current_batch_idx = batch_idx for m in (self.perception.preprocessor, self.perception.encoder, self.llm): @@ -483,13 +594,25 @@ def _training_step_batch(self, batch: dict | None, batch_idx: int): num_frames_global = num_frames num_frames_global = num_frames_global.clamp(min=1) + # The main and MTP fused losses both consume the full LM-head weight + # outside the owning FSDP module. Gather it once and share the regular + # tensor so their gradients accumulate into one reduce-scatter graph. + mtp_h = forward_outputs.get("mtp_per_depth_h", None) + shared_lm_weight = None + main_materialize = getattr(getattr(self, "_fused_linear_cross_entropy", None), "materialize_lm_weight", None) + mtp_materialize = getattr(getattr(self, "_mtp_loss_fn", None), "materialize_lm_weight", None) + if mtp_h is not None and callable(main_materialize) and callable(mtp_materialize): + lm_head = self.llm.get_output_embeddings() if hasattr(self.llm, "get_output_embeddings") else None + if lm_head is None: + lm_head = self.llm.lm_head + shared_lm_weight = main_materialize(lm_head.weight, grad_reduce_group=dp_group) + with loss_parallel(): - logits = forward_outputs["logits"] - loss_sum = torch.nn.functional.cross_entropy( - logits.reshape(-1, logits.size(-1)), # BSHD (B,T,V) or THD (1,T,V) -> (*, V) - inputs["target_ids"].reshape(-1), # BSHD (B,T) or THD (T,) -> (*,) - reduction="sum", - ignore_index=-100, + loss_sum, logits = self._compute_training_cross_entropy_sum( + forward_outputs, + inputs["target_ids"], + dp_group, + lm_weight=shared_lm_weight, ) loss = loss_sum * dp_size / num_frames_global if (dummy_audio_loss := inputs.get("dummy_audio_loss")) is not None: @@ -512,7 +635,6 @@ def _training_step_batch(self, batch: dict | None, batch_idx: int): # lm_head + CE work. ``mtp_loss`` keeps the same meaning as before: the weighted # auxiliary loss added to the training objective after the DP-size correction. mtp_metrics = {} - mtp_h = forward_outputs.get("mtp_per_depth_h", None) if mtp_h is not None: # Under packed THD multiple utterances share one token stream, so the # per-depth label roll must not predict the next sequence's first token @@ -533,6 +655,7 @@ def _training_step_batch(self, batch: dict | None, batch_idx: int): scaling_factor=self._mtp_loss_scaling_factor, num_label_tokens=num_frames_global, grad_reduce_group=dp_group, + lm_weight=shared_lm_weight, cu_seqlens=mtp_cu_seqlens, return_per_depth=True, ) @@ -570,6 +693,8 @@ def _training_step_batch(self, batch: dict | None, batch_idx: int): self.log("mtp_loss", mtp_metrics.pop("mtp_loss"), on_step=True, prog_bar=True, batch_size=B) self.log_dict(mtp_metrics, on_step=True, batch_size=B) self.log_dict({k: v for k, v in ans.items() if k != "loss"}, on_step=True, batch_size=B) + if (packing_efficiency := batch.get("packing_efficiency")) is not None: + self.log("packing_efficiency", packing_efficiency, on_step=True, batch_size=B) self.maybe_log_moe_metrics(batch_idx) return ans @@ -581,7 +706,19 @@ def _build_empty_training_batch(self) -> dict: if token_id is None: token_id = self.text_pad_id device = self.device - input_ids = torch.full((1, 2), int(token_id), dtype=torch.long, device=device) + packed_sequences = bool(self.cfg.get("packed_sequences", False)) + input_shape = (2,) if packed_sequences else (1, 2) + input_ids = torch.full(input_shape, int(token_id), dtype=torch.long, device=device) + if packed_sequences: + return { + "packed_audio_samples": torch.empty(0, dtype=torch.float32, device=device), + "audio_cu_seqlens": torch.zeros(1, dtype=torch.long, device=device), + "audio_lens": torch.empty(0, dtype=torch.long, device=device), + "input_ids": input_ids, + "loss_mask": torch.zeros_like(input_ids, dtype=torch.bool), + "text_cu_seqlens": torch.tensor([0, 2], dtype=torch.long, device=device), + "conversations": [], + } return { "audios": torch.empty(0, dtype=torch.float32, device=device), "audio_lens": torch.empty(0, dtype=torch.long, device=device), @@ -591,7 +728,10 @@ def _build_empty_training_batch(self) -> dict: } def _log_training_batch_debug(self, batch: dict | None, batch_idx: int) -> None: - max_logged = int(self.cfg.get("debug_log_training_batches", 2) or 0) + cfg = getattr(self, "cfg", None) + if cfg is None: + return + max_logged = int(cfg.get("debug_log_training_batches", 2) or 0) logged = getattr(self, "_debug_logged_training_batches", 0) if logged >= max_logged: return @@ -633,6 +773,7 @@ def shape_of(key: str): "training_batch_debug " f"rank={rank} batch_idx={batch_idx} " f"input_ids_shape={shape_of('input_ids')} audios_shape={shape_of('audios')} " + f"packed_audio_samples_shape={shape_of('packed_audio_samples')} " f"audio_lens_min={audio_lens_min} audio_lens_max={audio_lens_max} " f"audio_sec_max={audio_sec_max:.2f} nonpad_tokens={nonpad_tokens} loss_tokens={loss_tokens} " f"spk_targets_shape={shape_of('spk_targets')} " @@ -817,6 +958,9 @@ def test_step(self, *args: Any, **kwargs: Any): def backward(self, *args, **kwargs): self._setup_moe_fsdp_sync() + # Transformer Engine FP8 autocast is a forward-only context. Backward + # precision and scaling state come from the recorded forward graph; a + # fresh context here would update global amax bookkeeping twice. with loss_parallel(): super().backward(*args, **kwargs) @@ -1389,7 +1533,7 @@ def configure_model( if fsdp_mesh.size() > 1: self._use_fsdp = True - self.perception = fully_shard(self.perception, mesh=fsdp_mesh) + self.perception = _fully_shard_perception(self.perception, fsdp_mesh) # Enable MoE FSDP gradient accumulation optimization. # The MoEFSDPSyncMixin on the LLM defers gradient sync/resharding on @@ -1431,3 +1575,10 @@ def oomptimizer_schema(self) -> dict: {"name": "loss_mask", "type": NeuralType(("B", "T"), MaskType()), "seq_length": "output"}, ], } + + +def _fully_shard_perception(perception, mesh): + """FSDP2-shard perception and register its packed custom root forward.""" + perception = fully_shard(perception, mesh=mesh) + register_fsdp_forward_method(perception, "forward_sequence_packed") + return perception diff --git a/nemo/collections/speechlm2/modules/__init__.py b/nemo/collections/speechlm2/modules/__init__.py index d1a67990d7da..ea9cee999362 100644 --- a/nemo/collections/speechlm2/modules/__init__.py +++ b/nemo/collections/speechlm2/modules/__init__.py @@ -12,10 +12,11 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -from .perception import AudioPerceptionModule +from .perception import AudioPerceptionModule, IndependentDualEncoder from .speech_generation import TransformerARSpeechDecoder __all__ = [ 'AudioPerceptionModule', + 'IndependentDualEncoder', 'TransformerARSpeechDecoder', ] diff --git a/nemo/collections/speechlm2/modules/perception.py b/nemo/collections/speechlm2/modules/perception.py index 084395ec2039..6943ad4a026e 100644 --- a/nemo/collections/speechlm2/modules/perception.py +++ b/nemo/collections/speechlm2/modules/perception.py @@ -12,6 +12,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import contextlib import inspect import torch @@ -23,6 +24,11 @@ from nemo.collections.asr.models import ASRModel from nemo.collections.asr.modules.conformer_encoder import ConformerMultiLayerFeatureExtractor from nemo.collections.asr.parts.mixins import TranscribeConfig +from nemo.collections.asr.parts.packed_sequence import ( + PackedEncoderActivations, + pack_encoder_output, + unpack_encoder_output, +) from nemo.core import Exportable, NeuralModule, typecheck @@ -63,6 +69,17 @@ def encoder(self) -> nn.Module: return self._modules['encoder_multilayer'].encoder return self._modules['encoder'] + @property + def supports_sequence_packed_output(self) -> bool: + """Whether this exact encoder/adapter stack can preserve native THD activations.""" + return ( + 'encoder_multilayer' not in self._modules + and isinstance(self.modality_adapter, IdentityConnector) + and self.rote is None + and bool(getattr(self.encoder, 'supports_sequence_packed_output', False)) + and callable(getattr(self.encoder, 'forward_sequence_packed', None)) + ) + def __init__(self, cfg: DictConfig): super().__init__() # Initialize components @@ -107,6 +124,7 @@ def maybe_preprocess_audio( input_signal_length=None, processed_signal=None, processed_signal_length=None, + input_signal_cu_seqlens=None, ): has_input_signal = input_signal is not None and input_signal_length is not None has_processed_signal = processed_signal is not None and processed_signal_length is not None @@ -117,10 +135,18 @@ def maybe_preprocess_audio( ) if not has_processed_signal: - processed_signal, processed_signal_length = self.preprocessor( - input_signal=input_signal, - length=input_signal_length, - ) + if input_signal_cu_seqlens is None: + processed_signal, processed_signal_length = self.preprocessor( + input_signal=input_signal, + length=input_signal_length, + ) + else: + processed_signal = self.preprocessor.forward_packed( + input_signal=input_signal, + length=input_signal_length, + input_signal_cu_seqlens=input_signal_cu_seqlens, + ) + processed_signal_length = processed_signal.lengths return processed_signal, processed_signal_length def _apply_rote(self, encoder_emb, time_offset=None): @@ -163,10 +189,19 @@ def forward( return_encoder_emb=False, time_offset=None, spk_targets=None, + input_signal_cu_seqlens=None, ): processed_signal, processed_signal_length = self.maybe_preprocess_audio( - input_signal, input_signal_length, processed_signal, processed_signal_length + input_signal, + input_signal_length, + processed_signal, + processed_signal_length, + input_signal_cu_seqlens=input_signal_cu_seqlens, ) + if isinstance(processed_signal, PackedEncoderActivations): + processed_signal = unpack_encoder_output( + processed_signal, total_length=processed_signal.padded_length + ).transpose(1, 2) # Spec augment is not applied during evaluation/testing if self.spec_augmentation is not None and self.training: @@ -193,10 +228,56 @@ def forward( # b, c, t -> b, t, c encoded = self.proj(encoded.transpose(1, 2)) + result = (encoded, encoded_len) if return_encoder_emb: - return encoded, encoded_len, encoder_emb.transpose(1, 2) - else: - return encoded, encoded_len + result += (encoder_emb.transpose(1, 2),) + return result + + @typecheck.disable_checks() + def forward_sequence_packed( + self, + input_signal=None, + input_signal_length=None, + processed_signal=None, + processed_signal_length=None, + time_offset=None, + spk_targets=None, + input_signal_cu_seqlens=None, + ) -> PackedEncoderActivations: + """Encode audio natively as token-major variable-length sequences. + + This opt-in API intentionally supports only an identity modality adapter, + no multi-layer feature extraction, and no RoTE. Existing ``forward`` and + checkpoint state dictionaries are unchanged. + """ + if not self.supports_sequence_packed_output: + raise ValueError( + "Packed encoder sequences require an encoder with native packed support, " + "IdentityConnector, no multi-layer feature extraction, and rote=null." + ) + processed_signal, processed_signal_length = self.maybe_preprocess_audio( + input_signal, + input_signal_length, + processed_signal, + processed_signal_length, + input_signal_cu_seqlens=input_signal_cu_seqlens, + ) + if self.spec_augmentation is not None and self.training: + if isinstance(processed_signal, PackedEncoderActivations): + processed_signal = self.spec_augmentation.forward_packed(processed_signal) + else: + processed_signal = self.spec_augmentation(input_spec=processed_signal, length=processed_signal_length) + + encoder_kwargs = {"audio_signal": processed_signal, "length": processed_signal_length} + if spk_targets is not None: + if not self._encoder_accepts_spk_targets(self.encoder): + raise ValueError( + "`spk_targets` were provided, but the mounted perception encoder " + f"({type(self.encoder).__name__}) does not support speaker-target inputs." + ) + encoder_kwargs["spk_targets"] = spk_targets + encoded = self.encoder.forward_sequence_packed(**encoder_kwargs) + return encoded.with_data(self.proj(encoded.data)) class IdentityConnector(nn.Module): @@ -213,6 +294,175 @@ def forward(self, audio_signal, length=None, *args, **kwargs): return audio_signal, length +class IndependentDualEncoder(nn.Module): + """Run two acoustic encoders independently and concatenate their states. + + Both branches see the same preprocessed features and must have the same input + width and subsampling factor. Chunking is configured independently for each + branch. It is applied *after* feature stacking, so chunk boundaries do not + duplicate STFT padding and both outputs retain exactly the same frame grid. + + This is primarily intended for a trainable ASR encoder paired with a frozen + auxiliary encoder. A frozen branch is kept in evaluation mode and runs under + ``no_grad`` even when the parent perception module is training. + """ + + supports_sequence_packed_output = True + + def __init__( + self, + asr_encoder: nn.Module, + auxiliary_encoder: nn.Module, + *, + frame_shift_seconds: float, + asr_chunk_size_seconds: float | None = None, + auxiliary_chunk_size_seconds: float | None = None, + freeze_auxiliary: bool = True, + ) -> None: + super().__init__() + self.asr_encoder = asr_encoder + self.auxiliary_encoder = auxiliary_encoder + self.frame_shift_seconds = float(frame_shift_seconds) + self.asr_chunk_size_seconds = asr_chunk_size_seconds + self.auxiliary_chunk_size_seconds = auxiliary_chunk_size_seconds + self.freeze_auxiliary = bool(freeze_auxiliary) + self.auxiliary_encoder_config = None + + if self.frame_shift_seconds <= 0: + raise ValueError(f"frame_shift_seconds must be positive, got {frame_shift_seconds!r}.") + for name, value in ( + ("asr_chunk_size_seconds", asr_chunk_size_seconds), + ("auxiliary_chunk_size_seconds", auxiliary_chunk_size_seconds), + ): + if value is not None and float(value) <= 0: + raise ValueError(f"{name} must be positive or null, got {value!r}.") + + self._feat_in = self._matching_positive_int("_feat_in") + self.subsampling_factor = self._matching_positive_int("subsampling_factor") + self.d_model = self._encoder_width(asr_encoder) + self._encoder_width(auxiliary_encoder) + + if not all( + bool(getattr(encoder, "supports_sequence_packed_output", False)) + and callable(getattr(encoder, "forward_sequence_packed", None)) + for encoder in (asr_encoder, auxiliary_encoder) + ): + raise TypeError("IndependentDualEncoder requires two native sequence-packed encoders.") + + if self.freeze_auxiliary: + self.auxiliary_encoder.requires_grad_(False) + self.auxiliary_encoder.eval() + + def _matching_positive_int(self, attr: str) -> int: + values = [int(getattr(encoder, attr, -1) or -1) for encoder in (self.asr_encoder, self.auxiliary_encoder)] + if values[0] <= 0 or values[0] != values[1]: + raise ValueError(f"Both encoders must have the same positive {attr}; got {values}.") + return values[0] + + @staticmethod + def _encoder_width(encoder: nn.Module) -> int: + width = int(getattr(encoder, "_feat_out", getattr(encoder, "d_model", -1)) or -1) + if width <= 0: + raise ValueError(f"Could not determine output width for {type(encoder).__name__}.") + return width + + def _chunk_size_tokens(self, chunk_size_seconds: float | None) -> int | None: + if chunk_size_seconds is None: + return None + token_seconds = self.frame_shift_seconds * self.subsampling_factor + return max(1, round(float(chunk_size_seconds) / token_seconds)) + + @staticmethod + def _chunk_metadata(packed: PackedEncoderActivations, max_tokens: int) -> PackedEncoderActivations: + chunk_lengths = [] + for length in packed.lengths.detach().cpu().tolist(): + if length == 0: + chunk_lengths.append(0) + continue + chunk_lengths.extend([max_tokens] * (length // max_tokens)) + if length % max_tokens: + chunk_lengths.append(length % max_tokens) + lengths = torch.as_tensor(chunk_lengths, dtype=torch.int64, device=packed.data.device) + cu_seqlens = torch.cat( + [ + torch.zeros(1, dtype=torch.int32, device=packed.data.device), + lengths.cumsum(0, dtype=torch.int32), + ] + ).contiguous() + return PackedEncoderActivations( + data=packed.data, + lengths=lengths, + cu_seqlens=cu_seqlens, + max_seqlen=min(max_tokens, packed.max_seqlen), + padding_value=packed.padding_value, + padded_length=None, + ) + + def _forward_branch( + self, + encoder: nn.Module, + features: PackedEncoderActivations, + chunk_size_seconds: float | None, + ) -> PackedEncoderActivations: + max_tokens = self._chunk_size_tokens(chunk_size_seconds) + if max_tokens is None or features.max_seqlen <= max_tokens: + return encoder.forward_sequence_packed(features, features.lengths) + + pre_encode = getattr(encoder, "pre_encode", None) + wrapped = getattr(pre_encode, "_checkpoint_wrapped_module", pre_encode) + if type(wrapped).__name__ != "FeatureStacking": + raise TypeError( + "Post-stacking chunking requires subsampling='feature_stacking'; " + f"got {type(wrapped).__name__} for {type(encoder).__name__}." + ) + pre_encoded = pre_encode(features) + chunked = self._chunk_metadata(pre_encoded, max_tokens) + encoded_chunks = encoder.forward_sequence_packed( + chunked, + chunked.lengths, + bypass_pre_encode=True, + ) + return pre_encoded.with_data(encoded_chunks.data) + + def forward_sequence_packed(self, audio_signal, length=None) -> PackedEncoderActivations: + if not isinstance(audio_signal, PackedEncoderActivations): + if length is None: + raise ValueError("length is required for padded IndependentDualEncoder input.") + audio_signal = pack_encoder_output(audio_signal.transpose(1, 2), length) + elif length is not None and not torch.equal(length.to(audio_signal.lengths), audio_signal.lengths): + raise ValueError("length must match audio_signal.lengths for packed input.") + + asr = self._forward_branch(self.asr_encoder, audio_signal, self.asr_chunk_size_seconds) + grad_context = torch.no_grad() if self.freeze_auxiliary else contextlib.nullcontext() + with grad_context: + auxiliary = self._forward_branch( + self.auxiliary_encoder, + audio_signal, + self.auxiliary_chunk_size_seconds, + ) + if not torch.equal(asr.lengths, auxiliary.lengths): + raise RuntimeError( + "Independent encoder output lengths diverged despite matching subsampling factors: " + f"ASR={asr.lengths.detach().cpu().tolist()}, " + f"auxiliary={auxiliary.lengths.detach().cpu().tolist()}." + ) + return asr.with_data(torch.cat([asr.data, auxiliary.data], dim=-1)) + + def forward(self, audio_signal, length): + packed = self.forward_sequence_packed(audio_signal, length) + return unpack_encoder_output(packed).transpose(1, 2), packed.lengths + + def set_activation_checkpointing(self, enabled: bool) -> None: + _set_encoder_activation_checkpointing(self.asr_encoder, enabled) + if not self.freeze_auxiliary: + _set_encoder_activation_checkpointing(self.auxiliary_encoder, enabled) + + def train(self, mode: bool = True): + super().train(mode) + if self.freeze_auxiliary: + self.auxiliary_encoder.eval() + return self + + def _set_encoder_activation_checkpointing(encoder: nn.Module, enabled: bool) -> None: """Wrap the encoder's subsampling front-end and each transformer layer with ``checkpoint_wrapper`` when enabled. @@ -272,7 +522,7 @@ def preprocessor(self) -> nn.Module: return self.asr.preprocessor def __init__(self, cfg: DictConfig, pretrained_asr: str): - from nemo.collections.speechlm2.parts.pretrained import load_pretrained_nemo + from nemo.collections.speechlm2.parts.model_loading import load_pretrained_nemo super().__init__() # Initialize components diff --git a/nemo/collections/speechlm2/parts/cp_helpers.py b/nemo/collections/speechlm2/parts/cp_helpers.py index 5d4b7de4b23c..523bcbbcb32e 100644 --- a/nemo/collections/speechlm2/parts/cp_helpers.py +++ b/nemo/collections/speechlm2/parts/cp_helpers.py @@ -35,9 +35,11 @@ from torch import Tensor from torch.distributed.nn.functional import all_gather as differentiable_all_gather +from nemo.collections.asr.parts.packed_sequence import split_packed_data from nemo.collections.speechlm2.parts.encoder_chunking import ( _get_min_chunk_size_samples, encode_audio_with_optional_chunking, + materialize_packed_spk_targets, ) @@ -74,14 +76,18 @@ def encode_audio_with_cp_distribution( audios: Tensor, audio_lens: Tensor, *, + audio_cu_seqlens: Tensor | None = None, chunk_size_seconds: Optional[float], chunk_batch_size: Optional[int] = None, sampling_rate: int, cp_mesh=None, spk_targets: Tensor | None = None, spk_target_lengths: Tensor | None = None, + spk_target_cu_seqlens: Tensor | None = None, fsdp_sync_group=None, return_dummy_loss: bool = False, + sequence_packed: bool = False, + packed_cp_gather: bool = False, ) -> list[Tensor] | tuple[list[Tensor], Tensor | None]: """Distribute the audio encoder forward across CP ranks. @@ -109,8 +115,26 @@ def encode_audio_with_cp_distribution( loss term. Adding that term to the training loss preserves the autograd edge so FSDP forward/backward hooks fire on the text-only rank without affecting gradients numerically. + + ``packed_cp_gather`` separately opts into a token-flat CP collective that + pads only each rank's flattened token buffer, never ``B*max_L``. Keeping it + separate leaves the historical CP collective unchanged by default. + + Speaker targets may use the legacy dense ``[B, T_spk, N]`` representation + or a flat ``[T_spk_total, N]`` representation accompanied by + ``spk_target_cu_seqlens``. Flat targets are materialized only after entering + this perception compatibility boundary. """ - B_aud = int(audios.shape[0]) + spk_targets, spk_target_lengths = materialize_packed_spk_targets( + spk_targets, + spk_target_lengths, + spk_target_cu_seqlens, + ) + if audio_cu_seqlens is not None: + if audios.ndim != 1: + raise ValueError(f"Packed audios must be 1D, got shape {tuple(audios.shape)}.") + split_packed_data(audios, audio_lens, audio_cu_seqlens) + B_aud = int(audio_lens.numel()) fsdp_group_has_audio = _fsdp_group_has_audio(B_aud, audios.device, fsdp_sync_group) if B_aud == 0: dummy_loss = ( @@ -122,6 +146,7 @@ def encode_audio_with_cp_distribution( chunk_batch_size=chunk_batch_size, sampling_rate=sampling_rate, fsdp_sync_group=fsdp_sync_group, + sequence_packed=sequence_packed, ) if fsdp_group_has_audio else None @@ -134,6 +159,7 @@ def encode_audio_with_cp_distribution( perception, audios, audio_lens, + input_signal_cu_seqlens=audio_cu_seqlens, chunk_size_seconds=chunk_size_seconds, chunk_batch_size=chunk_batch_size, sampling_rate=sampling_rate, @@ -141,6 +167,7 @@ def encode_audio_with_cp_distribution( spk_target_lengths=spk_target_lengths, sync_group=fsdp_sync_group, return_dummy_loss=return_dummy_loss, + sequence_packed=sequence_packed, ) return ans @@ -155,10 +182,11 @@ def encode_audio_with_cp_distribution( if pad_n > 0: dummy_len = int(audio_lens.min().item()) - T_samp = audios.shape[1] - dummy_audios = torch.zeros(pad_n, T_samp, dtype=audios.dtype, device=device) dummy_lens = torch.full((pad_n,), dummy_len, dtype=audio_lens.dtype, device=device) - audios = torch.cat([audios, dummy_audios], dim=0) + if audio_cu_seqlens is None: + T_samp = audios.shape[1] + dummy_audios = torch.zeros(pad_n, T_samp, dtype=audios.dtype, device=device) + audios = torch.cat([audios, dummy_audios], dim=0) audio_lens = torch.cat([audio_lens, dummy_lens], dim=0) if spk_targets is not None: dummy_targets = torch.zeros( @@ -176,8 +204,19 @@ def encode_audio_with_cp_distribution( start = cp_rank * per_rank end = start + per_rank - local_audios = audios[start:end] local_audio_lens = audio_lens[start:end] + if audio_cu_seqlens is None: + local_audios = audios[start:end] + local_audio_cu_seqlens = None + else: + local_audios, local_audio_cu_seqlens = _slice_packed_audio_for_cp( + audios, + audio_cu_seqlens, + local_audio_lens, + start=start, + end=end, + real_batch_size=B_aud, + ) local_spk_targets = spk_targets[start:end] if spk_targets is not None else None local_spk_target_lengths = spk_target_lengths[start:end] if spk_target_lengths is not None else None @@ -185,6 +224,7 @@ def encode_audio_with_cp_distribution( perception, local_audios, local_audio_lens, + input_signal_cu_seqlens=local_audio_cu_seqlens, chunk_size_seconds=chunk_size_seconds, chunk_batch_size=chunk_batch_size, sampling_rate=sampling_rate, @@ -192,37 +232,63 @@ def encode_audio_with_cp_distribution( spk_target_lengths=local_spk_target_lengths, sync_group=fsdp_sync_group, return_dummy_loss=return_dummy_loss, + sequence_packed=sequence_packed, ) if return_dummy_loss: local_embs, dummy_loss = local_embs else: dummy_loss = None - # All-gather across CP. Variable-length: pad to a common max-L first. - H = local_embs[0].shape[-1] - local_max_L = max(e.shape[0] for e in local_embs) - max_L_t = torch.tensor(local_max_L, dtype=torch.long, device=device) - dist.all_reduce(max_L_t, op=dist.ReduceOp.MAX, group=cp_group) - max_L = int(max_L_t.item()) - - local_stack = torch.zeros(per_rank, max_L, H, device=device, dtype=local_embs[0].dtype) - local_lens = torch.zeros(per_rank, dtype=torch.long, device=device) - for i, e in enumerate(local_embs): - local_stack[i, : e.shape[0]] = e - local_lens[i] = e.shape[0] + if not packed_cp_gather: + # Backwards-compatible all-gather: pad every local row to a common max-L. + H = local_embs[0].shape[-1] + local_max_L = max(e.shape[0] for e in local_embs) + max_L_t = torch.tensor(local_max_L, dtype=torch.long, device=device) + dist.all_reduce(max_L_t, op=dist.ReduceOp.MAX, group=cp_group) + max_L = int(max_L_t.item()) + local_stack = torch.zeros(per_rank, max_L, H, device=device, dtype=local_embs[0].dtype) + local_lens = torch.zeros(per_rank, dtype=torch.long, device=device) + for i, embedding in enumerate(local_embs): + local_stack[i, : embedding.shape[0]] = embedding + local_lens[i] = embedding.shape[0] + gathered_lens = [torch.zeros_like(local_lens) for _ in range(cp_size)] + gathered_stack = differentiable_all_gather(local_stack, group=cp_group) + dist.all_gather(gathered_lens, local_lens, group=cp_group) + full_embs = [] + for rank in range(cp_size): + for idx in range(per_rank): + full_idx = rank * per_rank + idx + if full_idx >= B_aud: + break + row_length = int(gathered_lens[rank][idx].item()) + full_embs.append(gathered_stack[rank][idx, :row_length]) + return (full_embs, dummy_loss) if return_dummy_loss else full_embs + # All-gather flattened token buffers. Ranks need equal collective shapes, so + # pad only to the largest per-rank token count, never per_rank * max_row_length. + H = local_embs[0].shape[-1] + local_lens = torch.as_tensor([e.shape[0] for e in local_embs], dtype=torch.long, device=device) gathered_lens = [torch.zeros_like(local_lens) for _ in range(cp_size)] - gathered_stack = differentiable_all_gather(local_stack, group=cp_group) dist.all_gather(gathered_lens, local_lens, group=cp_group) + local_flat = torch.cat(local_embs, dim=0) + max_tokens_t = torch.tensor(local_flat.shape[0], dtype=torch.long, device=device) + dist.all_reduce(max_tokens_t, op=dist.ReduceOp.MAX, group=cp_group) + max_tokens = int(max_tokens_t.item()) + padded_flat = torch.zeros(max_tokens, H, device=device, dtype=local_flat.dtype) + padded_flat[: local_flat.shape[0]] = local_flat + gathered_flat = differentiable_all_gather(padded_flat, group=cp_group) + full_embs: list[Tensor] = [] for r in range(cp_size): + offset = 0 for i in range(per_rank): full_idx = r * per_rank + i if full_idx >= B_aud: break # dummy slot L = int(gathered_lens[r][i].item()) - full_embs.append(gathered_stack[r][i, :L]) + full_embs.append(gathered_flat[r][offset : offset + L]) + offset += L return (full_embs, dummy_loss) if return_dummy_loss else full_embs @@ -244,6 +310,7 @@ def _dummy_audio_loss_for_fsdp_sync( chunk_batch_size: Optional[int], sampling_rate: int, fsdp_sync_group=None, + sequence_packed: bool = False, ) -> Tensor | None: if chunk_batch_size is not None: _, dummy_loss = encode_audio_with_optional_chunking( @@ -255,6 +322,7 @@ def _dummy_audio_loss_for_fsdp_sync( sampling_rate=sampling_rate, sync_group=fsdp_sync_group, return_dummy_loss=True, + sequence_packed=sequence_packed, ) return dummy_loss @@ -269,6 +337,34 @@ def _dummy_audio_loss_for_fsdp_sync( dummy_lens, chunk_size_seconds=chunk_size_seconds, sampling_rate=sampling_rate, + sequence_packed=sequence_packed, ) dummy_loss = sum(emb.float().sum() for emb in dummy_embs) return dummy_loss * 0.0 + + +def _slice_packed_audio_for_cp( + audios: Tensor, + audio_cu_seqlens: Tensor, + local_audio_lens: Tensor, + *, + start: int, + end: int, + real_batch_size: int, +) -> tuple[Tensor, Tensor]: + real_start = min(start, real_batch_size) + real_end = min(end, real_batch_size) + sample_start = int(audio_cu_seqlens[real_start].item()) + sample_end = int(audio_cu_seqlens[real_end].item()) + local_audios = audios[sample_start:sample_end] + real_rows = max(0, real_end - real_start) + dummy_samples = int(local_audio_lens[real_rows:].sum().item()) + if dummy_samples: + local_audios = torch.cat([local_audios, audios.new_zeros(dummy_samples)]) + local_cu_seqlens = torch.cat( + [ + torch.zeros(1, dtype=torch.long, device=audios.device), + local_audio_lens.cumsum(dim=0, dtype=torch.long), + ] + ) + return local_audios, local_cu_seqlens diff --git a/nemo/collections/speechlm2/parts/encoder_chunking.py b/nemo/collections/speechlm2/parts/encoder_chunking.py index 33541f365c80..a4daa6109695 100644 --- a/nemo/collections/speechlm2/parts/encoder_chunking.py +++ b/nemo/collections/speechlm2/parts/encoder_chunking.py @@ -20,19 +20,24 @@ from torch import Tensor from torch.nn.utils.rnn import pad_sequence +from nemo.collections.asr.parts.packed_sequence import split_encoder_output, split_packed_data + def encode_audio_with_optional_chunking( perception: Callable, input_signal: Tensor, input_signal_length: Tensor, *, + input_signal_cu_seqlens: Tensor | None = None, chunk_size_seconds: float | None, sampling_rate: int, spk_targets: Tensor | None = None, spk_target_lengths: Tensor | None = None, + spk_target_cu_seqlens: Tensor | None = None, chunk_batch_size: int | None = None, sync_group=None, return_dummy_loss: bool = False, + sequence_packed: bool = False, ) -> list[Tensor] | tuple[list[Tensor], Tensor | None]: """Encode audio rows, splitting long rows into time chunks before the perception forward. @@ -45,8 +50,12 @@ def encode_audio_with_optional_chunking( Args: perception: Callable returning ``(audio_embs, audio_emb_lens)`` for a batched input, accepting ``input_signal=Tensor`` and ``input_signal_length=Tensor``. - input_signal: Audio batch with shape ``(B, T)`` (fp32), padded to the longest row. + input_signal: Padded audio batch with shape ``(B, T)``, or contiguous + waveform samples ``(sum(input_signal_length),)`` when + ``input_signal_cu_seqlens`` is provided. input_signal_length: Per-row valid sample counts with shape ``(B,)`` (int64). + input_signal_cu_seqlens: Optional cumulative sample offsets for packed + one-dimensional ``input_signal``. chunk_size_seconds: Target chunk length in seconds; ``None`` disables chunking. sampling_rate: Audio sampling rate, used to convert ``chunk_size_seconds`` to samples. spk_targets: Optional speaker-activity targets with shape ``(B, T_spk, N)``. @@ -54,6 +63,10 @@ def encode_audio_with_optional_chunking( audio chunks. spk_target_lengths: Optional valid speaker-target frame counts with shape ``(B,)``. Required for exact slicing of padded, mixed-length target batches. + spk_target_cu_seqlens: Optional cumulative row offsets ``(B + 1,)``. + When provided, ``spk_targets`` must be flat ``(sum(T_spk), N)``; + it is materialized only at the perception boundary that still + requires a dense speaker-target batch. chunk_batch_size: Optional maximum number of time chunks to send through ``perception`` in one forward. When unset, all chunks are encoded in the historical single forward. @@ -62,6 +75,9 @@ def encode_audio_with_optional_chunking( FSDP-sharded perception modules. return_dummy_loss: When ``True``, also return a zero-valued tensor that keeps dummy perception forwards attached to autograd. + sequence_packed: Call ``perception.forward_sequence_packed`` and preserve + compact token-major activations through the encoder. Defaults to ``False`` + for checkpoint and behavior compatibility. Returns: List of length ``B`` of fp32 embedding tensors with shape ``(T_emb_i, D)`` and @@ -70,6 +86,20 @@ def encode_audio_with_optional_chunking( original audio row. """ _validate_chunk_config(chunk_size_seconds, chunk_batch_size) + spk_targets, spk_target_lengths = materialize_packed_spk_targets( + spk_targets, + spk_target_lengths, + spk_target_cu_seqlens, + ) + packed_audio_rows = None + if input_signal_cu_seqlens is not None: + if input_signal.ndim != 1: + raise ValueError(f"Packed input_signal must be 1D, got shape {tuple(input_signal.shape)}.") + packed_audio_rows = list(split_packed_data(input_signal, input_signal_length, input_signal_cu_seqlens)) + if packed_audio_rows is not None and not sequence_packed: + input_signal = pad_sequence(packed_audio_rows, batch_first=True) + input_signal_cu_seqlens = None + packed_audio_rows = None if input_signal_length.numel() == 0: dummy_loss = _run_dummy_chunk_forwards( @@ -80,16 +110,18 @@ def encode_audio_with_optional_chunking( sampling_rate=sampling_rate, chunk_batch_size=chunk_batch_size, sync_group=sync_group, + sequence_packed=sequence_packed, ) return _maybe_return_dummy_loss([], dummy_loss, return_dummy_loss) chunk_size_samples = _get_chunk_size_samples(chunk_size_seconds, sampling_rate) perception_kwargs = {"input_signal": input_signal, "input_signal_length": input_signal_length} + if input_signal_cu_seqlens is not None: + perception_kwargs["input_signal_cu_seqlens"] = input_signal_cu_seqlens if spk_targets is not None: perception_kwargs["spk_targets"] = spk_targets if chunk_size_samples is None or input_signal_length.numel() == 0: - audio_embs, audio_emb_lens = perception(**perception_kwargs) - ans = _unpad_audio_embeddings(audio_embs, audio_emb_lens) + ans = _encode_perception_unpadded(perception, perception_kwargs, sequence_packed=sequence_packed) return _maybe_return_dummy_loss(ans, None, return_dummy_loss) min_chunk_size_samples = _get_min_chunk_size_samples(perception) @@ -105,17 +137,20 @@ def encode_audio_with_optional_chunking( ) input_signal_lengths = input_signal_length.tolist() if max(input_signal_lengths) <= chunk_size_samples and chunk_batch_size is None: - audio_embs, audio_emb_lens = perception(**perception_kwargs) - ans = _unpad_audio_embeddings(audio_embs, audio_emb_lens) + ans = _encode_perception_unpadded(perception, perception_kwargs, sequence_packed=sequence_packed) return _maybe_return_dummy_loss(ans, None, return_dummy_loss) chunks, chunk_lens, chunks_per_audio, chunk_spans = _split_audio_into_chunks( - input_signal=input_signal, + input_signal=packed_audio_rows if packed_audio_rows is not None else input_signal, input_signal_lengths=input_signal_lengths, chunk_size_samples=chunk_size_samples, min_chunk_size_samples=min_chunk_size_samples, ) - chunked_signal = pad_sequence(chunks, batch_first=True) + if input_signal_cu_seqlens is None: + chunked_signal = pad_sequence(chunks, batch_first=True) + chunked_cu_seqlens = None + else: + chunked_signal, chunked_cu_seqlens = _pack_audio_rows(chunks) chunked_lens = torch.as_tensor(chunk_lens, device=input_signal_length.device, dtype=input_signal_length.dtype) # Absolute start time (seconds) of each chunk within its source audio. # RoTE (when enabled) uses this so a chunked long audio keeps a continuous time index across chunk boundaries. @@ -134,11 +169,15 @@ def encode_audio_with_optional_chunking( "input_signal_length": chunked_lens, "time_offset": time_offset, } + if chunked_cu_seqlens is not None: + chunked_perception_kwargs["input_signal_cu_seqlens"] = chunked_cu_seqlens if chunked_spk_targets is not None: chunked_perception_kwargs["spk_targets"] = chunked_spk_targets if chunk_batch_size is None: - chunked_embs, chunked_emb_lens = perception(**chunked_perception_kwargs) - ans = _recombine_chunked_audio_embeddings(chunked_embs, chunked_emb_lens, chunks_per_audio) + chunked_embs = _encode_perception_unpadded( + perception, chunked_perception_kwargs, sequence_packed=sequence_packed + ) + ans = _recombine_chunked_audio_embedding_list(chunked_embs, chunks_per_audio) return _maybe_return_dummy_loss(ans, None, return_dummy_loss) chunked_embs, dummy_loss = _encode_chunk_microbatches( @@ -146,11 +185,66 @@ def encode_audio_with_optional_chunking( chunked_perception_kwargs, chunk_batch_size=chunk_batch_size, sync_group=sync_group, + sequence_packed=sequence_packed, ) ans = _recombine_chunked_audio_embedding_list(chunked_embs, chunks_per_audio) return _maybe_return_dummy_loss(ans, dummy_loss, return_dummy_loss) +def materialize_packed_spk_targets( + spk_targets: Tensor | None, + spk_target_lengths: Tensor | None, + spk_target_cu_seqlens: Tensor | None, +) -> tuple[Tensor | None, Tensor | None]: + """Normalize padded or flat speaker targets to the perception API's dense form.""" + if spk_target_cu_seqlens is None: + if spk_targets is not None and spk_targets.ndim != 3: + raise ValueError( + "Padded spk_targets must have shape [B, T, N]; flat [T_total, N] targets require " + f"spk_target_cu_seqlens, got shape {tuple(spk_targets.shape)}." + ) + return spk_targets, spk_target_lengths + if spk_targets is None: + raise ValueError("spk_target_cu_seqlens was provided without spk_targets") + if spk_targets.ndim != 2: + raise ValueError("Packed spk_targets must have shape [T_total, N], " f"got {tuple(spk_targets.shape)}.") + if spk_target_cu_seqlens.ndim != 1 or spk_target_cu_seqlens.numel() < 2: + raise ValueError( + "spk_target_cu_seqlens must have shape [B + 1], " f"got {tuple(spk_target_cu_seqlens.shape)}." + ) + offsets = spk_target_cu_seqlens.to(dtype=torch.long) + if int(offsets[0].item()) != 0 or int(offsets[-1].item()) != spk_targets.shape[0]: + raise ValueError( + "spk_target_cu_seqlens must start at 0 and end at the flat target length, " + f"got endpoints ({int(offsets[0].item())}, {int(offsets[-1].item())}) " + f"for {spk_targets.shape[0]} frames." + ) + packed_lengths = offsets.diff() + if bool((packed_lengths < 0).any()): + raise ValueError(f"Packed speaker-target lengths must be non-negative, got {packed_lengths.tolist()}.") + if spk_target_lengths is not None: + expected_lengths = spk_target_lengths.to(device=packed_lengths.device, dtype=packed_lengths.dtype) + if not torch.equal(expected_lengths, packed_lengths): + raise ValueError( + "spk_target_length and spk_target_cu_seqlens disagree: " + f"{expected_lengths.tolist()} vs {packed_lengths.tolist()}." + ) + else: + spk_target_lengths = packed_lengths + rows = list(torch.split(spk_targets, packed_lengths.tolist(), dim=0)) + missing_rttm_rows = torch.stack([(row == -1.0).all() for row in rows]) + padded_targets = pad_sequence(rows, batch_first=True) + # Explicit RTTM rows need zero padding to represent silence, while an all--1 + # row is a sentinel requesting the embedded diarizer. Preserve that sentinel + # across dense materialization so mixed target lengths cannot change routing. + padded_targets = torch.where( + missing_rttm_rows[:, None, None], + torch.full_like(padded_targets, -1.0), + padded_targets, + ) + return padded_targets, spk_target_lengths + + def _maybe_return_dummy_loss( audio_embs: list[Tensor], dummy_loss: Tensor | None, @@ -174,13 +268,20 @@ def _preserve_module_buffers(module: Callable): yield return - buffers = [(buffer, buffer.detach().clone()) for buffer in module.buffers()] + buffers = [(buffer, buffer.detach().clone(), buffer._version) for buffer in module.buffers()] try: yield finally: with torch.no_grad(): - for buffer, value in buffers: - buffer.copy_(value) + for buffer, value, version in buffers: + # A blind copy bumps the autograd version even for immutable + # buffers. PEE's [n_spk, d_model] diarization kernel is saved by + # matmul in every real microbatch, so copying it after a synced + # dummy forward makes the eventual backward fail with an + # in-place-modification error. Restore only buffers the dummy + # forward actually mutated. + if buffer._version != version: + buffer.copy_(value) def _encode_chunk_microbatches( @@ -189,10 +290,11 @@ def _encode_chunk_microbatches( *, chunk_batch_size: int, sync_group=None, + sequence_packed: bool = False, ) -> tuple[list[Tensor], Tensor | None]: """Encode chunks in smaller forwards while keeping FSDP ranks synchronized.""" input_signal = chunked_perception_kwargs["input_signal"] - num_chunks = int(input_signal.shape[0]) + num_chunks = int(chunked_perception_kwargs["input_signal_length"].numel()) local_microbatches = (num_chunks + chunk_batch_size - 1) // chunk_batch_size total_microbatches = _sync_max_count(local_microbatches, input_signal.device, sync_group) @@ -203,21 +305,34 @@ def _encode_chunk_microbatches( end = min(start + chunk_batch_size, num_chunks) if start < end: mb_kwargs = _slice_perception_kwargs(chunked_perception_kwargs, start, end) - mb_embs, mb_lens = perception(**mb_kwargs) - chunked_embs.extend(_unpad_audio_embeddings(mb_embs, mb_lens)) + chunked_embs.extend(_encode_perception_unpadded(perception, mb_kwargs, sequence_packed=sequence_packed)) continue dummy_kwargs = _slice_perception_kwargs(chunked_perception_kwargs, 0, 1) with _preserve_module_buffers(perception): - mb_embs, _ = perception(**dummy_kwargs) - zero = mb_embs.float().sum() * 0.0 + mb_embs = _encode_perception_unpadded(perception, dummy_kwargs, sequence_packed=sequence_packed) + zero = sum(emb.float().sum() for emb in mb_embs) * 0.0 dummy_loss = zero if dummy_loss is None else dummy_loss + zero return chunked_embs, dummy_loss def _slice_perception_kwargs(perception_kwargs: dict[str, Tensor], start: int, end: int) -> dict[str, Tensor]: - return {name: value[start:end] for name, value in perception_kwargs.items()} + cu_seqlens = perception_kwargs.get("input_signal_cu_seqlens") + if cu_seqlens is None: + return {name: value[start:end] for name, value in perception_kwargs.items()} + + sample_start = int(cu_seqlens[start].item()) + sample_end = int(cu_seqlens[end].item()) + sliced = {} + for name, value in perception_kwargs.items(): + if name == "input_signal": + sliced[name] = value[sample_start:sample_end] + elif name == "input_signal_cu_seqlens": + sliced[name] = value[start : end + 1] - sample_start + else: + sliced[name] = value[start:end] + return sliced def _sync_max_count(local_count: int, device: torch.device, sync_group=None) -> int: @@ -237,6 +352,7 @@ def _run_dummy_chunk_forwards( sampling_rate: int, chunk_batch_size: int | None, sync_group=None, + sequence_packed: bool = False, ) -> Tensor | None: """Run synced zero-valued perception forwards for audio-free ranks.""" if chunk_batch_size is None or sync_group is None or not (dist.is_available() and dist.is_initialized()): @@ -256,8 +372,12 @@ def _run_dummy_chunk_forwards( dummy_loss = None for _ in range(total_microbatches): with _preserve_module_buffers(perception): - dummy_embs, _ = perception(input_signal=dummy_audio, input_signal_length=dummy_lens) - zero = dummy_embs.float().sum() * 0.0 + dummy_embs = _encode_perception_unpadded( + perception, + {"input_signal": dummy_audio, "input_signal_length": dummy_lens}, + sequence_packed=sequence_packed, + ) + zero = sum(emb.float().sum() for emb in dummy_embs) * 0.0 dummy_loss = zero if dummy_loss is None else dummy_loss + zero return dummy_loss @@ -330,7 +450,7 @@ def _get_spk_target_stride(perception: Callable) -> int: def _split_audio_into_chunks( - input_signal: Tensor, + input_signal: Tensor | list[Tensor], input_signal_lengths: list[int], chunk_size_samples: int, min_chunk_size_samples: int, @@ -344,7 +464,7 @@ def _split_audio_into_chunks( that ``chunks_per_audio`` stays aligned with the input batch. Args: - input_signal: ``(B, T)`` audio batch. + input_signal: ``(B, T)`` audio batch or a list of exact waveform views. input_signal_lengths: Per-row valid sample counts (length ``B``). chunk_size_samples: Target chunk length in samples. min_chunk_size_samples: Minimum chunk length below which a tail chunk is folded @@ -492,6 +612,24 @@ def _unpad_audio_embeddings(audio_embs: Tensor, audio_emb_lens: Tensor) -> list[ return [emb[:emblen] for emb, emblen in zip(audio_embs, audio_emb_lens)] +def _encode_perception_unpadded( + perception: Callable, + perception_kwargs: dict[str, Tensor], + *, + sequence_packed: bool, +) -> list[Tensor]: + if not sequence_packed: + audio_embs, audio_emb_lens = perception(**perception_kwargs) + return _unpad_audio_embeddings(audio_embs, audio_emb_lens) + if not bool(getattr(perception, 'supports_sequence_packed_output', False)): + raise ValueError( + "packed_encoder_sequences=true, but the mounted perception stack does not support native packed output. " + "Use TransformerEncoder/ParallelExpertEncoder with IdentityConnector and rote=null." + ) + packed = perception.forward_sequence_packed(**perception_kwargs) + return split_encoder_output(packed) + + def _recombine_chunked_audio_embeddings( chunked_embs: Tensor, chunked_emb_lens: Tensor, @@ -524,3 +662,9 @@ def _recombine_chunked_audio_embedding_list( audio_embs.append(parts[0] if len(parts) == 1 else torch.cat(parts, dim=0)) chunk_idx += num_chunks return audio_embs + + +def _pack_audio_rows(rows: list[Tensor]) -> tuple[Tensor, Tensor]: + lengths = torch.tensor([row.numel() for row in rows], dtype=torch.long, device=rows[0].device) + cu_seqlens = torch.cat([lengths.new_zeros(1), lengths.cumsum(dim=0)]) + return torch.cat(rows), cu_seqlens diff --git a/nemo/collections/speechlm2/parts/gc.py b/nemo/collections/speechlm2/parts/gc.py new file mode 100644 index 000000000000..af0a11482dbd --- /dev/null +++ b/nemo/collections/speechlm2/parts/gc.py @@ -0,0 +1,54 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from nemo.utils import logging + + +class GarbageCollectionManager: + """Manage deterministic Python garbage collection during distributed training. + + When enabled, automatic garbage collection is replaced at fit start by + NeMo Automodel's generation-1 collector. The manager owns the optimizer-step + counter so model implementations only need to forward lifecycle events. + """ + + def __init__(self, gc_every_steps: int | None) -> None: + if gc_every_steps is not None and ( + isinstance(gc_every_steps, bool) or not isinstance(gc_every_steps, int) or gc_every_steps <= 0 + ): + raise ValueError(f"model.gc_every_steps must be a positive integer or null, got {gc_every_steps!r}") + self.gc_every_steps = gc_every_steps + self._collector = None + self._optimizer_step_count = 0 + + def on_fit_start(self) -> None: + """Disable automatic GC and initialize the configured manual collector.""" + if self.gc_every_steps is None: + return + + from nemo_automodel.components.training.garbage_collection import GarbageCollection + + self._collector = GarbageCollection(gc_every_steps=self.gc_every_steps) + self._optimizer_step_count = 0 + logging.info( + "Automatic Python GC disabled; generation-1 collection will run every %d optimizer steps", + self.gc_every_steps, + ) + + def on_optimizer_step(self) -> None: + """Advance the manual collector after a completed optimizer step.""" + if self._collector is None: + return + self._optimizer_step_count += 1 + self._collector.run(self._optimizer_step_count) diff --git a/nemo/collections/speechlm2/parts/model_loading.py b/nemo/collections/speechlm2/parts/model_loading.py new file mode 100644 index 000000000000..a32354984c52 --- /dev/null +++ b/nemo/collections/speechlm2/parts/model_loading.py @@ -0,0 +1,45 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Dependency-neutral helpers for loading pretrained NeMo models.""" + +from pathlib import Path + + +def load_pretrained_nemo(cls, model_path_or_name: str): + """Load a pretrained NeMo model from a local archive or registered model name.""" + if Path(model_path_or_name).exists() and model_path_or_name.endswith(".nemo"): + # Local .nemo restore_from() does not resolve the config target, so resolve + # the concrete class first, matching from_pretrained() behavior. + cfg = cls.restore_from(model_path_or_name, return_config=True) + target = cfg.get("target", None) if hasattr(cfg, "get") else None + if target is not None: + from nemo.core.classes.common import _get_allowed_target_class + + resolved_cls = _get_allowed_target_class(target) + concrete_cls = resolved_cls + while hasattr(concrete_cls, "__wrapped__"): + concrete_cls = concrete_cls.__wrapped__ + if not isinstance(concrete_cls, type) or not issubclass(concrete_cls, cls): + raise TypeError(f"Checkpoint target {target!r} is not a subclass of {cls.__name__}.") + cls = resolved_cls + return cls.restore_from(model_path_or_name) + return cls.from_pretrained(model_path_or_name) + + +def load_pretrained_nemo_config(cls, model_path_or_name: str): + """Load a NeMo model config without loading model weights.""" + if Path(model_path_or_name).exists() and model_path_or_name.endswith(".nemo"): + return cls.restore_from(model_path_or_name, return_config=True) + return cls.from_pretrained(model_path_or_name, return_config=True) diff --git a/nemo/collections/speechlm2/parts/packed_sequences.py b/nemo/collections/speechlm2/parts/packed_sequences.py index c57af2f8adbb..b1fa0c68b226 100644 --- a/nemo/collections/speechlm2/parts/packed_sequences.py +++ b/nemo/collections/speechlm2/parts/packed_sequences.py @@ -27,8 +27,10 @@ from __future__ import annotations +import math +from collections.abc import Callable from dataclasses import dataclass -from typing import Any, Optional +from typing import Any import torch from torch import Tensor @@ -45,6 +47,119 @@ def pack_audio_into_text_embeds( placeholder_id: int, cp_size: int = 1, tp_size: int = 1, + token_alignment: int = 1, + ignore_index: int = -100, +) -> dict[str, Tensor]: + """Left-unpad dense text embeddings, splice audio frames, and pack THD. + + This compatibility wrapper accepts the historical dense ``[B, S, H]`` + embedding input. The SALMAutomodel packed path uses + :func:`prepare_packed_llm_inputs` with ``embed_tokens`` instead, so it can + remove left padding before the embedding lookup. + """ + ids_unpad, embs_unpad, tgts_unpad = _unpad_inputs(input_ids, embeds, target_ids, padding_id) + assert tgts_unpad is not None + return _pack_audio_into_unpadded_text_embeds( + input_ids=ids_unpad, + embeds=embs_unpad, + target_ids=tgts_unpad, + replacements=replacements, + padding_id=padding_id, + placeholder_id=placeholder_id, + cp_size=cp_size, + tp_size=tp_size, + token_alignment=token_alignment, + ignore_index=ignore_index, + ) + + +def _split_and_embed_text_ids( + input_ids: Tensor, + target_ids: Tensor, + padding_id: int, + placeholder_id: int, + embed_tokens: Callable[[Tensor], Tensor], + text_cu_seqlens: Tensor | None = None, +) -> tuple[list[Tensor], list[Tensor], list[Tensor]]: + """Embed only real text positions from padded or already-packed rows. + + Dense rows are left-unpadded before the embedding lookup. Flat rows are + already padding-free and are split using ``text_cu_seqlens``. Audio + placeholders are mapped to token 0 exactly as in the historical dense + SALMAutomodel path; their embeddings are overwritten by audio frames later. + """ + if target_ids.shape != input_ids.shape: + raise ValueError( + "input_ids and target_ids must have the same shape, " + f"got {tuple(input_ids.shape)} and {tuple(target_ids.shape)}" + ) + if input_ids.ndim == 2: + if text_cu_seqlens is not None: + raise ValueError("text_cu_seqlens must be omitted for padded [B, S] input_ids") + if input_ids.shape[0] == 0 or input_ids.shape[1] == 0: + raise ValueError(f"Cannot pack empty input_ids with shape {tuple(input_ids.shape)}") + non_padding = input_ids != padding_id + has_non_padding = non_padding.any(dim=1) + starts = non_padding.to(torch.int64).argmax(dim=1) + # Match _unpad_inputs for an all-padding row: retain its final slot. + starts = torch.where(has_non_padding, starts, torch.full_like(starts, input_ids.shape[1] - 1)) + columns = torch.arange(input_ids.shape[1], device=input_ids.device).unsqueeze(0) + keep = columns >= starts.unsqueeze(1) + row_lengths = keep.sum(dim=1).tolist() + flat_ids = input_ids[keep] + flat_targets = target_ids[keep] + elif input_ids.ndim == 1: + row_lengths = _validate_packed_text_offsets(input_ids, text_cu_seqlens) + flat_ids = input_ids + flat_targets = target_ids + else: + raise ValueError(f"input_ids must have shape [B, S] or [T], got {tuple(input_ids.shape)}") + + flat_ids_to_embed = torch.where(flat_ids == placeholder_id, 0, flat_ids) + flat_embeds = embed_tokens(flat_ids_to_embed) + if flat_embeds.ndim != 2 or flat_embeds.shape[0] != flat_ids.shape[0]: + raise ValueError( + "embed_tokens must map flat token IDs [T] to embeddings [T, H], " + f"got input shape {tuple(flat_ids.shape)} and output shape {tuple(flat_embeds.shape)}" + ) + embs_unpad = list(torch.split(flat_embeds, row_lengths, dim=0)) + ids_unpad = list(torch.split(flat_ids, row_lengths, dim=0)) + tgts_unpad = list(torch.split(flat_targets, row_lengths, dim=0)) + return ids_unpad, embs_unpad, tgts_unpad + + +def _validate_packed_text_offsets(input_ids: Tensor, text_cu_seqlens: Tensor | None) -> list[int]: + """Validate flat text row offsets and return per-row lengths.""" + if text_cu_seqlens is None: + raise ValueError("Flat input_ids [T] require text_cu_seqlens [B + 1]") + if text_cu_seqlens.ndim != 1 or text_cu_seqlens.numel() < 2: + raise ValueError( + "text_cu_seqlens must be a 1-D tensor with at least two entries, " + f"got shape {tuple(text_cu_seqlens.shape)}" + ) + offsets = text_cu_seqlens.to(dtype=torch.long) + if int(offsets[0].item()) != 0 or int(offsets[-1].item()) != input_ids.numel(): + raise ValueError( + "text_cu_seqlens must start at 0 and end at the flat input length, " + f"got endpoints ({int(offsets[0].item())}, {int(offsets[-1].item())}) " + f"for {input_ids.numel()} tokens" + ) + lengths = offsets.diff() + if bool((lengths <= 0).any()): + raise ValueError(f"Every packed text row must be non-empty, got lengths {lengths.tolist()}") + return lengths.tolist() + + +def _pack_audio_into_unpadded_text_embeds( + input_ids: list[Tensor], + embeds: list[Tensor], + target_ids: list[Tensor], + replacements: list[Tensor], + padding_id: int, + placeholder_id: int, + cp_size: int = 1, + tp_size: int = 1, + token_alignment: int = 1, ignore_index: int = -100, ) -> dict[str, Tensor]: """Splice audio frames into per-utterance text embeddings and pack into THD. @@ -55,15 +170,15 @@ def pack_audio_into_text_embeds( can be called without any further shift. Args: - input_ids: ``[B, S]`` int64; left-padded. - embeds: ``[B, S, H]`` text-token embeddings (placeholder slots - are pre-zeroed by the caller; they get overwritten). - target_ids: ``[B, S]`` int64; ``-100`` outside assistant spans. + input_ids: list of tight ``[S_i]`` int64 token-ID tensors. + embeds: list of tight ``[S_i, H]`` text-embedding tensors + (placeholder slots are overwritten). + target_ids: list of tight ``[S_i]`` labels; ``-100`` outside + assistant spans. replacements: list of ``[L_i, H]`` audio-frame embeddings, one per placeholder occurrence in row-major order. - padding_id: pad-token id in ``input_ids`` (used to strip left-pad - and to mark padding positions as ``ignore_index`` in - labels). + padding_id: pad-token id, retained for any internal pad-token + positions after left unpadding. placeholder_id: the ``<|audio|>`` token id. cp_size: per-utterance flat lengths are rounded up to a multiple of ``2 * cp_size`` (TE-CP requirement). @@ -71,6 +186,9 @@ def pack_audio_into_text_embeds( ``T_total % tp_size == 0`` (sequence-parallel). When CP is active, the bump also preserves the ``2 * cp_size`` per-utterance alignment. + token_alignment: the token count seen by each CP rank is rounded up to + this multiple. Transformer Engine FP8 requires 8; the + default 1 preserves existing behavior. ignore_index: label fill for audio-frame slots, padding slots, and the last position of every utterance. @@ -87,13 +205,17 @@ def pack_audio_into_text_embeds( - ``max_seqlen`` int32 scalar, ``max(seq_lens_padded)`` - ``qkv_format`` ``"thd"`` """ - B = input_ids.shape[0] - H = embeds.shape[-1] - device = embeds.device - dtype = embeds.dtype - - # Strip left-padding so per-utt sequences are tight before splicing. - ids_unpad, embs_unpad, tgts_unpad = _unpad_inputs(input_ids, embeds, target_ids, padding_id) + B = len(input_ids) + if B == 0: + raise ValueError("Cannot pack an empty SALM minibatch") + if not (len(embeds) == len(target_ids) == B): + raise ValueError("input_ids, embeds, and target_ids must contain the same number of rows") + H = embeds[0].shape[-1] + device = embeds[0].device + dtype = embeds[0].dtype + ids_unpad = input_ids + embs_unpad = embeds + tgts_unpad = target_ids seq_embs: list[Tensor] = [] seq_labs: list[Tensor] = [] @@ -143,28 +265,34 @@ def pack_audio_into_text_embeds( f"Used {rep_idx} of {len(replacements)} audio replacements — " f"placeholder occurrences in input_ids do not match replacements length." ) + if not isinstance(token_alignment, int) or token_alignment < 1: + raise ValueError(f"token_alignment must be a positive integer, got {token_alignment!r}") # Round each utterance's length up to a multiple of 2*cp_size (TE-CP # interleaves 2 chunks per rank); skip rounding when cp_size == 1. Then - # bump the last so the total is divisible by tp_size for sequence - # parallelism, preserving the CP alignment when CP is active. + # bump the last so the total satisfies both TP and backend alignment while + # preserving the per-utterance CP alignment when CP is active. if cp_size > 1: cp_mult = 2 * cp_size padded_lens = [((L + cp_mult - 1) // cp_mult) * cp_mult for L in real_lens] else: padded_lens = list(real_lens) - if tp_size > 1: + # Context parallelism shards the packed token dimension before the LLM, so + # the global total must contain ``token_alignment`` tokens per CP rank. + # Retain the existing TP divisibility contract at the same time. + total_alignment = math.lcm(tp_size, token_alignment * cp_size) + if total_alignment > 1: total_len = sum(padded_lens) if cp_size > 1: cp_mult = 2 * cp_size - tp_bump = 0 - while (total_len + tp_bump) % tp_size != 0: - tp_bump += cp_mult - padded_lens[-1] += tp_bump + alignment_bump = 0 + while (total_len + alignment_bump) % total_alignment != 0: + alignment_bump += cp_mult + padded_lens[-1] += alignment_bump else: - rem = total_len % tp_size + rem = total_len % total_alignment if rem != 0: - padded_lens[-1] += tp_size - rem + padded_lens[-1] += total_alignment - rem # Materialize the flat THD batch. flat_emb_segs: list[Tensor] = [] @@ -271,19 +399,25 @@ def _shard_packed_for_cp( def prepare_packed_llm_inputs( input_ids: Tensor, - text_embs: Tensor, + text_embs: Tensor | None, audio_embs: list[Tensor], target_ids: Tensor, padding_id: int, placeholder_id: int, - device_mesh: Optional[Any] = None, + device_mesh: Any | None = None, mtp_num_depths: int = 0, + embed_tokens: Callable[[Tensor], Tensor] | None = None, + text_cu_seqlens: Tensor | None = None, + token_alignment: int = 1, ) -> dict[str, Any]: """Pack a SALM minibatch and (optionally) shard it across CP ranks. Args: - input_ids: Token IDs of shape [batch, sequence]. - text_embs: Text embeddings of shape [batch, sequence, hidden]. + input_ids: Token IDs of shape [batch, sequence] in padded mode or + [total_text_tokens] in native packed mode. + text_embs: Legacy dense text embeddings of shape + [batch, sequence, hidden]. Pass ``None`` together with + ``embed_tokens`` to compact token IDs before embedding. audio_embs: Audio replacement tensors, each of shape [audio_frames, hidden]. target_ids: Unshifted labels of shape [batch, sequence]. padding_id: Token ID used for left padding. @@ -292,6 +426,13 @@ def prepare_packed_llm_inputs( mtp_num_depths: Number of future-token MTP input/target tensors to prepare. These are emitted only when CP is active because rank-local rolling is otherwise incorrect. + embed_tokens: Callable mapping flat IDs ``[T]`` to embeddings + ``[T, H]``. Exactly one of ``text_embs`` and ``embed_tokens`` + must be provided. + text_cu_seqlens: Cumulative text row offsets [batch + 1], required + when ``input_ids`` is flat and omitted for padded inputs. + token_alignment: Required multiple for the packed-token count seen by + each CP rank; default 1 preserves the existing behavior. Returns: Mapping containing rank-local THD model inputs. ``input_embeds`` has @@ -311,18 +452,43 @@ def prepare_packed_llm_inputs( if "tp" in names and device_mesh["tp"].size() > 1: tp_size = device_mesh["tp"].size() - packed = pack_audio_into_text_embeds( - input_ids=input_ids, - embeds=text_embs, - target_ids=target_ids, - replacements=audio_embs, - padding_id=padding_id, - placeholder_id=placeholder_id, - cp_size=cp_size, - tp_size=tp_size, - ) + if (text_embs is None) == (embed_tokens is None): + raise ValueError("Exactly one of text_embs and embed_tokens must be provided") + if embed_tokens is not None: + ids_unpad, embs_unpad, tgts_unpad = _split_and_embed_text_ids( + input_ids=input_ids, + target_ids=target_ids, + padding_id=padding_id, + placeholder_id=placeholder_id, + embed_tokens=embed_tokens, + text_cu_seqlens=text_cu_seqlens, + ) + packed = _pack_audio_into_unpadded_text_embeds( + input_ids=ids_unpad, + embeds=embs_unpad, + target_ids=tgts_unpad, + replacements=audio_embs, + padding_id=padding_id, + placeholder_id=placeholder_id, + cp_size=cp_size, + tp_size=tp_size, + token_alignment=token_alignment, + ) + else: + packed = pack_audio_into_text_embeds( + input_ids=input_ids, + embeds=text_embs, + target_ids=target_ids, + replacements=audio_embs, + padding_id=padding_id, + placeholder_id=placeholder_id, + cp_size=cp_size, + tp_size=tp_size, + token_alignment=token_alignment, + ) num_tokens = packed["seq_lens"].sum() - num_examples = torch.tensor(input_ids.shape[0], dtype=torch.long, device=input_ids.device) + batch_size = input_ids.shape[0] if text_cu_seqlens is None else text_cu_seqlens.numel() - 1 + num_examples = torch.tensor(batch_size, dtype=torch.long, device=input_ids.device) mtp_inputs = None if cp_mesh is not None and mtp_num_depths > 0: diff --git a/nemo/collections/speechlm2/parts/parallel.py b/nemo/collections/speechlm2/parts/parallel.py index 6195cb1b22c4..892b0ba834ce 100644 --- a/nemo/collections/speechlm2/parts/parallel.py +++ b/nemo/collections/speechlm2/parts/parallel.py @@ -18,6 +18,7 @@ import os import warnings from datetime import timedelta +from pathlib import Path from typing import Any, Dict, Optional import torch @@ -26,12 +27,86 @@ from lightning.pytorch.strategies.model_parallel import ModelParallelStrategy from typing_extensions import override - # Blackwell sm_120, where TE 2.14's cuDNN fused-attention backward kernel # silently amplifies THD/padding_causal gradients 8x-960x per layer. _SM120 = (12, 0) +def _validate_missing_optimizer_state( + *, + target_keys: set[str], + checkpoint_keys: set[str], + parameter_names: set[str], + optimizer_key: str, +) -> list[str]: + """Allow only wholly absent per-parameter optimizer state. + + PyTorch optimizers create state lazily. A parameter that has never received + a gradient therefore has no checkpoint entries, while + ``get_optimizer_state_dict`` initializes placeholders for every parameter + when preparing a fresh restore target. DCP's strict planner treats that + expected asymmetry as a missing-key error. + + Missing *complete* parameter states are safe to leave initialized locally. + A partially present state (for example ``step`` without ``exp_avg``), or a + missing key outside ``optimizer..``, still indicates an + incompatible/corrupt checkpoint and is rejected. + """ + missing_keys = target_keys - checkpoint_keys + if not missing_keys: + return [] + + prefixes = {name: f"{optimizer_key}.state.{name}" for name in parameter_names} + owned_target_keys: dict[str, set[str]] = {name: set() for name in parameter_names} + for key in target_keys: + owners = [name for name, prefix in prefixes.items() if key == prefix or key.startswith(f"{prefix}.")] + if owners: + # Parameter FQNs are normally not prefixes of each other. Choosing + # the longest match also handles that edge case deterministically. + owned_target_keys[max(owners, key=len)].add(key) + + missing_parameters = [] + classified_missing_keys = set() + for name, expected_keys in owned_target_keys.items(): + missing_for_parameter = expected_keys - checkpoint_keys + if not missing_for_parameter: + continue + if missing_for_parameter != expected_keys: + present = sorted(expected_keys & checkpoint_keys) + missing = sorted(missing_for_parameter) + raise RuntimeError( + f"Checkpoint contains partial optimizer state for parameter {name!r}: " + f"present={present[:3]} missing={missing[:3]}" + ) + missing_parameters.append(name) + classified_missing_keys.update(missing_for_parameter) + + unexpected_missing = missing_keys - classified_missing_keys + if unexpected_missing: + raise RuntimeError( + "Checkpoint is missing optimizer metadata or unrecognized state keys: " f"{sorted(unexpected_missing)[:5]}" + ) + return sorted(missing_parameters) + + +def _optimizer_load_planner(optimizer_state: dict, metadata, optimizer_key: str): + """Return a DCP planner that tolerates only never-initialized parameters.""" + from torch.distributed.checkpoint.default_planner import DefaultLoadPlanner + + strict_planner = DefaultLoadPlanner() + strict_planner.set_up_planner(optimizer_state, metadata, is_coordinator=False) + parameter_state = optimizer_state[optimizer_key].get("state", {}) + missing_parameters = _validate_missing_optimizer_state( + target_keys=set(strict_planner.state_dict), + checkpoint_keys=set(metadata.state_dict_metadata), + parameter_names=set(parameter_state), + optimizer_key=optimizer_key, + ) + if not missing_parameters: + return strict_planner, missing_parameters + return DefaultLoadPlanner(allow_partial_load=True), missing_parameters + + def validate_parallelism_compatibility( *, packed_sequences: bool, @@ -102,7 +177,7 @@ def validate_parallelism_compatibility( if check_backward and nvte_fused_attn != "0": msg = ( "SALMAutomodel: ``packed_sequences=true`` with ``attn=te`` and " - "``NVTE_FUSED_ATTN`` not set to ``\"0\"`` (got " + '``NVTE_FUSED_ATTN`` not set to ``"0"`` (got ' f"{nvte_fused_attn!r}). TE 2.14's cuDNN fused-attention " "backward kernel amplifies THD/padding_causal gradients " "8x-960x per layer on Blackwell sm_120; the resulting ``inf`` " @@ -198,6 +273,64 @@ class AutomodelParallelStrategy(ModelParallelStrategy): timeout: Process group initialization timeout. """ + @override + def load_checkpoint(self, checkpoint_path): + """Load DCP optimizer state while preserving lazy-state semantics. + + Model tensors remain strict. Optimizer loading alone permits a complete + per-parameter state to be absent when the parameter never received a + gradient before the save; partial states and missing optimizer metadata + remain hard errors. + """ + from lightning.pytorch.strategies.model_parallel import _METADATA_FILENAME, _is_sharded_checkpoint + + path = Path(self.broadcast(checkpoint_path)) + if not _is_sharded_checkpoint(path): + return super().load_checkpoint(path) + + from torch.distributed.checkpoint import FileSystemReader, load + from torch.distributed.checkpoint.state_dict import ( + StateDictOptions, + get_model_state_dict, + get_optimizer_state_dict, + set_optimizer_state_dict, + ) + + assert self.model is not None + assert self.lightning_module is not None + module_state = {"state_dict": get_model_state_dict(self.model)} + load(module_state, checkpoint_id=path) + self.model.load_state_dict(module_state["state_dict"], strict=self.lightning_module.strict_loading) + + state_dict_options = StateDictOptions(cpu_offload=True) + metadata = FileSystemReader(path).read_metadata() + for idx, optimizer in enumerate(self.optimizers): + optimizer_key = f"optimizer_{idx}" + optimizer_state = {optimizer_key: get_optimizer_state_dict(self.model, optimizer)} + planner, missing_parameters = _optimizer_load_planner(optimizer_state, metadata, optimizer_key) + load(optimizer_state, checkpoint_id=path, planner=planner) + set_optimizer_state_dict( + self.model, + optimizer, + optim_state_dict=optimizer_state[optimizer_key], + options=state_dict_options, + ) + if missing_parameters and self.global_rank == 0: + warnings.warn( + f"Initialized empty optimizer state for {len(missing_parameters)} parameter(s) that had no " + "state in the checkpoint because they had not received a gradient. " + f"Examples: {missing_parameters[:3]}", + stacklevel=2, + ) + + # Lightning drops its temporary loaded-checkpoint reference at the end + # of resume. Keep this metadata-only payload alive for consumers whose + # restored state can outlive that connector reference; model and + # optimizer tensors were loaded separately through DCP above. + checkpoint = torch.load(path / _METADATA_FILENAME) + self._checkpoint_keepalive = checkpoint + return checkpoint + def __init__( self, dp_size: Optional[int] = None, @@ -243,6 +376,7 @@ def __init__( self._activation_checkpointing_perception = activation_checkpointing_perception self._moe_mesh = None self._distributed_setup = None + self._checkpoint_keepalive = None @property def moe_mesh(self): diff --git a/nemo/collections/speechlm2/parts/pretrained.py b/nemo/collections/speechlm2/parts/pretrained.py index c9f5d4fe39f4..6b31c09dba30 100644 --- a/nemo/collections/speechlm2/parts/pretrained.py +++ b/nemo/collections/speechlm2/parts/pretrained.py @@ -25,48 +25,18 @@ from nemo.collections.asr.models import ASRModel from nemo.collections.asr.modules.parallel_expert_encoder import ParallelExpertEncoderPT from nemo.collections.speechlm2.modules import AudioPerceptionModule +from nemo.collections.speechlm2.parts.model_loading import load_pretrained_nemo, load_pretrained_nemo_config from nemo.collections.speechlm2.parts.precision import fp32_precision from nemo.collections.tts.models import AudioCodecModel from nemo.utils import logging from nemo.utils.compat import python313_pathlib_pickle_compat -def load_pretrained_nemo(cls, model_path_or_name: str): - """ - Load pretrained NeMo 1.0 model (inheriting from ModelPT). Works with ASR, TTS, codec models. - - Setting ``pretrained_weights=False`` returns a model that has identical architecture with the checkpoint, - but is randomly initialized. - """ - if Path(model_path_or_name).exists() and model_path_or_name.endswith(".nemo"): - # Local .nemo restore_from() doesn't resolve the config's `target` (instantiates - # the abstract base). Resolve the concrete class first, like from_pretrained(). - cfg = cls.restore_from(model_path_or_name, return_config=True) - target = cfg.get("target", None) if hasattr(cfg, "get") else None - if target is not None: - from nemo.core.classes.common import _get_allowed_target_class - - resolved_cls = _get_allowed_target_class(target) - concrete_cls = resolved_cls - while hasattr(concrete_cls, "__wrapped__"): - concrete_cls = concrete_cls.__wrapped__ - if not isinstance(concrete_cls, type) or not issubclass(concrete_cls, cls): - raise TypeError(f"Checkpoint target {target!r} is not a subclass of {cls.__name__}.") - cls = resolved_cls - return cls.restore_from(model_path_or_name) - else: - return cls.from_pretrained(model_path_or_name) - - -def load_pretrained_nemo_config(cls, model_path_or_name: str): - """Load a NeMo model config without loading model weights.""" - if Path(model_path_or_name).exists() and model_path_or_name.endswith(".nemo"): - return cls.restore_from(model_path_or_name, return_config=True) - return cls.from_pretrained(model_path_or_name, return_config=True) - - def load_pretrained_hf( - model_path_or_name: str, pretrained_weights: bool = True, dtype=torch.float32, trust_remote_code: bool = False + model_path_or_name: str, + pretrained_weights: bool = True, + dtype=torch.float32, + trust_remote_code: bool = False, ): """ Load pretrained HuggingFace AutoModelForCausalLM. @@ -322,12 +292,96 @@ def setup_speech_encoder(model: torch.nn.Module, pretrained_weights: bool = True # When a multilayer/Qformer connector is used, the encoder lives at # ``encoder_multilayer.encoder.*`` rather than ``encoder.*``; remap ASR # state-dict keys so pretrained encoder weights actually load. - if isinstance(model.perception.modality_adapter, (QformerConnector, MultiLayerProjectionConnector)): + if isinstance( + model.perception.modality_adapter, + (QformerConnector, MultiLayerProjectionConnector), + ): asr_sd = {("encoder_multilayer." + k if k.startswith("encoder.") else k): v for k, v in asr_sd.items()} model.perception.load_state_dict(asr_sd, strict=False) if model.cfg.get("pe_encoder_path", None) not in (None, "", False): + if model.cfg.get("speaker_encoder", None) not in (None, "", False): + raise ValueError("pe_encoder_path and speaker_encoder are mutually exclusive.") setup_parallel_expert_encoder(model) + elif model.cfg.get("speaker_encoder", None) not in (None, "", False): + setup_independent_speaker_encoder(model) + + +def setup_independent_speaker_encoder(model: torch.nn.Module): + """Add a standalone speaker Transformer beside the pretrained ASR encoder. + + ``model.speaker_encoder.path`` points at a rendered artifact directory with + ``model_config.yaml`` and ``model.safetensors``. The two encoders execute + independently inside :class:`IndependentDualEncoder`; their same-rate states + are concatenated before the existing perception-to-LLM projection. + """ + from nemo.collections.speechlm2.modules.perception import IdentityConnector, IndependentDualEncoder + + cfg = model.cfg.speaker_encoder + artifact = Path(str(cfg.get("path", ""))) + config_path = artifact / "model_config.yaml" + weights_path = artifact / "model.safetensors" + if not artifact.is_dir() or not config_path.is_file() or not weights_path.is_file(): + raise FileNotFoundError( + "model.speaker_encoder.path must contain model_config.yaml and model.safetensors; " f"got {artifact}." + ) + if model.cfg.get("encoder_chunk_size_seconds", None) is not None: + raise ValueError( + "Independent per-encoder chunking requires model.encoder_chunk_size_seconds=null; " + "set model.speaker_encoder.asr_chunk_size_seconds and chunk_size_seconds instead." + ) + if not isinstance(model.perception.modality_adapter, IdentityConnector) or model.perception.rote is not None: + raise ValueError("IndependentDualEncoder requires IdentityConnector and rote=null.") + if "encoder_multilayer" in model.perception._modules: + raise ValueError("IndependentDualEncoder does not support multi-layer perception adapters.") + + speaker_config = OmegaConf.load(config_path) + speaker = model.perception.from_config_dict(speaker_config) + state = load_file(str(weights_path), device="cpu") + speaker.load_state_dict(state, strict=True) + + frame_shift_seconds = ( + model.perception.preprocessor.featurizer.hop_length / model.perception.preprocessor.featurizer.sample_rate + ) + dual = IndependentDualEncoder( + model.perception.encoder, + speaker, + frame_shift_seconds=frame_shift_seconds, + asr_chunk_size_seconds=cfg.get("asr_chunk_size_seconds", None), + auxiliary_chunk_size_seconds=cfg.get("chunk_size_seconds", None), + freeze_auxiliary=cfg.get("frozen", True), + ) + dual.auxiliary_encoder_config = OmegaConf.to_container(speaker_config, resolve=True) + + old_proj = model.perception.proj + if not isinstance(old_proj, torch.nn.Linear): + raise TypeError( + "IndependentDualEncoder currently requires the perception stack to end in nn.Linear; " + f"got {type(old_proj).__name__}." + ) + model.perception.encoder = dual + model.perception.proj = torch.nn.Linear( + dual.d_model, + old_proj.out_features, + bias=old_proj.bias is not None, + device=old_proj.weight.device, + dtype=old_proj.weight.dtype, + ) + with open_dict(model.cfg): + if "d_model" in model.cfg.perception.modality_adapter: + model.cfg.perception.modality_adapter.d_model = dual.d_model + + logging.info( + "Mounted independent speaker encoder from %s beside ASR encoder " + "(widths: ASR=%d speaker=%d combined=%d; chunks: ASR=%s speaker=%s seconds; frozen=%s).", + artifact, + IndependentDualEncoder._encoder_width(dual.asr_encoder), + IndependentDualEncoder._encoder_width(dual.auxiliary_encoder), + dual.d_model, + dual.asr_chunk_size_seconds, + dual.auxiliary_chunk_size_seconds, + dual.freeze_auxiliary, + ) def setup_parallel_expert_encoder(model: torch.nn.Module): @@ -369,6 +423,40 @@ def setup_parallel_expert_encoder(model: torch.nn.Module): strict=True, config_overrides=model.cfg.get("pe_encoder_overrides", None), ) + obsolete_chunk_keys = [ + key + for key in ("pe_asr_chunk_size_seconds", "pe_diar_chunk_size_seconds") + if model.cfg.get(key, None) is not None + ] + if obsolete_chunk_keys: + raise ValueError( + f"{', '.join(f'model.{key}' for key in obsolete_chunk_keys)} are no longer supported; " + "use model.encoder_chunk_size_seconds to chunk both ParallelExpertEncoder branches." + ) + + chunk_size = model.cfg.get("encoder_chunk_size_seconds", None) + if chunk_size is not None: + chunk_size = float(chunk_size) + if chunk_size <= 0: + raise ValueError(f"model.encoder_chunk_size_seconds must be positive or null, got {chunk_size}.") + if ( + model.cfg.get("packed_encoder_sequences", False) + and model.cfg.get("encoder_chunk_batch_size", None) is not None + ): + raise ValueError( + "model.encoder_chunk_batch_size is not supported for packed ParallelExpertEncoder internal chunking; " + "set it to null or disable model.packed_encoder_sequences to use outer waveform chunk batching." + ) + if not hasattr(pe_encoder, "chunk_size_seconds"): + raise TypeError( + f"{type(pe_encoder).__name__} does not support model.encoder_chunk_size_seconds; " + "missing runtime attribute 'chunk_size_seconds'." + ) + pe_encoder.chunk_size_seconds = chunk_size + if hasattr(pe_encoder, "_bundle_config"): + pe_encoder._bundle_config.chunk_size_seconds = chunk_size + logging.info("Set ParallelExpertEncoder chunk_size_seconds=%s", chunk_size) + if (spk_kernel_scale := model.cfg.get("spk_kernel_scale", None)) is not None: pe_encoder.spk_kernel_scale = float(spk_kernel_scale) @@ -598,7 +686,7 @@ def _load_checkpoint_state(checkpoint_path: str) -> dict: return load_file(os.path.join(checkpoint_path, "model.safetensors")) else: - return torch.load(checkpoint_path, map_location='cpu')['state_dict'] + return torch.load(checkpoint_path, map_location="cpu")["state_dict"] def init_perception_from_checkpoint(model: torch.nn.Module, checkpoint_path: str): @@ -655,6 +743,7 @@ def load_pretrained_model(model: torch.nn.Module, checkpoint_path: str): import gc import os + from nemo.utils import logging logging.info(f"Loading pretrained s2s model from {checkpoint_path}") @@ -668,7 +757,11 @@ def load_pretrained_model(model: torch.nn.Module, checkpoint_path: str): loaded_keys = [] missing_keys = [] - with safe_open(os.path.join(checkpoint_path, "model.safetensors"), framework="pt", device="cpu") as f: + with safe_open( + os.path.join(checkpoint_path, "model.safetensors"), + framework="pt", + device="cpu", + ) as f: available_keys = f.keys() for key in available_keys: if key in model_state_dict: diff --git a/scripts/dataloading/_validate_dataloader/full_mode.py b/scripts/dataloading/_validate_dataloader/full_mode.py index 71520c766c7a..c0ffe1807ccd 100644 --- a/scripts/dataloading/_validate_dataloader/full_mode.py +++ b/scripts/dataloading/_validate_dataloader/full_mode.py @@ -49,9 +49,15 @@ def build_validation_dataset(full_cfg, tokenizer, *, mode: str, section: str = " from nemo.collections.speechlm2.data.salm_dataset import SALMDataset data_cfg = full_cfg.get("data", {}) + model_cfg = full_cfg.get("model", {}) kwargs = {"tokenizer": tokenizer, "strict_audio_loading": True} if (multispeaker_cfg := data_cfg.get("multispeaker_cfg")) is not None: kwargs["multispeaker_cfg"] = multispeaker_cfg + pack_audio = bool(model_cfg.get("use_nemo_automodel", False) and model_cfg.get("packed_encoder_sequences", False)) + if pack_audio: + kwargs["pack_audio"] = True + if (batch_tokens := data_cfg.get(section, {}).get("batch_tokens")) is not None: + kwargs["batch_tokens"] = batch_tokens # FallbackDataset would replay the prior batch after a decode failure and # hide the exact error this validation mode is intended to detect. return SALMDataset(**kwargs) diff --git a/scripts/speechlm2/distributed_oomptimizer.py b/scripts/speechlm2/distributed_oomptimizer.py index 3a5f5d2b36f6..412366da6f65 100755 --- a/scripts/speechlm2/distributed_oomptimizer.py +++ b/scripts/speechlm2/distributed_oomptimizer.py @@ -81,6 +81,7 @@ """ import importlib +import inspect import json import math import os @@ -105,6 +106,23 @@ from nemo.utils.trainer_utils import resolve_trainer_cfg +def _run_training_step(model: pl.LightningModule, batch, batch_idx: int): + if hasattr(model, "_training_step_batch"): + return model._training_step_batch(batch, batch_idx) + + signature = inspect.signature(model.training_step) + params = [ + param + for param in signature.parameters.values() + if param.kind in (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD) + and param.default is inspect.Parameter.empty + ] + has_varargs = any(param.kind == inspect.Parameter.VAR_POSITIONAL for param in signature.parameters.values()) + if has_varargs or len(params) >= 2: + return model.training_step(batch, batch_idx) + return model.training_step(batch) + + class ProfilingBatchGenerator: """ ProfilingBatchGenerator is used to generate artificial mini-batches for model training @@ -1591,7 +1609,7 @@ def step(): try: click.echo(f"\tCurrent gap: {gen.current_rel_gap}... ", nl=False) optimizer.zero_grad() - out = model.training_step(batch, batch_idx) + out = _run_training_step(model, batch, batch_idx) out['loss'].sum().backward() optimizer.step() peak_allocated = torch.cuda.max_memory_allocated() @@ -1599,13 +1617,7 @@ def step(): oom = True status = "OOM!" except RuntimeError as e: - error_msg = str(e) - oom_like = ( - "cuFFT error: CUFFT_INTERNAL_ERROR" in error_msg - or "CUDA out of memory" in error_msg - or "CUDACachingAllocator" in error_msg - ) - if not oom_like: + if not _is_oom_like(e): raise oom = True status = "OOM!" @@ -1690,7 +1702,7 @@ def run(self) -> None: try: self.optimizer.zero_grad() batch = self.gen(self.seq_len_in, self.seq_len_out) - out = self.model.training_step(batch, batch_idx) + out = _run_training_step(self.model, batch, batch_idx) out['loss'].sum().backward() self.optimizer.step() torch.cuda.synchronize(self.device) diff --git a/tests/collections/asr/test_asr_modules.py b/tests/collections/asr/test_asr_modules.py index 2083204d5cab..f9887a157881 100644 --- a/tests/collections/asr/test_asr_modules.py +++ b/tests/collections/asr/test_asr_modules.py @@ -13,11 +13,25 @@ # See the License for the specific language governing permissions and # limitations under the License. +import copy +import random + import pytest import torch from omegaconf import OmegaConf from nemo.collections.asr import modules +from nemo.collections.asr.parts.packed_sequence import ( + PackedEncoderActivations, + pack_encoder_output, + split_encoder_output, +) +from nemo.collections.asr.parts.preprocessing.features import ( + FilterbankFeatures, + normalize_batch, + normalize_packed_batch, +) +from nemo.collections.asr.parts.submodules.subsampling import FeatureStacking from nemo.collections.asr.parts.utils.rnnt_utils import Hypothesis from nemo.core.utils import numba_utils from nemo.core.utils.numba_utils import __NUMBA_MINIMUM_VERSION__ @@ -25,6 +39,421 @@ class TestASRModulesBasicTests: + @pytest.mark.unit + @pytest.mark.parametrize( + ("exact_pad", "preemph", "normalize", "pad_to", "frame_splicing"), + [ + (False, 0.97, "per_feature", 0, 1), + (False, None, "all_features", 16, 1), + (True, 0.97, None, 0, 1), + (False, 0.97, "per_feature", 16, 2), + ], + ) + @pytest.mark.parametrize( + "device", + [ + "cpu", + pytest.param("cuda", marks=pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA unavailable")), + ], + ) + def test_AudioToMelSpectrogramPreprocessor_packed_waveform_matches_padded( + self, exact_pad, preemph, normalize, pad_to, frame_splicing, device + ): + preprocessor = ( + modules.AudioToMelSpectrogramPreprocessor( + normalize=normalize, + dither=0, + pad_to=pad_to, + exact_pad=exact_pad, + preemph=preemph, + frame_splicing=frame_splicing, + ) + .eval() + .to(device) + ) + lengths = torch.tensor([4096, 2500, 701], dtype=torch.long, device=device) + torch.manual_seed(7) + audios = torch.randn(3, int(lengths.max()), device=device) + for row, length in zip(audios, lengths): + row[int(length) :] = 0.0 + packed_audio_samples = torch.cat([row[: int(length)] for row, length in zip(audios, lengths)]) + audio_cu_seqlens = torch.cat([lengths.new_zeros(1), lengths.cumsum(dim=0, dtype=torch.long)]) + + expected, expected_lens = preprocessor(input_signal=audios, length=lengths) + actual = preprocessor.forward_packed( + input_signal=packed_audio_samples, + length=lengths, + input_signal_cu_seqlens=audio_cu_seqlens, + ) + + assert isinstance(actual, PackedEncoderActivations) + assert torch.equal(actual.lengths, expected_lens) + assert actual.padded_length == expected.shape[2] + assert actual.data.shape == (int(expected_lens.sum()), expected.shape[1]) + for row, (features, valid_length) in enumerate(zip(split_encoder_output(actual), expected_lens.tolist())): + torch.testing.assert_close( + features, + expected[row, :, :valid_length].transpose(0, 1), + rtol=1e-5, + atol=2e-6, + ) + + @pytest.mark.unit + def test_AudioToMelSpectrogramPreprocessor_packed_waveform_uses_one_vectorized_stft(self, monkeypatch): + preprocessor = modules.AudioToMelSpectrogramPreprocessor(dither=0, pad_to=16).eval() + lengths = torch.tensor([4096, 2500, 701], dtype=torch.long) + padded = torch.randn(3, int(lengths.max())) + packed = torch.cat([row[: int(length)] for row, length in zip(padded, lengths)]) + cu_seqlens = torch.cat([lengths.new_zeros(1), lengths.cumsum(dim=0)]) + calls = 0 + stft_input_shape = None + original_stft = preprocessor.featurizer.stft + + def count_stft(*args, **kwargs): + nonlocal calls, stft_input_shape + calls += 1 + stft_input_shape = args[0].shape + return original_stft(*args, **kwargs) + + def reject_dense_frontend(*args, **kwargs): + raise AssertionError("packed waveform preprocessing must not call the dense frontend") + + monkeypatch.setattr(preprocessor.featurizer, "stft", count_stft) + monkeypatch.setattr(preprocessor, "forward", reject_dense_frontend) + actual = preprocessor.forward_packed(packed, lengths, cu_seqlens) + + assert calls == 1 + assert stft_input_shape[0] == 1 + assert stft_input_shape[1] <= packed.numel() + lengths.numel() * ( + preprocessor.featurizer.n_fft + preprocessor.featurizer.hop_length + ) + assert actual.total_tokens == int(actual.lengths.sum()) + + @pytest.mark.unit + def test_AudioToMelSpectrogramPreprocessor_packed_waveform_isolates_boundaries(self): + preprocessor = modules.AudioToMelSpectrogramPreprocessor( + normalize="per_feature", dither=0, pad_to=0, preemph=0.97 + ).eval() + first = torch.zeros(1600) + first[-1] = 1.0 + second = torch.zeros(1000) + second[0] = -1.0 + lengths = torch.tensor([first.numel(), second.numel()], dtype=torch.long) + packed = torch.cat([first, second]) + cu_seqlens = torch.tensor([0, first.numel(), packed.numel()], dtype=torch.long) + + actual = preprocessor.forward_packed( + input_signal=packed, + length=lengths, + input_signal_cu_seqlens=cu_seqlens, + ) + + for row, (features, waveform) in enumerate(zip(split_encoder_output(actual), (first, second))): + expected, expected_lens = preprocessor(input_signal=waveform.unsqueeze(0), length=lengths[row : row + 1]) + assert actual.lengths[row] == expected_lens[0] + valid_length = int(expected_lens[0]) + torch.testing.assert_close(features, expected[0, :, :valid_length].transpose(0, 1), rtol=1e-5, atol=2e-6) + + @pytest.mark.unit + @pytest.mark.parametrize("exact_pad", [False, True]) + @pytest.mark.parametrize( + "device", + [ + "cpu", + pytest.param("cuda", marks=pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA unavailable")), + ], + ) + def test_AudioToMelSpectrogramPreprocessor_packed_waveform_preserves_zero_length_rows(self, exact_pad, device): + preprocessor = ( + modules.AudioToMelSpectrogramPreprocessor(dither=0, pad_to=0, exact_pad=exact_pad).eval().to(device) + ) + lengths = torch.tensor([0, 1600, 2400], dtype=torch.long, device=device) + generator = torch.Generator(device=device).manual_seed(1) + audios = torch.randn(3, 2400, device=device, generator=generator) + audios[0].zero_() + audios[1, 1600:].zero_() + packed = torch.cat([audios[1, :1600], audios[2]]) + cu_seqlens = torch.tensor([0, 0, 1600, 4000], dtype=torch.long, device=device) + + expected, expected_lens = preprocessor(input_signal=audios, length=lengths) + actual = preprocessor.forward_packed(packed, lengths, cu_seqlens) + + assert torch.equal(actual.lengths, expected_lens) + for row, features in enumerate(split_encoder_output(actual)): + valid_length = int(expected_lens[row]) + torch.testing.assert_close(features, expected[row, :, :valid_length].transpose(0, 1), rtol=1e-5, atol=3e-6) + + @pytest.mark.unit + def test_AudioToMelSpectrogramPreprocessor_packed_dither_is_seed_deterministic(self): + preprocessor = modules.AudioToMelSpectrogramPreprocessor( + normalize="per_feature", dither=1e-5, pad_to=0 + ).train() + lengths = torch.tensor([1600, 900], dtype=torch.long) + packed = torch.linspace(-0.5, 0.5, int(lengths.sum())) + cu_seqlens = torch.tensor([0, 1600, 2500], dtype=torch.long) + padded = torch.zeros(2, 1600) + padded[0] = packed[:1600] + padded[1, :900] = packed[1600:] + + torch.manual_seed(11) + expected, expected_lens = preprocessor(input_signal=padded.clone(), length=lengths) + torch.manual_seed(11) + first = preprocessor.forward_packed(packed.clone(), lengths, cu_seqlens) + torch.manual_seed(11) + second = preprocessor.forward_packed(packed.clone(), lengths, cu_seqlens) + + assert torch.equal(first.lengths, expected_lens) + assert torch.equal(first.lengths, second.lengths) + assert torch.equal(first.data, second.data) + for row, (features, valid_length) in enumerate(zip(split_encoder_output(first), expected_lens.tolist())): + # Padded batches consume dither RNG for padding positions, so later rows + # are tolerance-equivalent rather than bitwise-identical. + torch.testing.assert_close( + features, + expected[row, :, :valid_length].transpose(0, 1), + rtol=1e-3, + atol=5e-4, + ) + + @pytest.mark.unit + def test_AudioToMelSpectrogramPreprocessor_packed_waveform_validates_metadata(self): + preprocessor = modules.AudioToMelSpectrogramPreprocessor(dither=0, pad_to=0) + packed = torch.zeros(6) + lengths = torch.tensor([4, 2], dtype=torch.long) + + with pytest.raises(ValueError, match="must equal length"): + preprocessor.forward_packed(packed, lengths, torch.tensor([0, 3, 6])) + with pytest.raises(ValueError, match="ends at 6"): + preprocessor.forward_packed(packed[:5], lengths, torch.tensor([0, 4, 6])) + with pytest.raises(TypeError, match="integer dtype"): + preprocessor.forward_packed(packed, lengths, torch.tensor([0.0, 4.0, 6.0])) + + @pytest.mark.unit + def test_AudioToMelSpectrogramPreprocessor_packed_padding_value_matches_dense_feature_stacking(self): + preprocessor = modules.AudioToMelSpectrogramPreprocessor( + features=8, normalize="per_feature", dither=0, pad_to=16, pad_value=-3.0 + ).eval() + lengths = torch.tensor([4096, 2500, 701]) + audios = torch.randn(3, int(lengths.max())) + audios.masked_fill_(torch.arange(audios.shape[1])[None, :] >= lengths[:, None], 0.0) + packed_audio = torch.cat([row[: int(length)] for row, length in zip(audios, lengths)]) + audio_cu_seqlens = torch.cat([lengths.new_zeros(1), lengths.cumsum(0)]) + stacking = FeatureStacking(subsampling_factor=4, feat_in=8, feat_out=16).eval() + + dense_mels, mel_lengths = preprocessor(input_signal=audios, length=lengths) + packed_mels = preprocessor.forward_packed(packed_audio, lengths, audio_cu_seqlens) + dense_stacked, output_lengths = stacking(dense_mels, mel_lengths) + expected = pack_encoder_output(dense_stacked, output_lengths) + actual = stacking.forward_packed(packed_mels) + + assert packed_mels.padding_value == -3.0 + torch.testing.assert_close(actual.data, expected.data, rtol=1e-5, atol=2e-6) + + @pytest.mark.unit + @pytest.mark.parametrize("linear_spec", [False, True]) + def test_FilterbankFeatures_packed_all_empty_has_the_right_width(self, linear_spec): + featurizer = FilterbankFeatures(nfilt=23, frame_splicing=2, dither=0, pad_to=0).eval() + packed = featurizer.forward_packed( + torch.empty(0), torch.tensor([0, 0]), torch.tensor([0, 0, 0]), linear_spec=linear_spec + ) + + expected_width = featurizer.n_fft // 2 + 1 if linear_spec else 46 + assert packed.data.shape == (0, expected_width) + assert packed.lengths.tolist() == [0, 0] + assert packed.cu_seqlens.tolist() == [0, 0, 0] + + @pytest.mark.unit + def test_FilterbankFeatures_forward_packed_honors_use_grads(self): + lengths = torch.tensor([1600, 900]) + cu_seqlens = torch.cat([lengths.new_zeros(1), lengths.cumsum(0)]) + samples = torch.randn(int(lengths.sum()), requires_grad=True) + frozen = FilterbankFeatures(nfilt=8, normalize=None, dither=0, pad_to=0, use_grads=False).eval() + differentiable = FilterbankFeatures(nfilt=8, normalize=None, dither=0, pad_to=0, use_grads=True).eval() + + assert not frozen.forward_packed(samples, lengths, cu_seqlens).data.requires_grad + output = differentiable.forward_packed(samples, lengths, cu_seqlens) + output.data.sum().backward() + assert samples.grad is not None + assert torch.isfinite(samples.grad).all() + + @pytest.mark.unit + def test_FilterbankFeatures_packed_linear_spectrum_matches_nondefault_dense_stft(self): + featurizer = FilterbankFeatures( + n_window_size=400, + n_window_stride=137, + n_fft=512, + nfilt=8, + normalize=None, + dither=0, + pad_to=0, + ).eval() + lengths = torch.tensor([4097, 2503]) + audios = torch.randn(2, int(lengths.max())) + audios.masked_fill_(torch.arange(audios.shape[1])[None, :] >= lengths[:, None], 0.0) + packed_audio = torch.cat([row[: int(length)] for row, length in zip(audios, lengths)]) + cu_seqlens = torch.cat([lengths.new_zeros(1), lengths.cumsum(0)]) + + expected, expected_lengths = featurizer(audios, lengths, linear_spec=True) + actual = featurizer.forward_packed(packed_audio, lengths, cu_seqlens, linear_spec=True) + expected = pack_encoder_output(expected.transpose(1, 2), expected_lengths) + + torch.testing.assert_close(actual.data, expected.data, rtol=1e-5, atol=3e-6) + + @pytest.mark.unit + def test_FilterbankFeatures_packed_fixed_normalization_matches_dense(self): + normalize = { + "fixed_mean": [value for row in ([0.1] * 8, [-0.2] * 8) for value in row], + "fixed_std": [value for row in ([1.5] * 8, [0.75] * 8) for value in row], + } + featurizer = FilterbankFeatures(nfilt=8, normalize=normalize, dither=0, pad_to=0).eval() + lengths = torch.tensor([4096, 2500]) + audios = torch.randn(2, int(lengths.max())) + audios.masked_fill_(torch.arange(audios.shape[1])[None, :] >= lengths[:, None], 0.0) + packed_audio = torch.cat([row[: int(length)] for row, length in zip(audios, lengths)]) + cu_seqlens = torch.cat([lengths.new_zeros(1), lengths.cumsum(0)]) + + expected, expected_lengths = featurizer(audios, lengths) + actual = featurizer.forward_packed(packed_audio, lengths, cu_seqlens) + expected = pack_encoder_output(expected.transpose(1, 2), expected_lengths) + + torch.testing.assert_close(actual.data, expected.data, rtol=1e-5, atol=3e-6) + + @pytest.mark.unit + @pytest.mark.parametrize("use_vectorized_code", [False, True]) + def test_SpectrogramAugmentation_forward_packed_matches_dense_vectorized_masks(self, use_vectorized_code): + augment = modules.SpectrogramAugmentation( + freq_masks=2, + time_masks=2, + freq_width=3, + time_width=4, + rect_masks=1, + rect_time=3, + rect_freq=2, + rng=random.Random(13), + mask_value=-2.0, + use_vectorized_spec_augment=use_vectorized_code, + ).train() + lengths = torch.tensor([12, 7, 0]) + features = torch.randn(3, 8, 12) + features.masked_fill_(torch.arange(12)[None, None, :] >= lengths[:, None, None], 0.0) + packed = pack_encoder_output(features.transpose(1, 2), lengths) + packed_augment = copy.deepcopy(augment) + augment.spec_augment.use_vectorized_code = True + + torch.manual_seed(29) + expected = augment(input_spec=features.clone(), length=lengths) + torch.manual_seed(29) + actual = packed_augment.forward_packed(packed) + expected = pack_encoder_output(expected.transpose(1, 2), lengths) + + torch.testing.assert_close(actual.data, expected.data, rtol=0.0, atol=0.0) + + @pytest.mark.unit + def test_SpectrogramAugmentation_packed_cutout_preserves_frontend_padded_rng_range(self): + preprocessor = modules.AudioToMelSpectrogramPreprocessor( + features=8, normalize=None, dither=0, pad_to=16 + ).eval() + lengths = torch.tensor([1600, 900]) + audios = torch.randn(2, int(lengths.max())) + audios.masked_fill_(torch.arange(audios.shape[1])[None, :] >= lengths[:, None], 0.0) + packed_audio = torch.cat([row[: int(length)] for row, length in zip(audios, lengths)]) + cu_seqlens = torch.cat([lengths.new_zeros(1), lengths.cumsum(0)]) + dense_mels, mel_lengths = preprocessor(input_signal=audios, length=lengths) + packed_mels = preprocessor.forward_packed(packed_audio, lengths, cu_seqlens) + augment = modules.SpectrogramAugmentation( + freq_masks=0, + time_masks=0, + rect_masks=2, + rect_time=5, + rect_freq=2, + rng=random.Random(31), + ) + packed_augment = copy.deepcopy(augment) + + expected = augment(input_spec=dense_mels.clone(), length=mel_lengths) + actual = packed_augment.forward_packed(packed_mels) + expected = pack_encoder_output(expected.transpose(1, 2), mel_lengths) + + assert packed_mels.padded_length == dense_mels.shape[2] + assert torch.equal(actual.data == 0.0, expected.data == 0.0) + torch.testing.assert_close(actual.data, expected.data, rtol=1e-5, atol=2e-6) + + @pytest.mark.unit + def test_SpectrogramAugmentation_packed_cutout_accepts_all_empty_batch(self): + augment = modules.SpectrogramAugmentation(rect_masks=1, rect_time=5, rect_freq=2) + packed = PackedEncoderActivations( + torch.empty(0, 8), + torch.tensor([0, 0]), + torch.tensor([0, 0, 0], dtype=torch.int32), + 0, + ) + + assert augment.forward_packed(packed) is packed + + @pytest.mark.unit + @pytest.mark.parametrize( + "normalize_type", + [ + None, + "per_feature", + "all_features", + { + "fixed_mean": [value for row in ([0.1] * 4, [-0.2] * 4) for value in row], + "fixed_std": [value for row in ([1.5] * 4, [0.75] * 4) for value in row], + }, + ], + ) + def test_normalize_packed_batch_preserves_dense_partial_stack_padding(self, normalize_type): + lengths = torch.tensor([7, 4]) + features = torch.randn(2, 4, 7) + features.masked_fill_(torch.arange(7)[None, None, :] >= lengths[:, None, None], -3.0) + packed = PackedEncoderActivations( + pack_encoder_output(features.transpose(1, 2), lengths).data, + lengths, + torch.tensor([0, 7, 11], dtype=torch.int32), + 7, + padding_value=-3.0, + padded_length=7, + ) + stacking = FeatureStacking(subsampling_factor=3, feat_in=4, feat_out=6).eval() + + if normalize_type is None: + dense_normalized = features + packed_normalized = packed + else: + dense_normalized, _, _ = normalize_batch(features, lengths, normalize_type) + packed_normalized = normalize_packed_batch(packed, normalize_type) + expected, output_lengths = stacking(dense_normalized, lengths) + expected = pack_encoder_output(expected, output_lengths) + actual = stacking.forward_packed(packed_normalized) + + torch.testing.assert_close(actual.data, expected.data, rtol=1e-5, atol=2e-6) + + @pytest.mark.unit + def test_FilterbankFeatures_packed_narrowband_uses_configured_rng(self): + featurizer = FilterbankFeatures( + nfilt=8, + normalize=None, + dither=0, + pad_to=0, + rng=random.Random(17), + nb_augmentation_prob=0.5, + nb_max_freq=2000, + ).train() + lengths = torch.tensor([4096, 2500, 1600]) + audios = torch.randn(3, int(lengths.max())) + audios.masked_fill_(torch.arange(audios.shape[1])[None, :] >= lengths[:, None], 0.0) + packed_audio = torch.cat([row[: int(length)] for row, length in zip(audios, lengths)]) + cu_seqlens = torch.cat([lengths.new_zeros(1), lengths.cumsum(0)]) + state = featurizer._rng.getstate() + + expected, expected_lengths = featurizer(audios, lengths) + featurizer._rng.setstate(state) + actual = featurizer.forward_packed(packed_audio, lengths, cu_seqlens) + expected = pack_encoder_output(expected.transpose(1, 2), expected_lengths) + + torch.testing.assert_close(actual.data, expected.data, rtol=1e-5, atol=2e-6) + @pytest.mark.unit def test_AudioToMelSpectrogramPreprocessor_config(self): # Test that dataclass matches signature of module diff --git a/tests/collections/asr/test_asr_rope.py b/tests/collections/asr/test_asr_rope.py index 1edc3e9b1043..e928bda8b579 100644 --- a/tests/collections/asr/test_asr_rope.py +++ b/tests/collections/asr/test_asr_rope.py @@ -227,6 +227,18 @@ def test_non_contiguous_inputs(self): assert torch.allclose(q_rot_nc, q_rot_c, atol=1e-6) assert torch.allclose(k_rot_nc, k_rot_c, atol=1e-6) + @pytest.mark.run_only_on('GPU') + @pytest.mark.unit + def test_cache_is_identical_across_construction_devices(self): + cpu = RotaryPositionalEncoding(d_k=64, max_len=8192) + cpu.extend_pe(8192, device=torch.device('cpu'), dtype=torch.bfloat16) + + cuda = RotaryPositionalEncoding(d_k=64, max_len=8192).to('cuda') + cuda.extend_pe(8192, device=torch.device('cuda'), dtype=torch.bfloat16) + + assert torch.equal(cuda.cos.cpu(), cpu.cos) + assert torch.equal(cuda.sin.cpu(), cpu.sin) + class TestRoPEMultiHeadAttention: @pytest.mark.unit diff --git a/tests/collections/asr/test_packed_normalization_precision.py b/tests/collections/asr/test_packed_normalization_precision.py new file mode 100644 index 000000000000..91e963e47439 --- /dev/null +++ b/tests/collections/asr/test_packed_normalization_precision.py @@ -0,0 +1,76 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Focused regression tests for low-precision packed normalization statistics.""" + +from __future__ import annotations + +import pytest +import torch + +from nemo.collections.asr.parts.packed_sequence import pack_encoder_output +from nemo.collections.asr.parts.preprocessing.features import normalize_batch, normalize_packed_batch + + +def make_features(device: str, dtype: torch.dtype) -> tuple[torch.Tensor, torch.Tensor]: + torch.manual_seed(2746317213) + lengths = torch.tensor([1200, 2500, 4500], dtype=torch.int64, device=device) + time = torch.arange(4500, dtype=torch.float32, device=device) + channels = torch.arange(128, dtype=torch.float32, device=device).unsqueeze(1) + features = torch.zeros((3, 128, 4500), dtype=dtype, device=device) + for index, length in enumerate(lengths.tolist()): + row = ( + -12.0 + + channels * 0.015 + + 2.0 * torch.sin(time.unsqueeze(0) * (0.003 + channels * 0.00001) + index) + + 1.5 * torch.randn((128, 4500), device=device) + ) + features[index, :, :length] = row[:, :length].to(dtype) + return features, lengths + + +@pytest.mark.parametrize("normalize_type", ["per_feature", "all_features"]) +@pytest.mark.parametrize( + "device", + [ + "cpu", + pytest.param("cuda", marks=pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA unavailable")), + ], +) +def test_low_precision_packed_statistics_match_fp32_reference(device, normalize_type): + features, lengths = make_features(device, torch.bfloat16) + packed = pack_encoder_output(features.transpose(1, 2), lengths) + + actual = normalize_packed_batch(packed, normalize_type) + expected, _, _ = normalize_batch(features.float(), lengths, normalize_type) + expected = pack_encoder_output(expected.transpose(1, 2), lengths).data.to(torch.bfloat16) + + assert actual.data.dtype == torch.bfloat16 + torch.testing.assert_close(actual.data, expected, rtol=1e-2, atol=1.5625e-2) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA unavailable") +def test_legacy_bf16_segmented_statistics_are_not_silently_equivalent_to_fp32_fix(): + features, lengths = make_features("cuda", torch.bfloat16) + packed = pack_encoder_output(features.transpose(1, 2), lengths) + stable = normalize_packed_batch(packed, "per_feature") + sequence_ids = torch.repeat_interleave(torch.arange(3, device="cuda"), lengths) + denominator = lengths.unsqueeze(1) + mean = torch.segment_reduce(packed.data, "sum", lengths=lengths, unsafe=True) / denominator + centered = packed.data - mean[sequence_ids] + variance = torch.segment_reduce(centered.square(), "sum", lengths=lengths, unsafe=True) / (denominator - 1) + std = torch.sqrt(variance).masked_fill(variance.isnan(), 0.0) + 1e-5 + legacy = centered / std[sequence_ids] + + cosine = torch.nn.functional.cosine_similarity(stable.data.float().flatten(), legacy.float().flatten(), dim=0) + assert cosine < 0.9 diff --git a/tests/collections/asr/test_packed_pee_grouped.py b/tests/collections/asr/test_packed_pee_grouped.py new file mode 100644 index 000000000000..ae0ab9ea858f --- /dev/null +++ b/tests/collections/asr/test_packed_pee_grouped.py @@ -0,0 +1,60 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch + +from nemo.collections.asr.parts.packed_sequence import PackedEncoderActivations, pack_encoder_output +from tests.collections.asr.test_parallel_expert_encoder_two_branch import ( + _MEL_FEATURES, + _N_SPK, + build_toy_packed_pe_encoder, +) + + +@pytest.mark.unit +def test_canonical_pee_packed_path_matches_dense_without_legacy_grouped_runtime(): + torch.manual_seed(0) + encoder = build_toy_packed_pe_encoder().eval() + mels = torch.randn(2, _MEL_FEATURES, 40) + lengths = torch.tensor([40, 17]) + targets = torch.zeros(2, 5, _N_SPK) + + with torch.no_grad(): + dense, dense_lengths = encoder(mels, lengths, spk_targets=targets) + packed = encoder.forward_sequence_packed(mels, lengths, spk_targets=targets) + + valid = torch.arange(dense.shape[-1])[None, :] < dense_lengths[:, None] + torch.testing.assert_close(packed.data, dense.transpose(1, 2)[valid], rtol=1e-5, atol=1e-6) + assert not hasattr(encoder, "pee") + assert not hasattr(encoder, "sequence_packed_execution_mode") + + +@pytest.mark.unit +def test_canonical_pee_accepts_token_flat_mels(): + torch.manual_seed(0) + encoder = build_toy_packed_pe_encoder().eval() + mels = torch.randn(2, _MEL_FEATURES, 40) + lengths = torch.tensor([40, 17]) + mels.masked_fill_(torch.arange(mels.shape[-1])[None, None, :] >= lengths[:, None, None], 0.0) + packed_mels = pack_encoder_output(mels.transpose(1, 2), lengths) + targets = torch.zeros(2, 5, _N_SPK) + + with torch.no_grad(): + dense_input = encoder.forward_sequence_packed(mels, lengths, spk_targets=targets) + packed_input = encoder.forward_sequence_packed(packed_mels, spk_targets=targets) + + assert isinstance(packed_input, PackedEncoderActivations) + assert torch.equal(packed_input.lengths, dense_input.lengths) + torch.testing.assert_close(packed_input.data, dense_input.data, rtol=1e-5, atol=1e-6) diff --git a/tests/collections/asr/test_packed_sequence_review_coverage.py b/tests/collections/asr/test_packed_sequence_review_coverage.py new file mode 100644 index 000000000000..9ac24cfb0d8c --- /dev/null +++ b/tests/collections/asr/test_packed_sequence_review_coverage.py @@ -0,0 +1,179 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import io +import tarfile + +import pytest +import torch +from omegaconf import OmegaConf + +from nemo.collections.asr.modules.parallel_expert_encoder import ParallelExpertEncoderPT +from nemo.collections.asr.modules.transformer_encoder import TransformerEncoder +from nemo.collections.asr.parts.packed_sequence import pack_encoder_output, unpack_encoder_output +from tests.collections.asr.test_parallel_expert_encoder_two_branch import ( + _MEL_FEATURES, + _N_SPK, + build_toy_packed_pe_encoder, + toy_packed_diarization_model_cfg, + toy_transformer_asr_encoder_cfg, +) + + +def test_sequence_packed_training_dropout_is_finite_and_reproducible_within_path(): + encoder = TransformerEncoder( + feat_in=8, + d_model=32, + n_heads=2, + n_layers=2, + subsampling_factor=2, + drop_rate=0.2, + dropout_pre_encoder=0.2, + dropout_emb=0.2, + self_attention_model="rope", + sync_max_audio_length=False, + ).train() + audio = torch.randn(3, 8, 12) + lengths = torch.tensor([12, 7, 3]) + + with torch.no_grad(): + torch.manual_seed(17) + first = encoder.forward_sequence_packed(audio, lengths).data + torch.manual_seed(17) + second = encoder.forward_sequence_packed(audio, lengths).data + + assert torch.isfinite(first).all() + torch.testing.assert_close(first, second) + + +def test_synthetic_canonical_pee_nemo_archive_loads_strictly_and_enables_packed_path(tmp_path): + torch.manual_seed(0) + source = build_toy_packed_pe_encoder().eval() + cfg = OmegaConf.create( + { + "target": "nemo.collections.asr.modules.parallel_expert_encoder.ParallelExpertEncoderPT", + "asr_encoder_type": "transformer", + "asr_encoder_cfg": toy_transformer_asr_encoder_cfg(), + "diarization_model_cfg": toy_packed_diarization_model_cfg(), + "asr_normalize_type": "per_feature", + "speaker_feature_config_version": 1, + "speaker_feature_mode": "continuous", + "speaker_activity_threshold": None, + "sync_max_audio_length": False, + } + ) + config_bytes = OmegaConf.to_yaml(cfg).encode() + weights = io.BytesIO() + torch.save({f"encoder.{key}": value for key, value in source.state_dict().items()}, weights) + archive = tmp_path / "synthetic_canonical_pee.nemo" + with tarfile.open(archive, "w") as tar: + config_info = tarfile.TarInfo("model_config.yaml") + config_info.size = len(config_bytes) + tar.addfile(config_info, io.BytesIO(config_bytes)) + weight_bytes = weights.getvalue() + weight_info = tarfile.TarInfo("model_weights.ckpt") + weight_info.size = len(weight_bytes) + tar.addfile(weight_info, io.BytesIO(weight_bytes)) + + restored = ParallelExpertEncoderPT.load_from_nemo(str(archive), strict=True).eval() + + assert set(restored.state_dict()) == set(source.state_dict()) + for key, value in source.state_dict().items(): + torch.testing.assert_close(restored.state_dict()[key], value) + with torch.no_grad(): + packed = restored.forward_sequence_packed( + torch.randn(2, _MEL_FEATURES, 24), + torch.tensor([24, 11]), + ) + assert packed.total_tokens == int(packed.lengths.sum()) + + +@pytest.mark.parametrize("target_mode", ["none", "mixed", "external"]) +def test_pee_packed_fusion_matches_dense_routing_modes(target_mode): + torch.manual_seed(0) + encoder = build_toy_packed_pe_encoder().eval() + mels = torch.randn(3, _MEL_FEATURES, 40) + lengths = torch.tensor([40, 23, 9]) + targets = None + if target_mode != "none": + targets = torch.zeros(3, 5, _N_SPK) + targets[0, :, 0] = 1.0 + targets[2, :, 2] = 1.0 + if target_mode == "mixed": + targets[1] = -1.0 + + with torch.no_grad(): + legacy, output_lengths = encoder(mels, lengths, spk_targets=targets) + packed = encoder.forward_sequence_packed(mels, lengths, spk_targets=targets) + + restored = unpack_encoder_output(packed, total_length=legacy.shape[-1]) + valid = torch.arange(legacy.shape[-1])[None, :] < output_lengths[:, None] + torch.testing.assert_close(restored[valid], legacy.transpose(1, 2)[valid], rtol=1e-4, atol=1e-5) + + +def test_pee_packed_speaker_threshold_edges_match_legacy(): + encoder = build_toy_packed_pe_encoder(speaker_activity_threshold=0.5).eval() + lengths = torch.tensor([3, 2]) + padded = torch.randn(2, 3, encoder.d_model) + packed = pack_encoder_output(padded, lengths) + threshold = encoder.speaker_activity_threshold + targets = torch.full((2, 3, _N_SPK), threshold) + targets[:, 0, 0] = threshold - 1e-6 + targets[:, 1, 1] = threshold + 1e-6 + + with torch.no_grad(): + legacy = encoder._fuse_diar_and_asr(padded.transpose(1, 2), targets).transpose(1, 2) + compact = encoder._fuse_diar_and_asr_packed(packed, targets) + + restored = unpack_encoder_output(compact, total_length=3) + valid = torch.arange(3)[None, :] < lengths[:, None] + torch.testing.assert_close(restored[valid], legacy[valid]) + + +def test_pee_packed_rejects_mismatched_branch_metadata(monkeypatch): + encoder = build_toy_packed_pe_encoder().eval() + diar = pack_encoder_output(torch.randn(2, 3, _N_SPK), torch.tensor([3, 1])) + asr = pack_encoder_output(torch.randn(2, 3, encoder.d_model), torch.tensor([3, 2])) + monkeypatch.setattr(encoder, "_run_diarization_packed", lambda features: diar) + monkeypatch.setattr(encoder, "_run_asr_packed", lambda features: asr) + + with pytest.raises(RuntimeError, match="metadata diverged"): + encoder.forward_sequence_packed(torch.randn(2, _MEL_FEATURES, 8), torch.tensor([8, 4])) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="PEE packed gradient parity requires CUDA") +def test_pee_packed_matches_dense_input_and_parameter_gradients(): + torch.manual_seed(0) + dense_encoder = build_toy_packed_pe_encoder(freeze_asr=False, freeze_diar=True).cuda().eval() + packed_encoder = build_toy_packed_pe_encoder(freeze_asr=False, freeze_diar=True).cuda().eval() + packed_encoder.load_state_dict(dense_encoder.state_dict(), strict=True) + dense_mels = torch.randn(2, _MEL_FEATURES, 32, device="cuda", requires_grad=True) + packed_mels = dense_mels.detach().clone().requires_grad_() + lengths = torch.tensor([32, 17], device="cuda") + targets = torch.zeros(2, 4, _N_SPK, device="cuda") + targets[0, :, 0] = 1.0 + + dense, output_lengths = dense_encoder(dense_mels, lengths, spk_targets=targets) + packed = packed_encoder.forward_sequence_packed(packed_mels, lengths, spk_targets=targets) + valid = torch.arange(dense.shape[-1], device="cuda")[None, :] < output_lengths[:, None] + dense.transpose(1, 2)[valid].float().square().mean().backward() + packed.data.float().square().mean().backward() + + torch.testing.assert_close(packed_mels.grad, dense_mels.grad, rtol=2e-3, atol=2e-4) + for name, dense_parameter in dense_encoder.named_parameters(): + if not name.startswith(("asr_encoder.", "asr_norm.")) or not dense_parameter.requires_grad: + continue + packed_grad = dict(packed_encoder.named_parameters())[name].grad + assert dense_parameter.grad is not None and packed_grad is not None + torch.testing.assert_close(packed_grad, dense_parameter.grad, rtol=2e-3, atol=2e-4) diff --git a/tests/collections/asr/test_packed_sequence_round2.py b/tests/collections/asr/test_packed_sequence_round2.py new file mode 100644 index 000000000000..7bb58ff72bfd --- /dev/null +++ b/tests/collections/asr/test_packed_sequence_round2.py @@ -0,0 +1,109 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch +from torch.utils._pytree import tree_flatten + +from nemo.collections.asr.parts.packed_sequence import PackedEncoderActivations, pack_encoder_output +from tests.collections.asr.test_parallel_expert_encoder_two_branch import ( + _MEL_FEATURES, + _N_SPK, + build_toy_packed_pe_encoder, +) + + +def test_packed_encoder_activations_is_registered_as_pytree(): + packed = pack_encoder_output(torch.randn(2, 4, 3), torch.tensor([4, 2])) + + leaves, _ = tree_flatten(packed) + + assert all(leaf is not packed for leaf in leaves) + assert any(leaf is packed.data for leaf in leaves) + assert any(leaf is packed.lengths for leaf in leaves) + assert any(leaf is packed.cu_seqlens for leaf in leaves) + + +def test_packed_output_with_data_reuses_validated_metadata_and_preserves_gradients(): + packed = pack_encoder_output(torch.randn(2, 4, 3), torch.tensor([4, 2])) + replacement = torch.randn(6, 5, requires_grad=True) + + updated = packed.with_data(replacement) + + assert updated.lengths is packed.lengths + assert updated.cu_seqlens is packed.cu_seqlens + assert updated.max_seqlen == packed.max_seqlen + updated.data.square().sum().backward() + assert replacement.grad is not None + with pytest.raises(ValueError, match="replacement data"): + packed.with_data(torch.randn(5, 3)) + + +def test_canonical_pee_packed_output_preserves_compact_metadata(): + torch.manual_seed(0) + encoder = build_toy_packed_pe_encoder().eval() + mels = torch.randn(2, _MEL_FEATURES, 24) + lengths = torch.tensor([24, 11]) + targets = torch.zeros(2, 3, _N_SPK) + + with torch.no_grad(): + output = encoder.forward_sequence_packed(mels, lengths, spk_targets=targets) + + assert isinstance(output, PackedEncoderActivations) + assert output.total_tokens == int(output.lengths.sum()) + assert output.cu_seqlens.tolist() == [0, *output.lengths.cumsum(0).tolist()] + + +def test_canonical_pee_dense_contract_is_unchanged_after_packed_use(): + encoder = build_toy_packed_pe_encoder().eval() + mels = torch.randn(2, _MEL_FEATURES, 24) + lengths = torch.tensor([24, 11]) + targets = torch.zeros(2, 3, _N_SPK) + state_keys = set(encoder.state_dict()) + + with torch.no_grad(): + packed = encoder.forward_sequence_packed(mels, lengths, spk_targets=targets) + dense, dense_lengths = encoder(mels, lengths, spk_targets=targets) + + restored = torch.cat( + [dense[index, :, : int(length)].transpose(0, 1) for index, length in enumerate(dense_lengths)] + ) + torch.testing.assert_close(packed.data, restored, rtol=1e-5, atol=1e-6) + assert set(encoder.state_dict()) == state_keys + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="PEE ASR-gradient parity requires CUDA") +def test_canonical_pee_packed_matches_dense_trainable_asr_gradients(): + torch.manual_seed(0) + dense_encoder = build_toy_packed_pe_encoder(freeze_asr=False, freeze_diar=True).cuda().eval() + packed_encoder = build_toy_packed_pe_encoder(freeze_asr=False, freeze_diar=True).cuda().eval() + packed_encoder.load_state_dict(dense_encoder.state_dict(), strict=True) + dense_mels = torch.randn(2, _MEL_FEATURES, 32, device="cuda", requires_grad=True) + packed_mels = dense_mels.detach().clone().requires_grad_() + lengths = torch.tensor([32, 17], device="cuda") + targets = torch.zeros(2, 4, _N_SPK, device="cuda") + + dense, output_lengths = dense_encoder(dense_mels, lengths, spk_targets=targets) + packed = packed_encoder.forward_sequence_packed(packed_mels, lengths, spk_targets=targets) + valid = torch.arange(dense.shape[-1], device="cuda")[None, :] < output_lengths[:, None] + dense.transpose(1, 2)[valid].float().square().mean().backward() + packed.data.float().square().mean().backward() + + torch.testing.assert_close(packed_mels.grad, dense_mels.grad, rtol=2e-3, atol=2e-4) + for name, dense_parameter in dense_encoder.named_parameters(): + if not name.startswith(("asr_encoder.", "asr_norm.")) or not dense_parameter.requires_grad: + continue + packed_grad = dict(packed_encoder.named_parameters())[name].grad + assert dense_parameter.grad is not None and packed_grad is not None + torch.testing.assert_close(packed_grad, dense_parameter.grad, rtol=2e-3, atol=2e-4) diff --git a/tests/collections/asr/test_packed_transformer_encoder.py b/tests/collections/asr/test_packed_transformer_encoder.py new file mode 100644 index 000000000000..83c70eb3f435 --- /dev/null +++ b/tests/collections/asr/test_packed_transformer_encoder.py @@ -0,0 +1,605 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import builtins +import copy +import gc +from types import SimpleNamespace + +import pytest +import torch +from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import checkpoint_wrapper + +from nemo.collections.asr.modules import transformer_encoder_utils as transformer_encoder_utils_module +from nemo.collections.asr.modules.transformer_encoder import ( + MultiHeadAttention, + TransformerEncoder, + TransformerEncoderConfig, +) +from nemo.collections.asr.parts.packed_sequence import ( + PackedEncoderActivations, + pack_encoder_output, + packed_encoder_position_ids, + split_encoder_output, + split_packed_data, + unpack_encoder_output, +) + + +def _supports_cuda_varlen_flash_attention(): + if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 8: + return False + return transformer_encoder_utils_module._get_flash_attention_varlen() is not None + + +requires_cuda_varlen_flash_attention = pytest.mark.skipif( + not _supports_cuda_varlen_flash_attention(), reason="requires an SM80+ CUDA varlen FlashAttention provider" +) + + +def test_packed_encoder_output_round_trip_and_positions(): + padded = torch.arange(3 * 5 * 4, dtype=torch.float32).reshape(3, 5, 4)[..., ::2] + assert not padded.is_contiguous() + lengths = torch.tensor([5, 2, 0]) + + packed = pack_encoder_output(padded, lengths) + + assert packed.data.shape == (7, 2) + assert packed.lengths.dtype == torch.int64 + assert packed.cu_seqlens.dtype == torch.int32 + assert packed.cu_seqlens.tolist() == [0, 5, 7, 7] + assert packed.max_seqlen == 5 + assert packed_encoder_position_ids(packed).tolist() == [0, 1, 2, 3, 4, 0, 1] + assert [part.shape[0] for part in split_encoder_output(packed)] == [5, 2, 0] + torch.testing.assert_close( + unpack_encoder_output(packed), padded * (torch.arange(5)[None, :, None] < lengths[:, None, None]) + ) + + +def test_packed_encoder_output_all_empty_is_differentiable(): + padded = torch.randn(2, 0, 4, requires_grad=True) + packed = pack_encoder_output(padded, torch.zeros(2, dtype=torch.long)) + + assert packed.data.shape == (0, 4) + assert packed.cu_seqlens.tolist() == [0, 0, 0] + assert packed.max_seqlen == 0 + packed.data.sum().backward() + assert padded.grad is not None + + +def test_split_packed_data_preserves_empty_rows_and_validates_integer_metadata(): + data = torch.arange(6) + lengths = torch.tensor([2, 0, 4]) + rows = split_packed_data(data, lengths, torch.tensor([0, 2, 2, 6])) + + assert [row.tolist() for row in rows] == [[0, 1], [], [2, 3, 4, 5]] + with pytest.raises(TypeError, match="integer dtype"): + split_packed_data(data, lengths.float(), torch.tensor([0, 2, 2, 6])) + with pytest.raises(TypeError, match="integer dtype"): + split_packed_data(data, lengths, torch.tensor([0.0, 2.0, 2.0, 6.0])) + + +@pytest.mark.parametrize( + ("lengths", "match"), + [ + (torch.tensor([1]), "shape"), + (torch.tensor([-1, 1]), "between"), + (torch.tensor([3, 1]), "between"), + (torch.tensor([1.0, 1.0]), "integer dtype"), + (torch.tensor([True, False]), "integer dtype"), + ], +) +def test_pack_encoder_output_rejects_invalid_lengths(lengths, match): + padded = torch.randn(2, 2, 4) + with pytest.raises((TypeError, ValueError), match=match): + pack_encoder_output(padded, lengths) + + +def test_pack_encoder_output_uses_prevalidated_metadata_constructor(monkeypatch): + from nemo.collections.asr.parts import packed_sequence as packed_sequence_module + + def fail_if_revalidated(*args, **kwargs): + raise AssertionError("pack_encoder_output must not revalidate CUDA metadata") + + monkeypatch.setattr(packed_sequence_module, "_validate_packed_encoder_activations", fail_if_revalidated) + packed = pack_encoder_output(torch.randn(2, 3, 4), torch.tensor([3, 1])) + + assert packed.cu_seqlens.tolist() == [0, 3, 4] + assert packed.max_seqlen == 3 + + +def test_packed_encoder_output_validates_manual_metadata(): + with pytest.raises(ValueError, match="differences equal to lengths"): + PackedEncoderActivations( + data=torch.randn(3, 4), + lengths=torch.tensor([1, 2]), + cu_seqlens=torch.tensor([0, 2, 3], dtype=torch.int32), + max_seqlen=2, + ) + + +def test_packed_encoder_output_rejects_noncontiguous_cu_seqlens(): + cu_seqlens = torch.tensor([0, 99, 1, 99, 3, 99], dtype=torch.int32)[::2] + assert not cu_seqlens.is_contiguous() + with pytest.raises(ValueError, match="cu_seqlens must be contiguous"): + PackedEncoderActivations( + data=torch.randn(3, 4), + lengths=torch.tensor([1, 2]), + cu_seqlens=cu_seqlens, + max_seqlen=2, + ) + + +def test_packed_encoder_output_pytree_transforms_do_not_revalidate_temporary_leaves(): + packed = pack_encoder_output(torch.randn(2, 3, 4), torch.tensor([3, 1])) + + leaves, spec = torch.utils._pytree.tree_flatten(packed) + restored = torch.utils._pytree.tree_unflatten(leaves, spec) + placeholder = torch.utils._pytree.tree_map(lambda _: None, packed) + + assert restored.data is packed.data + assert restored.lengths is packed.lengths + assert restored.cu_seqlens is packed.cu_seqlens + assert restored.max_seqlen == packed.max_seqlen + assert restored.padding_value == packed.padding_value + assert restored.padded_length == packed.padded_length + assert placeholder.data is None + assert placeholder.lengths is None + assert placeholder.cu_seqlens is None + assert placeholder.max_seqlen is None + assert placeholder.padding_value is None + assert placeholder.padded_length is None + + +def _make_encoder(*, position: str, attention: str = "full", qk_norm: bool = False, rotary_fraction: float = 1.0): + return TransformerEncoder( + feat_in=8, + d_model=32, + n_heads=2, + n_layers=2, + subsampling_factor=2, + drop_rate=0.0, + dropout_pre_encoder=0.0, + dropout_emb=0.0, + self_attention_model=position, + attn_mode=attention, + qk_norm=qk_norm, + qkv_bias=True, + rotary_fraction=rotary_fraction, + sync_max_audio_length=False, + ).eval() + + +@pytest.mark.parametrize("position", ["rope", "abs_pos", "no_pos", "rel_pos"]) +@pytest.mark.parametrize("attention", ["full", "causal"]) +@pytest.mark.parametrize("qk_norm", [False, True]) +def test_sequence_packed_matches_padded_valid_outputs_cpu(position, attention, qk_norm): + kwargs = {"rotary_fraction": 0.5} if position == "rope" else {} + torch.manual_seed(0) + encoder = _make_encoder(position=position, attention=attention, qk_norm=qk_norm, **kwargs) + audio = torch.randn(3, 8, 12) + lengths = torch.tensor([12, 7, 4]) + + with torch.no_grad(): + padded, output_lengths = encoder(audio, lengths) + packed = encoder.forward_sequence_packed(audio, lengths) + + restored = unpack_encoder_output(packed, total_length=padded.shape[-1]) + valid = torch.arange(padded.shape[-1])[None, :] < output_lengths[:, None] + torch.testing.assert_close(restored[valid], padded.transpose(1, 2)[valid], rtol=1e-5, atol=1e-6) + assert packed.total_tokens == int(output_lengths.sum()) + assert encoder.layers[0].attn._last_sequence_packed_backend == "flex_attention_reference" + assert encoder.layers[0].attn._last_sequence_packed_provider is None + + +@pytest.mark.parametrize("position", ["rope", "abs_pos", "no_pos", "rel_pos"]) +def test_sequence_packed_accepts_token_flat_features_without_dense_pre_encode(position): + torch.manual_seed(0) + encoder = _make_encoder(position=position) + audio = torch.randn(3, 8, 12) + lengths = torch.tensor([12, 7, 4]) + audio.masked_fill_(torch.arange(audio.shape[-1])[None, None, :] >= lengths[:, None, None], 0.0) + packed_features = pack_encoder_output(audio.transpose(1, 2), lengths) + + with torch.no_grad(): + from_dense = encoder.forward_sequence_packed(audio, lengths) + from_packed = encoder.forward_sequence_packed(packed_features, lengths) + + assert from_packed.total_tokens == int(from_packed.lengths.sum()) + assert torch.equal(from_packed.lengths, from_dense.lengths) + torch.testing.assert_close(from_packed.data, from_dense.data, rtol=1e-5, atol=1e-6) + + +def test_sequence_packed_token_flat_feature_outputs_and_gradients_match_dense_input(): + torch.manual_seed(0) + dense_encoder = _make_encoder(position="rope").train() + packed_encoder = copy.deepcopy(dense_encoder) + lengths = torch.tensor([12, 7, 4]) + dense_features = torch.randn(3, 8, 12) + dense_features.masked_fill_(torch.arange(12)[None, None, :] >= lengths[:, None, None], 0.0) + dense_features.requires_grad_() + packed_source = pack_encoder_output(dense_features.detach().transpose(1, 2), lengths) + packed_features = PackedEncoderActivations( + packed_source.data.clone().requires_grad_(), + packed_source.lengths, + packed_source.cu_seqlens, + packed_source.max_seqlen, + ) + + dense_output = dense_encoder.forward_sequence_packed(dense_features, lengths) + packed_output = packed_encoder.forward_sequence_packed(packed_features, lengths) + dense_output.data.square().mean().backward() + packed_output.data.square().mean().backward() + + torch.testing.assert_close(packed_output.data, dense_output.data, rtol=1e-5, atol=1e-6) + dense_valid_grads = pack_encoder_output(dense_features.grad.transpose(1, 2), lengths).data + torch.testing.assert_close(packed_features.data.grad, dense_valid_grads, rtol=1e-5, atol=1e-6) + for (dense_name, dense_parameter), (packed_name, packed_parameter) in zip( + dense_encoder.named_parameters(), packed_encoder.named_parameters() + ): + assert dense_name == packed_name + if dense_parameter.grad is not None: + torch.testing.assert_close(packed_parameter.grad, dense_parameter.grad, rtol=1e-5, atol=1e-6) + + +def test_sequence_packed_token_flat_feature_stacking_preserves_activation_checkpointing(): + encoder = _make_encoder(position="rope").train() + encoder.pre_encode = checkpoint_wrapper(encoder.pre_encode) + lengths = torch.tensor([12, 7]) + padded = torch.randn(2, 12, 8) + padded.masked_fill_(torch.arange(12)[None, :, None] >= lengths[:, None, None], 0.0) + packed = pack_encoder_output(padded, lengths) + packed = PackedEncoderActivations( + packed.data.detach().clone().requires_grad_(), packed.lengths, packed.cu_seqlens, packed.max_seqlen + ) + + output = encoder.forward_sequence_packed(packed, lengths) + output.data.square().mean().backward() + + assert packed.data.grad is not None + assert encoder.pre_encode._checkpoint_wrapped_module.proj.weight.grad is not None + + +@pytest.mark.parametrize(("position", "qk_norm"), [("rope", True), ("rel_pos", False)]) +@pytest.mark.parametrize("fused_qkv", [False, True]) +def test_sequence_packed_all_empty_keeps_attention_parameters_in_backward(position, qk_norm, fused_qkv): + encoder = _make_encoder(position=position, qk_norm=qk_norm).train() + inputs = torch.empty(2, 0, encoder.d_model, requires_grad=True) + + packed = encoder.forward_sequence_packed( + inputs, + torch.zeros(2, dtype=torch.long), + bypass_pre_encode=True, + fused_qkv=fused_qkv, + ) + packed.data.sum().backward() + + assert inputs.grad is not None + for layer in encoder.layers: + for name, parameter in layer.attn.named_parameters(): + assert parameter.grad is not None, f"missing gradient for attention parameter {name}" + assert torch.count_nonzero(parameter.grad) == 0 + + +def test_sequence_packed_boundaries_isolate_other_utterances_and_causal_future(): + torch.manual_seed(0) + encoder = _make_encoder(position="rope", attention="causal") + encoded = torch.randn(2, 6, encoder.d_model) + lengths = torch.tensor([6, 4]) + changed = encoded.clone() + changed[0, 4:] = torch.randn_like(changed[0, 4:]) * 100 + changed[1] = torch.randn_like(changed[1]) * 100 + + with torch.no_grad(): + original = encoder.forward_sequence_packed(encoded, lengths, bypass_pre_encode=True) + mutated = encoder.forward_sequence_packed(changed, lengths, bypass_pre_encode=True) + + torch.testing.assert_close(original.data[:4], mutated.data[:4], rtol=1e-5, atol=1e-6) + + +def test_sequence_packed_fused_qkv_matches_default_and_preserves_checkpoint_keys(): + torch.manual_seed(0) + encoder = _make_encoder(position="rope", qk_norm=True) + inputs = torch.randn(2, 6, encoder.d_model) + lengths = torch.tensor([6, 3]) + state_keys = set(encoder.state_dict()) + + with torch.no_grad(): + independent = encoder.forward_sequence_packed(inputs, lengths, bypass_pre_encode=True) + fused = encoder.forward_sequence_packed(inputs, lengths, bypass_pre_encode=True, fused_qkv=True) + + torch.testing.assert_close(fused.data, independent.data, rtol=1e-5, atol=1e-6) + assert set(encoder.state_dict()) == state_keys + + +def test_sequence_packed_layers_receive_only_valid_tokens(monkeypatch): + encoder = _make_encoder(position="rope") + encoded = torch.randn(3, 7, encoder.d_model) + lengths = torch.tensor([7, 3, 1]) + observed = [] + original = encoder.layers[0].ffn.forward + + def record(x): + observed.append(tuple(x.shape)) + return original(x) + + monkeypatch.setattr(encoder.layers[0].ffn, "forward", record) + with torch.no_grad(): + encoder.forward_sequence_packed(encoded, lengths, bypass_pre_encode=True) + + assert observed == [(11, encoder.d_model)] + + +def test_sequence_packed_varlen_dispatch_contract(monkeypatch): + cfg = TransformerEncoderConfig(d_model=32, n_heads=2, self_attention_model="no_pos", qkv_bias=False) + attention = MultiHeadAttention(cfg) + recorded = {} + + def fake_flash(q, k, v, cu_q, cu_k, max_q, max_k, **kwargs): + recorded.update(q=q, k=k, v=v, cu_q=cu_q, cu_k=cu_k, max_q=max_q, max_k=max_k, kwargs=kwargs) + return v + + monkeypatch.setattr( + transformer_encoder_utils_module, + "_select_flash_attention_varlen", + lambda q, *, static_eligible: fake_flash, + ) + lengths = torch.tensor([3, 2]) + cu_seqlens = torch.tensor([0, 3, 5], dtype=torch.int32) + + output = attention.forward_sequence_packed( + torch.randn(5, 32), + lengths=lengths, + cu_seqlens=cu_seqlens, + max_seqlen=3, + causal=True, + sequence_offsets=(0, 3, 5), + ) + + assert output.shape == (5, 32) + for name in ("q", "k", "v"): + assert recorded[name].shape == (5, 2, 16) + assert recorded[name].is_contiguous() + assert recorded["cu_q"] is recorded["cu_k"] is cu_seqlens + assert recorded["cu_q"].dtype == torch.int32 and recorded["cu_q"].is_contiguous() + assert recorded["max_q"] == recorded["max_k"] == 3 + assert recorded["kwargs"] == {"dropout_p": 0.0, "softmax_scale": None, "causal": True} + assert attention._last_sequence_packed_provider == "external" + + +def test_sequence_packed_flash_device_probe_is_cached(monkeypatch): + provider = object() + capability_probes = [] + provider_probes = [] + tensor = SimpleNamespace( + is_cuda=True, + dtype=torch.bfloat16, + shape=(1, 2, 16), + device=torch.device("cuda:0"), + ) + + def get_device_capability(device): + capability_probes.append(device) + return (9, 0) + + def get_provider(): + provider_probes.append(None) + return provider + + monkeypatch.setattr(torch.version, "cuda", "12.8") + monkeypatch.setattr(torch.cuda, "get_device_capability", get_device_capability) + monkeypatch.setattr(transformer_encoder_utils_module, "_get_flash_attention_varlen", get_provider) + transformer_encoder_utils_module._get_flash_attention_varlen_for_device.cache_clear() + try: + assert transformer_encoder_utils_module._can_use_flash_attention_varlen_layout(tensor, head_dim=16) + assert transformer_encoder_utils_module._can_use_flash_attention_varlen_layout(tensor, head_dim=16) + assert capability_probes == [torch.device("cuda:0")] + assert provider_probes == [None] + finally: + transformer_encoder_utils_module._get_flash_attention_varlen_for_device.cache_clear() + + +def test_flash_attention_varlen_aten_provider_is_reported(monkeypatch): + if getattr(torch.ops.aten, "_flash_attention_forward", None) is None: + pytest.skip("PyTorch build does not expose the ATen FlashAttention operator") + original_import = builtins.__import__ + + def import_without_external_flash(name, *args, **kwargs): + if name == "flash_attn": + raise ImportError("simulate flash-attn not installed") + return original_import(name, *args, **kwargs) + + transformer_encoder_utils_module._get_flash_attention_varlen.cache_clear() + monkeypatch.setattr(builtins, "__import__", import_without_external_flash) + try: + provider = transformer_encoder_utils_module._get_flash_attention_varlen() + assert provider is not None + assert provider._sequence_packed_provider == "aten" + finally: + transformer_encoder_utils_module._get_flash_attention_varlen.cache_clear() + + +def test_sequence_packed_adds_no_state_dict_keys_and_loads_strictly(): + encoder = _make_encoder(position="rope") + before = set(encoder.state_dict()) + with torch.no_grad(): + encoder.forward_sequence_packed(torch.randn(2, 5, encoder.d_model), torch.tensor([5, 2]), True) + after = set(encoder.state_dict()) + clone = _make_encoder(position="rope") + + result = clone.load_state_dict(encoder.state_dict(), strict=True) + + assert before == after + assert result.missing_keys == [] + assert result.unexpected_keys == [] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="FlexAttention backward requires CUDA") +def test_sequence_packed_supports_activation_checkpoint_wrapped_layers(): + encoder = _make_encoder(position="rope").to(device="cuda", dtype=torch.bfloat16).train() + state_keys = set(encoder.state_dict()) + for idx, layer in enumerate(encoder.layers): + encoder.layers[idx] = checkpoint_wrapper(layer) + inputs = torch.randn(2, 8, 10, device="cuda", dtype=torch.bfloat16, requires_grad=True) + + packed = encoder.forward_sequence_packed(inputs, torch.tensor([10, 4], device="cuda")) + packed.data.square().mean().backward() + + assert inputs.grad is not None + assert encoder.layers[0]._checkpoint_wrapped_module.attn.w_qkv.weight.grad is not None + assert set(encoder.state_dict()) == state_keys + + +@requires_cuda_varlen_flash_attention +def test_sequence_packed_thd_cuda_matches_padded_outputs_and_gradients(): + torch.manual_seed(0) + padded_encoder = _make_encoder(position="rope", qk_norm=True).to(device="cuda", dtype=torch.bfloat16).train() + packed_encoder = copy.deepcopy(padded_encoder) + padded_input = torch.randn(3, 12, padded_encoder.d_model, device="cuda", dtype=torch.bfloat16, requires_grad=True) + packed_input = padded_input.detach().clone().requires_grad_() + lengths = torch.tensor([12, 7, 3], device="cuda") + + padded, output_lengths = padded_encoder(padded_input, lengths, bypass_pre_encode=True) + packed = packed_encoder.forward_sequence_packed(packed_input, lengths, bypass_pre_encode=True) + restored = unpack_encoder_output(packed, total_length=padded.shape[-1]) + valid = torch.arange(padded.shape[-1], device="cuda")[None, :] < output_lengths[:, None] + padded_valid = padded.transpose(1, 2)[valid] + packed_valid = restored[valid] + padded_valid.float().square().mean().backward() + packed_valid.float().square().mean().backward() + + torch.testing.assert_close(packed_valid, padded_valid, rtol=2e-2, atol=2e-2) + torch.testing.assert_close(packed_input.grad, padded_input.grad, rtol=3e-2, atol=3e-2) + for name in ( + "layers.0.attn.w_qkv.weight", + "layers.0.attn.out_proj.weight", + "layers.0.attn.q_norm.weight", + "layers.0.ffn.net.0.weight", + ): + torch.testing.assert_close( + dict(packed_encoder.named_parameters())[name].grad, + dict(padded_encoder.named_parameters())[name].grad, + rtol=3e-2, + atol=3e-2, + ) + assert packed_encoder.layers[0].attn._last_sequence_packed_backend == "flash_attention_varlen" + assert packed_encoder.layers[0].attn._last_sequence_packed_provider in {"aten", "external"} + assert packed.data.shape == (int(lengths.sum()), packed_encoder.d_model) + + +@requires_cuda_varlen_flash_attention +@pytest.mark.parametrize( + ("dtype", "position", "attention", "qk_norm", "rotary_fraction"), + [ + (torch.float16, "rope", "full", False, 0.5), + (torch.bfloat16, "rope", "causal", False, 1.0), + (torch.bfloat16, "abs_pos", "causal", True, 1.0), + (torch.float16, "no_pos", "full", True, 1.0), + ], +) +def test_sequence_packed_cuda_fast_path_matrix(dtype, position, attention, qk_norm, rotary_fraction): + torch.manual_seed(0) + encoder = _make_encoder( + position=position, + attention=attention, + qk_norm=qk_norm, + rotary_fraction=rotary_fraction, + ).to(device="cuda", dtype=dtype) + inputs = torch.randn(3, 12, 32, device="cuda", dtype=dtype) + lengths = torch.tensor([12, 0, 5], device="cuda") + + with torch.no_grad(): + padded, output_lengths = encoder(inputs, lengths, bypass_pre_encode=True) + packed = encoder.forward_sequence_packed(inputs, lengths, bypass_pre_encode=True) + + restored = unpack_encoder_output(packed, total_length=padded.shape[-1]) + valid = torch.arange(padded.shape[-1], device="cuda")[None, :] < output_lengths[:, None] + torch.testing.assert_close(restored[valid], padded.transpose(1, 2)[valid], rtol=3e-2, atol=3e-2) + assert encoder.layers[0].attn._last_sequence_packed_backend == "flash_attention_varlen" + assert encoder.layers[0].attn._last_sequence_packed_provider in {"aten", "external"} + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA FlexAttention") +def test_sequence_packed_cuda_fp32_relative_position_fallback_gradients(): + torch.manual_seed(0) + padded_encoder = _make_encoder(position="rel_pos").cuda().train() + packed_encoder = copy.deepcopy(padded_encoder) + padded_input = torch.randn(2, 8, 32, device="cuda", requires_grad=True) + packed_input = padded_input.detach().clone().requires_grad_() + lengths = torch.tensor([8, 3], device="cuda") + + padded, _ = padded_encoder(padded_input, lengths, bypass_pre_encode=True) + packed = packed_encoder.forward_sequence_packed(packed_input, lengths, bypass_pre_encode=True) + valid = torch.arange(8, device="cuda")[None, :] < lengths[:, None] + padded.transpose(1, 2)[valid].square().mean().backward() + packed.data.square().mean().backward() + + torch.testing.assert_close(packed_input.grad[valid], padded_input.grad[valid], rtol=2e-4, atol=2e-5) + for suffix in ("linear_pos.weight", "pos_bias_u", "pos_bias_v"): + packed_grad = dict(packed_encoder.named_parameters())[f"layers.0.attn.{suffix}"].grad + padded_grad = dict(padded_encoder.named_parameters())[f"layers.0.attn.{suffix}"].grad + torch.testing.assert_close(packed_grad, padded_grad, rtol=2e-4, atol=2e-5) + assert packed_encoder.layers[0].attn._last_sequence_packed_backend == "flex_attention_reference" + + +@requires_cuda_varlen_flash_attention +def test_sequence_packed_reduces_forward_backward_peak_memory_for_uneven_batch(): + torch.manual_seed(0) + encoder = TransformerEncoder( + feat_in=64, + d_model=128, + n_heads=4, + n_layers=2, + subsampling_factor=1, + drop_rate=0.0, + dropout_pre_encoder=0.0, + dropout_emb=0.0, + self_attention_model="rope", + sync_max_audio_length=False, + ).to(device="cuda", dtype=torch.bfloat16) + encoder.train() + source = torch.randn(4, 64, 256, device="cuda", dtype=torch.bfloat16) + lengths = torch.tensor([256, 64, 32, 16], device="cuda") + + def run(sequence_packed: bool, measure: bool = False) -> int: + gc.collect() + torch.cuda.empty_cache() + encoder.zero_grad(set_to_none=True) + inputs = source.detach().clone().requires_grad_() + torch.cuda.synchronize() + baseline = torch.cuda.memory_allocated() + if measure: + torch.cuda.reset_peak_memory_stats() + if sequence_packed: + output = encoder.forward_sequence_packed(inputs, lengths).data + else: + output = encoder(inputs, lengths)[0] + output.float().sum().backward() + torch.cuda.synchronize() + peak = torch.cuda.max_memory_allocated() - baseline if measure else 0 + del output, inputs + return peak + + # Warm both kernel families so compilation and allocator setup are excluded. + run(sequence_packed=False) + run(sequence_packed=True) + padded_peak = run(sequence_packed=False, measure=True) + packed_peak = run(sequence_packed=True, measure=True) + + assert encoder.layers[0].attn._last_sequence_packed_backend == "flash_attention_varlen" + assert encoder.layers[0].attn._last_sequence_packed_provider in {"aten", "external"} + assert packed_peak < padded_peak * 0.7, ( + f"Expected native THD to materially reduce peak activation memory; " + f"padded={padded_peak:,} bytes, packed={packed_peak:,} bytes." + ) diff --git a/tests/collections/asr/test_parallel_expert_encoder_two_branch.py b/tests/collections/asr/test_parallel_expert_encoder_two_branch.py new file mode 100644 index 000000000000..8f318f6b54a2 --- /dev/null +++ b/tests/collections/asr/test_parallel_expert_encoder_two_branch.py @@ -0,0 +1,681 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import pytest +import torch +import torch.distributed as dist +from omegaconf import DictConfig, OmegaConf +from torch import nn + +import nemo.collections.asr.modules.parallel_expert_encoder as pee_module +from nemo.collections.asr.models import SortformerEncLabelModel +from nemo.collections.asr.modules.conformer_encoder import ConformerEncoder +from nemo.collections.asr.modules.parallel_expert_encoder import ( + ParallelExpertEncoder, + _clone_config, + _default_dtype, + _disable_dist_feature_sync, +) +from nemo.collections.asr.modules.transformer_encoder import TransformerEncoder +from nemo.collections.asr.parts.packed_sequence import pack_encoder_output, unpack_encoder_output +from nemo.collections.asr.parts.preprocessing.features import normalize_batch, normalize_packed_batch + +_MEL_FEATURES = 128 +_ASR_D_MODEL = 32 +_DIAR_FC_D_MODEL = 32 +_DIAR_TF_D_MODEL = 16 +_N_SPK = 4 +_SUBSAMPLING_FACTOR = 8 + + +def toy_asr_encoder_cfg() -> DictConfig: + return DictConfig( + { + "_target_": "nemo.collections.asr.modules.ConformerEncoder", + "feat_in": _MEL_FEATURES, + "feat_out": -1, + "n_layers": 1, + "d_model": _ASR_D_MODEL, + "subsampling": "dw_striding", + "subsampling_factor": _SUBSAMPLING_FACTOR, + "subsampling_conv_channels": 16, + "ff_expansion_factor": 4, + "self_attention_model": "rel_pos", + "n_heads": 4, + "att_context_size": [-1, -1], + "conv_kernel_size": 9, + "dropout": 0.0, + "dropout_pre_encoder": 0.0, + "dropout_emb": 0.0, + "dropout_att": 0.0, + } + ) + + +def toy_transformer_asr_encoder_cfg() -> DictConfig: + return DictConfig( + { + "_target_": "nemo.collections.asr.modules.transformer_encoder.TransformerEncoder", + "feat_in": _MEL_FEATURES, + "d_model": _ASR_D_MODEL, + "n_heads": 2, + "n_layers": 1, + "subsampling": "feature_stacking", + "subsampling_factor": _SUBSAMPLING_FACTOR, + "drop_rate": 0.0, + "dropout_pre_encoder": 0.0, + "dropout_emb": 0.0, + "qkv_bias": False, + "qk_norm": True, + "ff_expansion": 2.0, + "pre_block_norm": True, + "self_attention_model": "rope", + "attn_mode": "full", + "sync_max_audio_length": False, + } + ) + + +def toy_diarization_model_cfg() -> DictConfig: + defaults = {"fc_d_model": _DIAR_FC_D_MODEL, "tf_d_model": _DIAR_TF_D_MODEL} + return DictConfig( + { + "target": "nemo.collections.asr.models.sortformer_diar_models.SortformerEncLabelModel", + "sample_rate": 16000, + "pil_weight": 0.5, + "ats_weight": 0.5, + "max_num_of_spks": _N_SPK, + "streaming_mode": False, + "async_streaming": False, + "model_defaults": DictConfig(defaults), + "preprocessor": DictConfig( + { + "_target_": "nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor", + "normalize": "per_feature", + "window_size": 0.025, + "sample_rate": 16000, + "window_stride": 0.01, + "window": "hann", + "features": _MEL_FEATURES, + "n_fft": 512, + "frame_splicing": 1, + "dither": 0.00001, + } + ), + "encoder": DictConfig( + { + "_target_": "nemo.collections.asr.modules.ConformerEncoder", + "feat_in": _MEL_FEATURES, + "feat_out": -1, + "n_layers": 1, + "d_model": _DIAR_FC_D_MODEL, + "subsampling": "dw_striding", + "subsampling_factor": _SUBSAMPLING_FACTOR, + "subsampling_conv_channels": 16, + "causal_downsampling": False, + "ff_expansion_factor": 4, + "self_attention_model": "rel_pos", + "n_heads": 4, + "att_context_size": [-1, -1], + "conv_kernel_size": 9, + "conv_norm_type": "batch_norm", + "dropout": 0.0, + "dropout_pre_encoder": 0.0, + "dropout_emb": 0.0, + "dropout_att": 0.0, + } + ), + "transformer_encoder": DictConfig( + { + "_target_": "nemo.collections.asr.modules.transformer.transformer_encoders.TransformerEncoder", + "num_layers": 1, + "hidden_size": _DIAR_TF_D_MODEL, + "inner_size": 32, + "num_attention_heads": 4, + "attn_score_dropout": 0.0, + "attn_layer_dropout": 0.0, + "ffn_dropout": 0.0, + "hidden_act": "relu", + "pre_ln": False, + "pre_ln_final_layer_norm": True, + } + ), + "sortformer_modules": DictConfig( + { + "_target_": "nemo.collections.asr.modules.sortformer_modules.SortformerModules", + "num_spks": _N_SPK, + "dropout_rate": 0.0, + "fc_d_model": _DIAR_FC_D_MODEL, + "tf_d_model": _DIAR_TF_D_MODEL, + } + ), + "loss": DictConfig( + { + "_target_": "nemo.collections.asr.losses.bce_loss.BCELoss", + "weight": None, + "reduction": "mean", + } + ), + } + ) + + +def toy_packed_diarization_model_cfg() -> DictConfig: + cfg = toy_diarization_model_cfg() + cfg.encoder = toy_transformer_asr_encoder_cfg() + cfg.encoder.d_model = _DIAR_FC_D_MODEL + cfg.encoder.qk_norm = False + cfg.transformer_encoder.num_layers = 0 + cfg.transformer_encoder.pre_ln = False + return cfg + + +def build_toy_pe_encoder(**overrides) -> ParallelExpertEncoder: + kwargs = { + "asr_encoder_cfg": toy_asr_encoder_cfg(), + "diarization_model_cfg": toy_diarization_model_cfg(), + "asr_normalize_type": "per_feature", + "online_inference_length": 500, + } + kwargs.update(overrides) + return ParallelExpertEncoder(**kwargs) + + +def build_toy_packed_pe_encoder(**overrides) -> ParallelExpertEncoder: + """Construct the canonical PEE with native packed-capable branch encoders.""" + kwargs = { + "asr_encoder_type": "transformer", + "asr_encoder_cfg": toy_transformer_asr_encoder_cfg(), + "diarization_model_cfg": toy_packed_diarization_model_cfg(), + } + kwargs.update(overrides) + return build_toy_pe_encoder(**kwargs) + + +@pytest.mark.unit +def test_clone_config_is_deep_and_handles_none(): + config = OmegaConf.create({"a": {"b": 1}}) + clone = _clone_config(config) + clone.a.b = 2 + assert config.a.b == 1 + assert _clone_config(None) is None + + +@pytest.mark.unit +@pytest.mark.parametrize("target_dtype", [torch.float64, torch.float16]) +def test_default_dtype_sets_and_restores(target_dtype): + previous = torch.get_default_dtype() + with _default_dtype(target_dtype): + assert torch.get_default_dtype() == target_dtype + assert torch.get_default_dtype() == previous + + +@pytest.mark.unit +def test_disable_dist_feature_sync_noop_when_uninitialized(): + assert not dist.is_initialized() + original = dist.is_initialized + with _disable_dist_feature_sync(): + pass + assert dist.is_initialized is original + + +@pytest.mark.unit +def test_static_helpers_align_and_cast(): + diar = torch.arange(9, dtype=torch.float32).reshape(1, 3, 3) + aligned = ParallelExpertEncoder._align_diar_frames(diar, 5) + assert aligned.shape == (1, 5, 3) + assert torch.equal(aligned[:, -1], diar[:, -1]) + + module = nn.Linear(4, 4).to(torch.float64) + cast = ParallelExpertEncoder._match_module_io(torch.zeros(2, 4), module) + assert cast.dtype == torch.float64 + + +@pytest.mark.unit +def test_pe_encoder_builds_two_real_branches_and_freezes_diarizer(): + encoder = build_toy_pe_encoder() + assert isinstance(encoder.asr_encoder, ConformerEncoder) + assert encoder.asr_encoder_type == "fastconformer" + assert isinstance(encoder.diarization_model, SortformerEncLabelModel) + assert encoder.d_model == _ASR_D_MODEL + assert encoder.subsampling_factor == _SUBSAMPLING_FACTOR + assert encoder.n_spk == _N_SPK + assert encoder.diar_normalize_type == "per_feature" + assert all(not parameter.requires_grad for parameter in encoder.diarization_model.parameters()) + assert any(parameter.requires_grad for parameter in encoder.asr_encoder.parameters()) + + encoder.train() + assert encoder.asr_encoder.training + assert not encoder.diarization_model.training + + +@pytest.mark.unit +def test_pe_encoder_selects_native_transformer_asr_branch(): + encoder = build_toy_pe_encoder( + asr_encoder_type="transformer", + asr_encoder_cfg=toy_transformer_asr_encoder_cfg(), + ) + assert isinstance(encoder.asr_encoder, TransformerEncoder) + assert encoder.asr_encoder_type == "transformer" + assert encoder.d_model == _ASR_D_MODEL + assert encoder.subsampling_factor == _SUBSAMPLING_FACTOR + + +@pytest.mark.unit +@pytest.mark.parametrize( + ("asr_encoder_type", "asr_encoder_cfg", "expected_class"), + [ + ("fastconformer", toy_transformer_asr_encoder_cfg, "ConformerEncoder"), + ("transformer", toy_asr_encoder_cfg, "TransformerEncoder"), + ], +) +def test_pe_encoder_rejects_asr_encoder_type_config_mismatch(asr_encoder_type, asr_encoder_cfg, expected_class): + with pytest.raises(TypeError, match=rf"requires .*{expected_class}"): + build_toy_pe_encoder( + asr_encoder_type=asr_encoder_type, + asr_encoder_cfg=asr_encoder_cfg(), + ) + + +@pytest.mark.unit +def test_pe_encoder_rejects_unknown_asr_encoder_type(): + with pytest.raises(ValueError, match="asr_encoder_type must be one of"): + build_toy_pe_encoder(asr_encoder_type="auto") + + +@pytest.mark.unit +def test_freeze_asr_keeps_both_frozen_branches_in_eval(): + encoder = build_toy_pe_encoder(freeze_asr=True) + encoder.train() + assert not encoder.asr_encoder.training + assert not encoder.diarization_model.training + assert all(not parameter.requires_grad for parameter in encoder.asr_encoder.parameters()) + + +@pytest.mark.unit +def test_pe_encoder_rejects_incompatible_branch_frame_rates(): + diarization_config = toy_diarization_model_cfg() + diarization_config.encoder.subsampling_factor = 4 + with pytest.raises(ValueError, match="embedded diarization encoder subsampling factor"): + build_toy_pe_encoder(diarization_model_cfg=diarization_config) + + +@pytest.mark.unit +@pytest.mark.parametrize( + ("asr_encoder_type", "asr_encoder_cfg"), + [ + ("fastconformer", toy_asr_encoder_cfg), + ("transformer", toy_transformer_asr_encoder_cfg), + ], +) +def test_offline_forward_runs_both_branches(asr_encoder_type, asr_encoder_cfg): + encoder = build_toy_pe_encoder( + asr_encoder_type=asr_encoder_type, + asr_encoder_cfg=asr_encoder_cfg(), + ).eval() + mels = torch.randn(2, _MEL_FEATURES, 160) + lengths = torch.tensor([160, 137]) + with torch.no_grad(): + output, output_lengths = encoder(mels, lengths) + assert output.shape[:2] == (2, _ASR_D_MODEL) + assert output.shape[-1] == int(output_lengths.max()) + assert torch.isfinite(output).all() + + +@pytest.mark.unit +def test_offline_sortformer_receives_per_feature_normalized_mels(monkeypatch): + encoder = build_toy_pe_encoder().eval() + mels = 4.0 * torch.randn(2, _MEL_FEATURES, 80) + 17.0 + lengths = torch.tensor([80, 53]) + observed = [] + original_frontend = encoder.diarization_model.frontend_encoder + + def tracked_frontend(**kwargs): + observed.append(kwargs["processed_signal"].detach().clone()) + return original_frontend(**kwargs) + + monkeypatch.setattr(encoder.diarization_model, "frontend_encoder", tracked_frontend) + with torch.no_grad(): + encoder._run_diarization(mels, lengths) + + expected, _, _ = normalize_batch(mels, lengths, normalize_type="per_feature") + torch.testing.assert_close(observed[0], expected) + + +@pytest.mark.unit +def test_mixed_missing_rttm_rows_use_sortformer_predictions(monkeypatch): + encoder = build_toy_pe_encoder().eval() + mels = torch.randn(3, _MEL_FEATURES, 80) + lengths = torch.tensor([80, 72, 64]) + diarization = torch.rand(3, 10, _N_SPK) + asr_states = torch.randn(3, _ASR_D_MODEL, 10) + asr_lengths = torch.tensor([10, 9, 8]) + monkeypatch.setattr(encoder, "_run_diarization", lambda *_: diarization) + monkeypatch.setattr(encoder, "_run_asr", lambda *_: (asr_states, asr_lengths)) + + targets = torch.zeros(3, 10, _N_SPK) + targets[0, :, 0] = 1.0 + targets[1] = encoder.missing_rttm_target + targets[2, :, 2] = 1.0 + expected = encoder._fuse_diar_and_asr( + asr_states, + targets, + diarization_preds=diarization, + use_diarization=torch.tensor([False, True, False]), + ) + actual, actual_lengths = encoder(mels, lengths, spk_targets=targets) + torch.testing.assert_close(actual, expected) + assert torch.equal(actual_lengths, asr_lengths) + + +@pytest.mark.unit +@pytest.mark.parametrize(("training", "world_size"), [(True, 1), (False, 2)]) +def test_all_rttm_rows_still_run_sortformer_in_collective_safe_paths(monkeypatch, training, world_size): + encoder = build_toy_pe_encoder().train(training) + monkeypatch.setattr(dist, "is_available", lambda: True) + monkeypatch.setattr(dist, "is_initialized", lambda: world_size > 1) + monkeypatch.setattr(dist, "get_world_size", lambda: world_size) + mels = torch.randn(2, _MEL_FEATURES, 80) + lengths = torch.tensor([80, 64]) + diarization = torch.rand(2, 10, _N_SPK) + asr_states = torch.randn(2, _ASR_D_MODEL, 10) + asr_lengths = torch.tensor([10, 8]) + diarization_calls = 0 + + def run_diarization(*_): + nonlocal diarization_calls + diarization_calls += 1 + return diarization + + monkeypatch.setattr(encoder, "_run_diarization", run_diarization) + monkeypatch.setattr(encoder, "_run_asr", lambda *_: (asr_states, asr_lengths)) + + targets = torch.zeros(2, 10, _N_SPK) + targets[0, :, 0] = 1.0 + targets[1, :, 1] = 1.0 + expected = encoder._fuse_diar_and_asr(asr_states, targets) + actual, actual_lengths = encoder(mels, lengths, spk_targets=targets) + + assert diarization_calls == 1 + torch.testing.assert_close(actual, expected) + assert torch.equal(actual_lengths, asr_lengths) + + +@pytest.mark.unit +def test_all_rttm_rows_can_skip_sortformer_in_single_process_eval(monkeypatch): + encoder = build_toy_pe_encoder().eval() + mels = torch.randn(2, _MEL_FEATURES, 80) + lengths = torch.tensor([80, 64]) + asr_states = torch.randn(2, _ASR_D_MODEL, 10) + asr_lengths = torch.tensor([10, 8]) + monkeypatch.setattr( + encoder, + "_run_diarization", + lambda *_: (_ for _ in ()).throw(AssertionError("single-process eval unexpectedly ran Sortformer")), + ) + monkeypatch.setattr(encoder, "_run_asr", lambda *_: (asr_states, asr_lengths)) + + targets = torch.zeros(2, 10, _N_SPK) + targets[0, :, 0] = 1.0 + targets[1, :, 1] = 1.0 + expected = encoder._fuse_diar_and_asr(asr_states, targets) + actual, actual_lengths = encoder(mels, lengths, spk_targets=targets) + + torch.testing.assert_close(actual, expected) + assert torch.equal(actual_lengths, asr_lengths) + + +@pytest.mark.unit +def test_all_rttm_rows_still_run_sortformer_in_packed_training_path(monkeypatch): + encoder = build_toy_pe_encoder().train() + lengths = torch.tensor([80, 64]) + packed_input = pack_encoder_output(torch.randn(2, 80, _MEL_FEATURES), lengths) + output_lengths = torch.tensor([10, 8]) + asr_states = pack_encoder_output(torch.randn(2, 10, _ASR_D_MODEL), output_lengths) + diarization = pack_encoder_output(torch.rand(2, 10, _N_SPK), output_lengths) + diarization_calls = 0 + + def run_diarization(_): + nonlocal diarization_calls + diarization_calls += 1 + return diarization + + monkeypatch.setattr(encoder, "_run_diarization_packed", run_diarization) + monkeypatch.setattr(encoder, "_run_asr_packed", lambda _: asr_states) + + targets = torch.zeros(2, 10, _N_SPK) + targets[0, :, 0] = 1.0 + targets[1, :, 1] = 1.0 + expected = encoder._fuse_diar_and_asr_packed(asr_states, targets) + actual = encoder.forward_sequence_packed(packed_input, spk_targets=targets) + + assert diarization_calls == 1 + torch.testing.assert_close(actual.data, expected.data) + assert torch.equal(actual.lengths, expected.lengths) + + +@pytest.mark.unit +def test_speaker_threshold_and_kernel_scale_are_preserved(): + encoder = build_toy_pe_encoder(speaker_activity_threshold=0.5, spk_kernel_scale=0.25).eval() + asr_states = torch.randn(1, _ASR_D_MODEL, 3) + targets = torch.full((1, 3, _N_SPK), 0.5) + targets[:, 1, 0] = 0.5001 + fused = encoder._fuse_diar_and_asr(asr_states, targets) + + normalized = encoder.asr_norm(asr_states.transpose(1, 2)) + binary = (targets > 0.5).to(normalized.dtype) + infusion = encoder.diar_norm(binary) @ encoder.diar_kernel + expected = (normalized + 0.25 * infusion).transpose(1, 2) + torch.testing.assert_close(fused, expected) + + +@pytest.mark.unit +def test_high_resolution_diarization_is_pooled_to_asr_grid(): + diarization_config = toy_diarization_model_cfg() + diarization_config.high_resolution = True + diarization_config.output_subsampling_factor = 1 + encoder = build_toy_pe_encoder(diarization_model_cfg=diarization_config).eval() + lengths = torch.tensor([80, 53]) + mels = torch.randn(2, _MEL_FEATURES, 80) + packed_input = pack_encoder_output(mels.transpose(1, 2), lengths) + + with torch.no_grad(): + padded = encoder._run_diarization(mels, lengths) + packed = encoder._run_diarization_packed(packed_input) + asr = encoder._run_asr_packed(packed_input) + + assert torch.equal(packed.lengths, asr.lengths) + assert padded.shape[1] == asr.max_seqlen + restored = unpack_encoder_output(packed, total_length=padded.shape[1]) + valid = torch.arange(padded.shape[1])[None, :] < packed.lengths[:, None] + torch.testing.assert_close(restored[valid], padded[valid], rtol=1e-5, atol=1e-6) + + +@pytest.mark.unit +def test_high_resolution_fusion_pools_offline_but_not_aligned_online_predictions(): + diarization_config = toy_diarization_model_cfg() + diarization_config.high_resolution = True + diarization_config.output_subsampling_factor = 1 + encoder = build_toy_pe_encoder(diarization_model_cfg=diarization_config).eval() + asr_states = torch.randn(1, _ASR_D_MODEL, 3) + fine_predictions = torch.sigmoid(torch.arange(24 * _N_SPK).reshape(1, 24, _N_SPK).float() / 17.0 - 2.0) + aligned_predictions = encoder.diarization_model.sortformer_modules.downsample_preds( + fine_predictions, _SUBSAMPLING_FACTOR + ) + + offline_fused = encoder._fuse_diar_and_asr(asr_states, fine_predictions) + online_fused = encoder._fuse_diar_and_asr(asr_states, aligned_predictions) + + assert aligned_predictions.shape[1] == asr_states.shape[-1] + torch.testing.assert_close(offline_fused, online_fused) + + +@pytest.mark.unit +def test_packed_diarization_supports_optional_post_encoder(): + diarization_config = toy_packed_diarization_model_cfg() + diarization_config.transformer_encoder = None + encoder = build_toy_pe_encoder(diarization_model_cfg=diarization_config).eval() + lengths = torch.tensor([80, 53]) + mels = torch.randn(2, _MEL_FEATURES, 80) + + with torch.no_grad(): + predictions = encoder._run_diarization_packed(pack_encoder_output(mels.transpose(1, 2), lengths)) + + assert predictions.lengths.tolist() == [10, 7] + assert torch.isfinite(predictions.data).all() + + +@pytest.mark.unit +def test_packed_fallback_matches_padded_forward_for_dense_and_packed_inputs(): + torch.manual_seed(0) + encoder = build_toy_pe_encoder().eval() + lengths = torch.tensor([80, 53]) + mels = torch.randn(2, _MEL_FEATURES, 80) + mels[1, :, 53:] = 0.0 + targets = torch.zeros(2, 10, _N_SPK) + targets[0, :, 0] = 1.0 + targets[1] = -1.0 + + with torch.no_grad(): + padded, output_lengths = encoder(mels, lengths, spk_targets=targets) + packed_from_dense = encoder.forward_sequence_packed(mels, lengths, spk_targets=targets) + packed_input = pack_encoder_output(mels.transpose(1, 2), lengths) + packed_from_packed = encoder.forward_sequence_packed(packed_input, spk_targets=targets) + + restored = unpack_encoder_output(packed_from_dense, total_length=padded.shape[-1]) + valid = torch.arange(padded.shape[-1])[None, :] < output_lengths[:, None] + torch.testing.assert_close(restored[valid], padded.transpose(1, 2)[valid], rtol=1e-4, atol=1e-5) + torch.testing.assert_close(packed_from_packed.data, packed_from_dense.data, rtol=1e-5, atol=1e-6) + assert torch.equal(packed_from_packed.lengths, packed_from_dense.lengths) + + +@pytest.mark.unit +def test_native_packed_path_normalizes_diar_and_asr_independently_without_unpack( + monkeypatch, +): + encoder = build_toy_packed_pe_encoder().eval() + lengths = torch.tensor([80, 53]) + mels = torch.randn(2, _MEL_FEATURES, 80) + packed = pack_encoder_output(mels.transpose(1, 2), lengths) + + calls = [] + original_forward = encoder._forward_packed_branch + + def tracked_forward(branch, features, chunk_size_seconds): + calls.append((branch, features.data.detach().clone(), chunk_size_seconds)) + return original_forward(branch, features, chunk_size_seconds) + + monkeypatch.setattr(encoder, "_forward_packed_branch", tracked_forward) + monkeypatch.setattr( + pee_module, + "unpack_encoder_output", + lambda *_args, **_kwargs: (_ for _ in ()).throw(AssertionError("native packed path unpacked")), + ) + with torch.no_grad(): + output = encoder.forward_sequence_packed(packed) + + assert [branch for branch, _, _ in calls] == [ + encoder.diarization_model.encoder, + encoder.asr_encoder, + ] + expected_diar = normalize_packed_batch(packed, "per_feature") + torch.testing.assert_close(calls[0][1], expected_diar.data) + torch.testing.assert_close(calls[1][1], expected_diar.data) + assert torch.equal(output.lengths, torch.tensor([10, 7])) + assert torch.isfinite(output.data).all() + + +@pytest.mark.unit +def test_native_packed_output_matches_padded_two_branch_path(): + torch.manual_seed(0) + encoder = build_toy_packed_pe_encoder().eval() + lengths = torch.tensor([80, 53]) + mels = torch.randn(2, _MEL_FEATURES, 80) + mels[1, :, 53:] = 0.0 + with torch.no_grad(): + padded, padded_lengths = encoder(mels, lengths) + packed = encoder.forward_sequence_packed(mels, lengths) + restored = unpack_encoder_output(packed, total_length=padded.shape[-1]) + valid = torch.arange(padded.shape[-1])[None, :] < padded_lengths[:, None] + torch.testing.assert_close(restored[valid], padded.transpose(1, 2)[valid], rtol=1e-4, atol=1e-5) + + +@pytest.mark.unit +def test_native_packed_branches_share_chunk_size_after_feature_stacking(monkeypatch): + encoder = build_toy_packed_pe_encoder( + frame_shift_seconds=0.01, + chunk_size_seconds=0.16, + ).eval() + packed = pack_encoder_output(torch.randn(2, 80, _MEL_FEATURES), torch.tensor([80, 65])) + diar_calls = [] + asr_calls = [] + diar_forward = encoder.diarization_model.encoder.forward_sequence_packed + asr_forward = encoder.asr_encoder.forward_sequence_packed + + def tracked_diar(audio_signal, length, *args, **kwargs): + diar_calls.append( + ( + audio_signal.lengths.detach().cpu().tolist(), + kwargs.get("bypass_pre_encode", False), + ) + ) + return diar_forward(audio_signal, length, *args, **kwargs) + + def tracked_asr(audio_signal, length, *args, **kwargs): + asr_calls.append( + ( + audio_signal.lengths.detach().cpu().tolist(), + kwargs.get("bypass_pre_encode", False), + ) + ) + return asr_forward(audio_signal, length, *args, **kwargs) + + monkeypatch.setattr(encoder.diarization_model.encoder, "forward_sequence_packed", tracked_diar) + monkeypatch.setattr(encoder.asr_encoder, "forward_sequence_packed", tracked_asr) + with torch.no_grad(): + output = encoder.forward_sequence_packed(packed) + + expected_chunks = [([2, 2, 2, 2, 2, 2, 2, 2, 2, 1], True)] + assert diar_calls == expected_chunks + assert asr_calls == expected_chunks + assert output.lengths.tolist() == [10, 9] + + +@pytest.mark.unit +def test_packed_fallback_rejects_online_scope(): + encoder = build_toy_pe_encoder().eval() + with encoder.online_inference(), pytest.raises(RuntimeError, match="offline API"): + encoder.forward_sequence_packed(torch.randn(1, _MEL_FEATURES, 32), torch.tensor([32])) + + +@pytest.mark.unit +def test_activation_checkpointing_wraps_trainable_asr_layers_and_packed_backward(): + encoder = build_toy_packed_pe_encoder().train() + encoder.set_activation_checkpointing(True) + encoder.set_activation_checkpointing(True) + + assert getattr(encoder.asr_encoder.pre_encode, "_checkpoint_wrapped_module", None) is not None + assert all(getattr(layer, "_checkpoint_wrapped_module", None) is not None for layer in encoder.asr_encoder.layers) + assert all( + getattr(layer, "_checkpoint_wrapped_module", None) is None + for layer in encoder.diarization_model.encoder.layers + ) + + mels = torch.randn(1, 64, _MEL_FEATURES, requires_grad=True) + packed = pack_encoder_output(mels, torch.tensor([64])) + output = encoder._run_asr_packed(packed) + output.data.square().mean().backward() + assert mels.grad is not None + assert torch.isfinite(mels.grad).all() diff --git a/tests/collections/common/test_indexed_build_cli.py b/tests/collections/common/test_indexed_build_cli.py index f8db56beabe0..4e38518099d0 100644 --- a/tests/collections/common/test_indexed_build_cli.py +++ b/tests/collections/common/test_indexed_build_cli.py @@ -335,6 +335,8 @@ def test_full_mode_builds_production_salm_dataset(): dataset = validate_dataloader._build_validation_dataset(config, _Tokenizer(), mode="full") assert isinstance(dataset, SALMDataset) + assert dataset.pack_audio is True + assert dataset.batch_tokens == 1024 assert dataset.strict_audio_loading is True diff --git a/tests/collections/common/test_salm_text_only_compat.py b/tests/collections/common/test_salm_text_only_compat.py new file mode 100644 index 000000000000..c32a9c0e7fb2 --- /dev/null +++ b/tests/collections/common/test_salm_text_only_compat.py @@ -0,0 +1,72 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch +from lhotse import CutSet + +from nemo.collections.common.data.lhotse.text_adapters import NeMoSFTExample +from nemo.collections.speechlm2.data.salm_dataset import SALMDataset + + +class _Tokenizer: + pad = 0 + unk_id = 1 + + +@pytest.mark.unit +@pytest.mark.parametrize("pack_audio", [False, True]) +def test_strict_salm_batching_accepts_prompt_formatted_text_only_examples(pack_audio): + example = NeMoSFTExample(data={"conversations": []}) + example.input_ids = torch.tensor([7, 8], dtype=torch.long) + example.mask = torch.tensor([False, True]) + + batch = SALMDataset(_Tokenizer(), pack_audio=pack_audio)[CutSet([example])] + + assert batch["audio_lens"].numel() == 0 + assert batch["input_ids"].tolist() == [[7, 8]] + assert list(batch["conversations"]) == [example] + if pack_audio: + assert batch["packed_audio_samples"].numel() == 0 + assert batch["audio_cu_seqlens"].tolist() == [0] + else: + assert batch["audios"].numel() == 0 + + +@pytest.mark.unit +@pytest.mark.parametrize("pack_sequences", [False, True]) +def test_text_only_examples_with_multispeaker_cfg_have_no_speaker_targets( + pack_sequences, +): + example = NeMoSFTExample(data={"conversations": []}) + example.input_ids = torch.tensor([7, 8], dtype=torch.long) + example.mask = torch.tensor([False, True]) + + dataset = SALMDataset( + _Tokenizer(), + pack_sequences=pack_sequences, + multispeaker_cfg={"num_speakers": 2}, + ) + batch = dataset[CutSet([example])] + + assert batch["audio_lens"].numel() == 0 + assert list(batch["conversations"]) == [example] + assert "spk_targets" not in batch + if pack_sequences: + assert batch["input_ids"].tolist() == [7, 8] + assert batch["text_cu_seqlens"].tolist() == [0, 2] + assert batch["audio_cu_seqlens"].tolist() == [0] + else: + assert batch["input_ids"].tolist() == [[7, 8]] + assert batch["audios"].numel() == 0 diff --git a/tests/collections/speechlm2/_chunking_helpers.py b/tests/collections/speechlm2/_chunking_helpers.py index 19f7abccd359..2866a96e18b5 100644 --- a/tests/collections/speechlm2/_chunking_helpers.py +++ b/tests/collections/speechlm2/_chunking_helpers.py @@ -16,6 +16,9 @@ from types import SimpleNamespace import torch +from torch.nn.utils.rnn import pad_sequence + +from nemo.collections.asr.parts.packed_sequence import pack_encoder_output def chunking_test_devices(): @@ -46,6 +49,8 @@ def __init__(self, sampling_rate, hop_length): self.calls = [] self.time_offsets = [] self.spk_targets_calls = [] + self.sequence_packed_calls = 0 + self.supports_sequence_packed_output = True def forward(self, input_signal=None, input_signal_length=None, time_offset=None, spk_targets=None): self.calls.append((input_signal.detach().clone(), input_signal_length.detach().clone())) @@ -54,6 +59,26 @@ def forward(self, input_signal=None, input_signal_length=None, time_offset=None, max_len = int(input_signal_length.max().item()) if input_signal_length.numel() > 0 else 0 return input_signal[:, :max_len].unsqueeze(-1), input_signal_length.clone() + def forward_sequence_packed( + self, + input_signal=None, + input_signal_length=None, + time_offset=None, + spk_targets=None, + input_signal_cu_seqlens=None, + ): + self.sequence_packed_calls += 1 + if input_signal_cu_seqlens is not None: + offsets = input_signal_cu_seqlens.tolist() + input_signal = pad_sequence( + [input_signal[offsets[row] : offsets[row + 1]] for row in range(input_signal_length.numel())], + batch_first=True, + ) + padded, lengths = self.forward( + input_signal, input_signal_length, time_offset=time_offset, spk_targets=spk_targets + ) + return pack_encoder_output(padded, lengths) + class _ChunkingTestFeaturizer: def __init__(self, sampling_rate, hop_length): diff --git a/tests/collections/speechlm2/test_datamodule.py b/tests/collections/speechlm2/test_datamodule.py index a1c6ef1144e8..3a356bfdae21 100644 --- a/tests/collections/speechlm2/test_datamodule.py +++ b/tests/collections/speechlm2/test_datamodule.py @@ -22,6 +22,7 @@ from omegaconf import DictConfig import nemo.collections.speechlm2.data.datamodule as datamodule_module +from nemo.collections.common.data.fallback import FallbackDataset from nemo.collections.common.data.lhotse.broadcasting import BroadcastingDataLoader from nemo.collections.common.tokenizers.sentencepiece_tokenizer import SentencePieceTokenizer, create_spt_model from nemo.collections.speechlm2.data import DataModule @@ -92,6 +93,19 @@ def __getitem__(self, item): return item +class PolicyAwareDataset(torch.utils.data.Dataset): + def __init__(self, calls, fault_tolerant_audio_loading=None): + self.calls = calls + self.fault_tolerant_audio_loading = fault_tolerant_audio_loading + + def __getitem__(self, item): + return item + + def with_fault_tolerant_audio_loading(self, enabled): + self.calls.append(enabled) + return PolicyAwareDataset(self.calls, fault_tolerant_audio_loading=enabled) + + def test_datamodule_train_dataloader(data_config, tokenizer): data = DataModule(data_config, tokenizer=tokenizer, dataset=Identity()) dl = data.train_dataloader() @@ -104,6 +118,41 @@ def test_datamodule_train_dataloader(data_config, tokenizer): assert all(c.tag == "train" for c in batch) +@pytest.mark.parametrize("skip_missing_manifest_entries", [False, True]) +@pytest.mark.parametrize("fault_tolerant_audio_loading", [False, True]) +def test_datamodule_audio_loading_policy_is_independent_of_missing_manifest_policy( + skip_missing_manifest_entries, fault_tolerant_audio_loading +): + calls = [] + dataset = PolicyAwareDataset(calls) + cfg = DictConfig( + { + "train_ds": { + "skip_missing_manifest_entries": skip_missing_manifest_entries, + "fault_tolerant_audio_loading": fault_tolerant_audio_loading, + } + } + ) + datamodule = DataModule(cfg, tokenizer=None, dataset=dataset) + training = datamodule._dataset_for_config(cfg.train_ds, training=True) + configured = training.dataset if isinstance(training, FallbackDataset) else training + assert isinstance(training, FallbackDataset) is fault_tolerant_audio_loading + assert configured.fault_tolerant_audio_loading is fault_tolerant_audio_loading + + validation = datamodule._dataset_for_config(cfg.train_ds, training=False) + assert isinstance(validation, PolicyAwareDataset) + assert validation.fault_tolerant_audio_loading is fault_tolerant_audio_loading + assert calls == [fault_tolerant_audio_loading, fault_tolerant_audio_loading] + + +def test_datamodule_fault_tolerant_audio_loading_defaults_true(): + dataset = PolicyAwareDataset([]) + cfg = DictConfig({"train_ds": {}}) + configured = DataModule(cfg, tokenizer=None, dataset=dataset)._dataset_for_config(cfg.train_ds, training=True) + assert isinstance(configured, FallbackDataset) + assert configured.dataset.fault_tolerant_audio_loading is True + + def test_datamodule_train_dataloader_caches_broadcast_wrapper_and_passes_dp_group(data_config, tokenizer, monkeypatch): data = DataModule(data_config, tokenizer=tokenizer, dataset=Identity()) mesh = SimpleNamespace(mesh_dim_names=()) diff --git a/tests/collections/speechlm2/test_encoder_chunking.py b/tests/collections/speechlm2/test_encoder_chunking.py index ea7c49b3e4c6..60b71a34a10b 100644 --- a/tests/collections/speechlm2/test_encoder_chunking.py +++ b/tests/collections/speechlm2/test_encoder_chunking.py @@ -18,14 +18,50 @@ from nemo.collections.speechlm2.data.salm_dataset import MultiSpeakerConfig from nemo.collections.speechlm2.parts import encoder_chunking as encoder_chunking_module from nemo.collections.speechlm2.parts.encoder_chunking import ( + _preserve_module_buffers, _recombine_chunked_audio_embeddings, _split_audio_into_chunks, _split_spk_targets_into_chunks, encode_audio_with_optional_chunking, + materialize_packed_spk_targets, ) + from tests.collections.speechlm2._chunking_helpers import ChunkingTestPerception +def test_materialize_packed_spk_targets_preserves_missing_rttm_sentinel(): + packed = torch.tensor( + [ + [1.0, 0.0], + [0.0, 1.0], + [1.0, 1.0], + [-1.0, -1.0], + [1.0, 0.0], + [0.0, 1.0], + ] + ) + lengths = torch.tensor([3, 1, 2]) + + targets, actual_lengths = materialize_packed_spk_targets( + packed, + lengths, + torch.tensor([0, 3, 4, 6]), + ) + + assert torch.equal(actual_lengths, lengths) + assert torch.equal( + targets, + torch.tensor( + [ + [[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]], + [[-1.0, -1.0], [-1.0, -1.0], [-1.0, -1.0]], + [[1.0, 0.0], [0.0, 1.0], [0.0, 0.0]], + ] + ), + ) + assert (targets == -1.0).all(dim=(1, 2)).tolist() == [False, True, False] + + @pytest.mark.parametrize( ( "input_signal_lengths", @@ -421,6 +457,31 @@ def raise_synced_microbatch_count(count, *, op, group): torch.testing.assert_close(perception.scale.grad, torch.tensor(10.0)) +def test_preserve_module_buffers_does_not_invalidate_saved_immutable_buffer(): + class BufferUsingModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("kernel", torch.randn(8, 2048)) + self.register_buffer("forward_count", torch.zeros((), dtype=torch.long)) + + def forward(self, inputs): + self.forward_count.add_(1) + return inputs @ self.kernel.T + + module = BufferUsingModule() + inputs = torch.randn(2, 2048, requires_grad=True) + real_output = module(inputs) + kernel_version = module.kernel._version + + with _preserve_module_buffers(module): + dummy_output = module(inputs) + + assert module.forward_count.item() == 1 + assert module.kernel._version == kernel_version + (real_output.sum() + dummy_output.sum() * 0.0).backward() + assert inputs.grad is not None + + @pytest.mark.parametrize( ("audio_values", "audio_len", "expected_chunk_lens", "expected_spk_targets"), [ @@ -467,3 +528,101 @@ def test_encode_audio_with_optional_chunking_forwards_chunked_spk_targets( assert torch.equal(chunked_lens, torch.tensor(expected_chunk_lens, dtype=torch.long)) assert torch.equal(perception.spk_targets_calls[0], torch.tensor(expected_spk_targets)) assert torch.equal(embs[0].squeeze(-1), audios[0]) + + +@pytest.mark.parametrize( + ("chunk_size_seconds", "chunk_batch_size", "expected_calls"), + [(None, None, 1), (2.0, None, 1), (2.0, 2, 3)], +) +def test_sequence_packed_chunking_matches_legacy(chunk_size_seconds, chunk_batch_size, expected_calls): + audios = torch.tensor( + [ + [1.0, 2.0, 3.0, 4.0, 5.0, 6.0], + [7.0, 8.0, 9.0, 10.0, 0.0, 0.0], + ] + ) + audio_lens = torch.tensor([6, 4], dtype=torch.long) + legacy_perception = ChunkingTestPerception(sampling_rate=1, hop_length=1) + packed_perception = ChunkingTestPerception(sampling_rate=1, hop_length=1) + + legacy = encode_audio_with_optional_chunking( + legacy_perception, + audios, + audio_lens, + chunk_size_seconds=chunk_size_seconds, + chunk_batch_size=chunk_batch_size, + sampling_rate=1, + ) + packed = encode_audio_with_optional_chunking( + packed_perception, + audios, + audio_lens, + chunk_size_seconds=chunk_size_seconds, + chunk_batch_size=chunk_batch_size, + sampling_rate=1, + sequence_packed=True, + ) + + assert packed_perception.sequence_packed_calls == expected_calls + assert len(packed) == len(legacy) == 2 + for actual, expected in zip(packed, legacy): + torch.testing.assert_close(actual, expected) + + +@pytest.mark.parametrize("sequence_packed", [False, True]) +@pytest.mark.parametrize( + ("chunk_size_seconds", "chunk_batch_size"), + [(None, None), (2.0, None), (2.0, 2)], +) +def test_packed_audio_samples_match_padded_frontend(sequence_packed, chunk_size_seconds, chunk_batch_size): + audios = torch.tensor( + [ + [1.0, 2.0, 3.0, 4.0, 5.0, 6.0], + [7.0, 8.0, 9.0, 10.0, 0.0, 0.0], + ] + ) + audio_lens = torch.tensor([6, 4], dtype=torch.long) + packed_audio_samples = torch.cat([audios[0, :6], audios[1, :4]]) + audio_cu_seqlens = torch.tensor([0, 6, 10], dtype=torch.long) + padded_perception = ChunkingTestPerception(sampling_rate=1, hop_length=1) + packed_perception = ChunkingTestPerception(sampling_rate=1, hop_length=1) + + expected = encode_audio_with_optional_chunking( + padded_perception, + audios, + audio_lens, + chunk_size_seconds=chunk_size_seconds, + chunk_batch_size=chunk_batch_size, + sampling_rate=1, + sequence_packed=sequence_packed, + ) + actual = encode_audio_with_optional_chunking( + packed_perception, + packed_audio_samples, + audio_lens, + input_signal_cu_seqlens=audio_cu_seqlens, + chunk_size_seconds=chunk_size_seconds, + chunk_batch_size=chunk_batch_size, + sampling_rate=1, + sequence_packed=sequence_packed, + ) + + assert len(actual) == len(expected) == 2 + for actual_row, expected_row in zip(actual, expected): + torch.testing.assert_close(actual_row, expected_row, rtol=0.0, atol=0.0) + + +def test_sequence_packed_chunking_rejects_unsupported_perception(): + class UnsupportedPerception: + def __call__(self, **kwargs): + raise AssertionError("legacy path must not be called") + + with pytest.raises(ValueError, match="does not support native packed output"): + encode_audio_with_optional_chunking( + UnsupportedPerception(), + torch.zeros(1, 4), + torch.tensor([4]), + chunk_size_seconds=None, + sampling_rate=1, + sequence_packed=True, + ) diff --git a/tests/collections/speechlm2/test_gc.py b/tests/collections/speechlm2/test_gc.py new file mode 100644 index 000000000000..a921173de907 --- /dev/null +++ b/tests/collections/speechlm2/test_gc.py @@ -0,0 +1,55 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest + +from nemo.collections.speechlm2.parts.gc import GarbageCollectionManager + + +def test_garbage_collection_manager_owns_step_state(monkeypatch): + calls = [] + + class FakeGarbageCollection: + def __init__(self, gc_every_steps): + calls.append(("init", gc_every_steps)) + + def run(self, step_count): + calls.append(("run", step_count)) + + import nemo_automodel.components.training.garbage_collection as gc_module + + monkeypatch.setattr(gc_module, "GarbageCollection", FakeGarbageCollection) + manager = GarbageCollectionManager(gc_every_steps=10) + + manager.on_fit_start() + manager.on_optimizer_step() + manager.on_optimizer_step() + + assert calls == [("init", 10), ("run", 1), ("run", 2)] + + +def test_garbage_collection_manager_is_noop_when_disabled(): + manager = GarbageCollectionManager(gc_every_steps=None) + + manager.on_fit_start() + manager.on_optimizer_step() + + assert manager._collector is None + assert manager._optimizer_step_count == 0 + + +@pytest.mark.parametrize("value", [True, False, 0, -1, 1.5, "10"]) +def test_garbage_collection_manager_rejects_invalid_interval(value): + with pytest.raises(ValueError, match="gc_every_steps"): + GarbageCollectionManager(gc_every_steps=value) diff --git a/tests/collections/speechlm2/test_independent_dual_encoder.py b/tests/collections/speechlm2/test_independent_dual_encoder.py new file mode 100644 index 000000000000..c65e88a1d435 --- /dev/null +++ b/tests/collections/speechlm2/test_independent_dual_encoder.py @@ -0,0 +1,90 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch + +from nemo.collections.asr.modules.transformer_encoder import TransformerEncoder +from nemo.collections.asr.parts.packed_sequence import pack_encoder_output +from nemo.collections.speechlm2.modules.perception import IndependentDualEncoder + + +def make_encoder(d_model: int, *, subsampling_factor: int = 2) -> TransformerEncoder: + return TransformerEncoder( + feat_in=4, + d_model=d_model, + n_heads=2, + n_layers=2, + subsampling="feature_stacking", + subsampling_factor=subsampling_factor, + ff_expansion=1.0, + self_attention_model="rope", + pos_emb_max_len=64, + drop_rate=0.0, + dropout_pre_encoder=0.0, + sync_max_audio_length=False, + ) + + +def test_independent_dual_encoder_chunks_only_frozen_auxiliary_branch(): + torch.manual_seed(7) + asr = make_encoder(32) + auxiliary = make_encoder(32) + dual = IndependentDualEncoder( + asr, + auxiliary, + frame_shift_seconds=0.01, + asr_chunk_size_seconds=None, + auxiliary_chunk_size_seconds=0.04, + freeze_auxiliary=True, + ).train() + + features = torch.randn(2, 4, 17) + lengths = torch.tensor([17, 10], dtype=torch.int64) + packed_features = pack_encoder_output(features.transpose(1, 2), lengths) + with torch.no_grad(): + asr_reference = asr.forward_sequence_packed(packed_features, packed_features.lengths) + + auxiliary_calls = [] + auxiliary_forward = auxiliary.forward_sequence_packed + + def record_auxiliary_call(audio_signal, length, bypass_pre_encode=False, **kwargs): + auxiliary_calls.append((length.detach().clone(), bypass_pre_encode)) + return auxiliary_forward(audio_signal, length, bypass_pre_encode=bypass_pre_encode, **kwargs) + + auxiliary.forward_sequence_packed = record_auxiliary_call + output = dual.forward_sequence_packed(packed_features, packed_features.lengths) + + assert output.lengths.tolist() == [9, 5] + assert output.data.shape == (14, 64) + torch.testing.assert_close(output.data[:, :32], asr_reference.data) + assert len(auxiliary_calls) == 1 + assert auxiliary_calls[0][0].tolist() == [2, 2, 2, 2, 1, 2, 2, 1] + assert auxiliary_calls[0][1] is True + assert not auxiliary.training + assert all(not parameter.requires_grad for parameter in auxiliary.parameters()) + + output.data.square().mean().backward() + assert any(parameter.grad is not None for parameter in asr.parameters()) + assert all(parameter.grad is None for parameter in auxiliary.parameters()) + + +def test_independent_dual_encoder_rejects_mismatched_frame_rates(): + asr = make_encoder(32, subsampling_factor=2) + auxiliary = make_encoder(32, subsampling_factor=4) + try: + IndependentDualEncoder(asr, auxiliary, frame_shift_seconds=0.01) + except ValueError as error: + assert "subsampling_factor" in str(error) + else: + raise AssertionError("Expected mismatched subsampling factors to be rejected.") diff --git a/tests/collections/speechlm2/test_packed_sequence_resume.py b/tests/collections/speechlm2/test_packed_sequence_resume.py new file mode 100644 index 000000000000..32412222a903 --- /dev/null +++ b/tests/collections/speechlm2/test_packed_sequence_resume.py @@ -0,0 +1,48 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch + +from tests.collections.speechlm2.test_salm_automodel import _make_chunking_test_model + + +def test_salm_manual_model_and_optimizer_state_restore_with_runtime_packed_opt_in(tmp_path): + previous = _make_chunking_test_model(encoder_chunk_size_seconds=1.0, sampling_rate=2, device="cpu") + optimizer = torch.optim.AdamW(previous.parameters(), lr=1e-3) + previous.llm.model.embed_tokens.weight.square().sum().backward() + optimizer.step() + checkpoint_path = tmp_path / "previous_salm.ckpt" + torch.save( + {"state_dict": previous.state_dict(), "optimizer_states": [optimizer.state_dict()]}, + checkpoint_path, + ) + + resumed = _make_chunking_test_model(encoder_chunk_size_seconds=1.0, sampling_rate=2, device="cpu") + resumed.cfg["packed_encoder_sequences"] = True + resumed_optimizer = torch.optim.AdamW(resumed.parameters(), lr=1e-3) + checkpoint = torch.load(checkpoint_path, weights_only=True) + resumed.load_state_dict(checkpoint["state_dict"], strict=True) + resumed_optimizer.load_state_dict(checkpoint["optimizer_states"][0]) + + assert set(resumed.state_dict()) == set(previous.state_dict()) + assert len(resumed_optimizer.state) == len(optimizer.state) + device = next(resumed.parameters()).device + batch = { + "audios": torch.tensor([[1.0, 2.0, 3.0, 4.0, 5.0]], device=device), + "audio_lens": torch.tensor([5], dtype=torch.long, device=device), + "input_ids": torch.tensor([[resumed.audio_locator_tag_id, 10]], dtype=torch.long, device=device), + "loss_mask": torch.tensor([[False, True]], dtype=torch.bool, device=device), + } + resumed.prepare_inputs(batch) + assert resumed.perception.sequence_packed_calls == 1 diff --git a/tests/collections/speechlm2/test_parallel.py b/tests/collections/speechlm2/test_parallel.py index 94b54475720d..12ac2e677e89 100644 --- a/tests/collections/speechlm2/test_parallel.py +++ b/tests/collections/speechlm2/test_parallel.py @@ -13,12 +13,19 @@ # See the License for the specific language governing permissions and # limitations under the License. +import gc +import weakref +from types import SimpleNamespace + import pytest +import torch.distributed.checkpoint as dcp +import torch.distributed.checkpoint.state_dict as dcp_state_dict +from lightning.pytorch.strategies import model_parallel as lightning_model_parallel from lightning.pytorch.strategies.model_parallel import ModelParallelStrategy from nemo_automodel.components.distributed.config import FSDP2Config, MoEParallelizerConfig from omegaconf import DictConfig -from nemo.collections.speechlm2.parts.parallel import AutomodelParallelStrategy +from nemo.collections.speechlm2.parts.parallel import AutomodelParallelStrategy, _validate_missing_optimizer_state from nemo.utils.trainer_utils import _resolve_automodel_configs, resolve_trainer_cfg # --------------------------------------------------------------------------- @@ -41,9 +48,47 @@ def test_default_init(self): assert strategy._distributed_config is None assert strategy._moe_config is None assert strategy._moe_mesh is None + assert strategy._checkpoint_keepalive is None assert strategy.activation_checkpointing_llm is False assert strategy.activation_checkpointing_perception is False + def test_sharded_checkpoint_metadata_is_retained_after_return(self, monkeypatch, tmp_path): + class CheckpointMetadata: + pass + + class Model: + def load_state_dict(self, state_dict, strict): + assert state_dict == {} + assert strict is True + + class Reader: + def __init__(self, path): + assert path == tmp_path + + def read_metadata(self): + return object() + + strategy = AutomodelParallelStrategy() + strategy._model = Model() + strategy._lightning_module = SimpleNamespace(strict_loading=True) + strategy._optimizers = [] + strategy.broadcast = lambda path: path + monkeypatch.setattr(lightning_model_parallel, "_is_sharded_checkpoint", lambda path: path == tmp_path) + monkeypatch.setattr(dcp, "FileSystemReader", Reader) + monkeypatch.setattr(dcp, "load", lambda state, checkpoint_id: None) + monkeypatch.setattr(dcp_state_dict, "get_model_state_dict", lambda model: {}) + monkeypatch.setattr("torch.load", lambda path: CheckpointMetadata()) + + checkpoint = strategy.load_checkpoint(tmp_path) + checkpoint_ref = weakref.ref(checkpoint) + del checkpoint + gc.collect() + + assert checkpoint_ref() is strategy._checkpoint_keepalive + strategy._checkpoint_keepalive = None + gc.collect() + assert checkpoint_ref() is None + def test_accepts_activation_checkpointing_flags(self): strategy = AutomodelParallelStrategy( activation_checkpointing_llm=True, @@ -108,6 +153,42 @@ def test_distributed_sampler_kwargs_raises_before_setup(self): with pytest.raises(RuntimeError): _ = strategy.distributed_sampler_kwargs + def test_allows_wholly_absent_lazy_optimizer_parameter_state(self): + active = "optimizer_0.state.llm.model.layer.weight" + unused = "optimizer_0.state.llm.mtp.layer.weight" + fields = {"step", "exp_avg", "exp_avg_sq"} + target = {f"{prefix}.{field}" for prefix in (active, unused) for field in fields} + checkpoint = {f"{active}.{field}" for field in fields} + assert _validate_missing_optimizer_state( + target_keys=target, + checkpoint_keys=checkpoint, + parameter_names={"llm.model.layer.weight", "llm.mtp.layer.weight"}, + optimizer_key="optimizer_0", + ) == ["llm.mtp.layer.weight"] + + def test_rejects_partial_optimizer_parameter_state(self): + prefix = "optimizer_0.state.llm.mtp.layer.weight" + with pytest.raises(RuntimeError, match="partial optimizer state"): + _validate_missing_optimizer_state( + target_keys={ + f"{prefix}.step", + f"{prefix}.exp_avg", + f"{prefix}.exp_avg_sq", + }, + checkpoint_keys={f"{prefix}.step"}, + parameter_names={"llm.mtp.layer.weight"}, + optimizer_key="optimizer_0", + ) + + def test_rejects_missing_optimizer_metadata(self): + with pytest.raises(RuntimeError, match="missing optimizer metadata"): + _validate_missing_optimizer_state( + target_keys={"optimizer_0.param_groups.0.lr"}, + checkpoint_keys=set(), + parameter_names=set(), + optimizer_key="optimizer_0", + ) + # --------------------------------------------------------------------------- # _resolve_automodel_configs @@ -119,7 +200,10 @@ class TestResolveAutomodelConfigs: def test_plain_dict_to_fsdp2_config(self): strategy = AutomodelParallelStrategy( - distributed_config={"defer_fsdp_grad_sync": False, "sequence_parallel": True}, + distributed_config={ + "defer_fsdp_grad_sync": False, + "sequence_parallel": True, + }, ) _resolve_automodel_configs(strategy) assert isinstance(strategy.distributed_config, FSDP2Config) diff --git a/tests/collections/speechlm2/test_perception_packed_sequence.py b/tests/collections/speechlm2/test_perception_packed_sequence.py new file mode 100644 index 000000000000..3cc2a84bd0a5 --- /dev/null +++ b/tests/collections/speechlm2/test_perception_packed_sequence.py @@ -0,0 +1,238 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch + +from nemo.collections.asr.modules.audio_preprocessing import AudioToMelSpectrogramPreprocessor, SpectrogramAugmentation +from nemo.collections.asr.modules.transformer_encoder import TransformerEncoder +from nemo.collections.asr.parts.packed_sequence import unpack_encoder_output +from nemo.collections.speechlm2.modules.perception import AudioPerceptionModule, IdentityConnector +from tests.collections.asr.test_parallel_expert_encoder_two_branch import build_toy_packed_pe_encoder + + +class _FeaturePassthrough(torch.nn.Module): + def forward(self, input_signal, length): + return input_signal, length + + +def _make_perception() -> AudioPerceptionModule: + encoder = TransformerEncoder( + feat_in=8, + d_model=32, + n_heads=2, + n_layers=2, + subsampling_factor=2, + drop_rate=0.0, + dropout_pre_encoder=0.0, + dropout_emb=0.0, + self_attention_model='rope', + sync_max_audio_length=False, + ).eval() + perception = AudioPerceptionModule.__new__(AudioPerceptionModule) + torch.nn.Module.__init__(perception) + perception.preprocessor = _FeaturePassthrough() + perception._modules['encoder'] = encoder + perception.modality_adapter = IdentityConnector() + perception.proj = torch.nn.Linear(32, 24) + perception.spec_augmentation = None + perception.rote = None + return perception.eval() + + +def test_perception_sequence_packed_matches_legacy_and_preserves_state_dict(): + torch.manual_seed(0) + perception = _make_perception() + features = torch.randn(3, 8, 12) + lengths = torch.tensor([12, 7, 4]) + state_keys = set(perception.state_dict()) + + with torch.no_grad(): + legacy, output_lengths = perception(input_signal=features, input_signal_length=lengths) + legacy_with_encoder = perception( + input_signal=features, + input_signal_length=lengths, + return_encoder_emb=True, + ) + packed = perception.forward_sequence_packed(input_signal=features, input_signal_length=lengths) + + assert len(legacy_with_encoder) == 3 + torch.testing.assert_close(legacy_with_encoder[0], legacy) + assert torch.equal(legacy_with_encoder[1], output_lengths) + restored = unpack_encoder_output(packed, total_length=legacy.shape[1]) + valid = torch.arange(legacy.shape[1])[None, :] < output_lengths[:, None] + torch.testing.assert_close(restored[valid], legacy[valid], rtol=1e-5, atol=1e-6) + assert perception.supports_sequence_packed_output + assert set(perception.state_dict()) == state_keys + + +def test_perception_sequence_packed_rejects_adapter_that_cannot_preserve_thd(): + perception = _make_perception() + perception.modality_adapter = torch.nn.Identity() + + assert not perception.supports_sequence_packed_output + with pytest.raises(ValueError, match="IdentityConnector"): + perception.forward_sequence_packed( + input_signal=torch.randn(1, 8, 8), + input_signal_length=torch.tensor([8]), + ) + + +@pytest.mark.parametrize( + "device", + ["cpu", pytest.param("cuda", marks=pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA unavailable"))], +) +@pytest.mark.parametrize("encoder_kind", ["transformer", "pee"]) +def test_perception_packed_waveform_matches_padded_waveform_for_supported_encoders(encoder_kind, device): + torch.manual_seed(17) + perception = _make_waveform_perception(encoder_kind).to(device) + lengths = torch.tensor([4096, 2600, 1200], dtype=torch.long, device=device) + audios = torch.randn(3, int(lengths.max()), device=device) + for row, length in zip(audios, lengths): + row[int(length) :] = 0.0 + packed_audio_samples = torch.cat([row[: int(length)] for row, length in zip(audios, lengths)]) + audio_cu_seqlens = torch.cat([lengths.new_zeros(1), lengths.cumsum(0)]) + checkpoint = perception.state_dict() + + with torch.no_grad(): + expected, expected_lengths = perception(input_signal=audios, input_signal_length=lengths) + actual = perception.forward_sequence_packed( + input_signal=packed_audio_samples, + input_signal_length=lengths, + input_signal_cu_seqlens=audio_cu_seqlens, + ) + + restored = unpack_encoder_output(actual, total_length=expected.shape[1]) + valid = torch.arange(expected.shape[1], device=device)[None, :] < expected_lengths[:, None] + assert torch.equal(actual.lengths, expected_lengths) + atol = 1e-5 if encoder_kind == "pee" else 3e-6 + torch.testing.assert_close(restored[valid], expected[valid], rtol=2e-5, atol=atol) + assert set(perception.state_dict()) == set(checkpoint) + + reloaded = _make_waveform_perception(encoder_kind).to(device) + reloaded.load_state_dict(checkpoint, strict=True) + with torch.no_grad(): + reloaded_output = reloaded.forward_sequence_packed( + input_signal=packed_audio_samples, + input_signal_length=lengths, + input_signal_cu_seqlens=audio_cu_seqlens, + ) + torch.testing.assert_close(reloaded_output.data, actual.data, rtol=0.0, atol=0.0) + + +@pytest.mark.parametrize("encoder_kind", ["transformer", "pee"]) +def test_perception_legacy_forward_accepts_packed_waveform(encoder_kind): + torch.manual_seed(23) + perception = _make_waveform_perception(encoder_kind) + lengths = torch.tensor([4096, 2600, 1200]) + audios = torch.randn(3, int(lengths.max())) + audios.masked_fill_(torch.arange(audios.shape[1])[None, :] >= lengths[:, None], 0.0) + packed_audio = torch.cat([row[: int(length)] for row, length in zip(audios, lengths)]) + cu_seqlens = torch.cat([lengths.new_zeros(1), lengths.cumsum(0)]) + + with torch.no_grad(): + expected, expected_lengths = perception(input_signal=audios, input_signal_length=lengths) + actual, actual_lengths = perception( + input_signal=packed_audio, + input_signal_length=lengths, + input_signal_cu_seqlens=cu_seqlens, + ) + + assert torch.equal(actual_lengths, expected_lengths) + valid = torch.arange(expected.shape[1])[None, :] < expected_lengths[:, None] + atol = 1e-5 if encoder_kind == "pee" else 3e-6 + torch.testing.assert_close(actual[valid], expected[valid], rtol=2e-5, atol=atol) + + +@pytest.mark.parametrize("encoder_kind", ["transformer", "pee"]) +def test_perception_packed_waveform_all_empty_batch(encoder_kind, device): + torch_device = "cuda" if device == "GPU" and torch.cuda.is_available() else "cpu" + perception = _make_waveform_perception(encoder_kind).to(torch_device) + lengths = torch.tensor([0, 0], device=torch_device) + + with torch.no_grad(): + output = perception.forward_sequence_packed( + input_signal=torch.empty(0, device=torch_device), + input_signal_length=lengths, + input_signal_cu_seqlens=torch.tensor([0, 0, 0], device=torch_device), + ) + + assert output.data.shape == (0, 24) + assert output.lengths.tolist() == [0, 0] + assert output.cu_seqlens.tolist() == [0, 0, 0] + + +def test_perception_packed_waveform_training_dispatches_packed_spec_augment(monkeypatch): + perception = _make_waveform_perception("transformer").train() + perception.spec_augmentation = SpectrogramAugmentation(freq_masks=1, freq_width=2) + calls = 0 + original = perception.spec_augmentation.forward_packed + + def count_packed(input_spec): + nonlocal calls + calls += 1 + return original(input_spec) + + def reject_dense(*args, **kwargs): + raise AssertionError("packed waveform training must not densify for SpecAugment") + + monkeypatch.setattr(perception.spec_augmentation, "forward_packed", count_packed) + monkeypatch.setattr(perception.spec_augmentation, "forward", reject_dense) + lengths = torch.tensor([4096, 2600]) + packed_audio = torch.randn(int(lengths.sum())) + cu_seqlens = torch.cat([lengths.new_zeros(1), lengths.cumsum(0)]) + + output = perception.forward_sequence_packed( + input_signal=packed_audio, + input_signal_length=lengths, + input_signal_cu_seqlens=cu_seqlens, + ) + + assert calls == 1 + assert output.total_tokens == int(output.lengths.sum()) + + +def _make_waveform_perception(encoder_kind: str) -> AudioPerceptionModule: + if encoder_kind == "transformer": + encoder = TransformerEncoder( + feat_in=8, + d_model=32, + n_heads=2, + n_layers=2, + subsampling_factor=2, + drop_rate=0.0, + dropout_pre_encoder=0.0, + dropout_emb=0.0, + self_attention_model="rope", + sync_max_audio_length=False, + ) + features = 8 + else: + encoder = build_toy_packed_pe_encoder() + features = 128 + + perception = AudioPerceptionModule.__new__(AudioPerceptionModule) + torch.nn.Module.__init__(perception) + perception.preprocessor = AudioToMelSpectrogramPreprocessor( + features=features, + normalize="per_feature", + dither=0, + pad_to=0, + ) + perception._modules["encoder"] = encoder + perception.modality_adapter = IdentityConnector() + perception.proj = torch.nn.Linear(encoder.d_model, 24) + perception.spec_augmentation = None + perception.rote = None + return perception.eval() diff --git a/tests/collections/speechlm2/test_perception_packed_sequence_capabilities.py b/tests/collections/speechlm2/test_perception_packed_sequence_capabilities.py new file mode 100644 index 000000000000..eb5f16d512ce --- /dev/null +++ b/tests/collections/speechlm2/test_perception_packed_sequence_capabilities.py @@ -0,0 +1,89 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import copy + +import pytest +import torch + +from tests.collections.speechlm2.test_perception_packed_sequence import _make_perception + + +class _SpecAugment(torch.nn.Module): + def __init__(self): + super().__init__() + self.calls = 0 + + def forward(self, input_spec, length): + self.calls += 1 + return input_spec + 0.125 + + +class _NoPackedEncoder(torch.nn.Module): + def forward(self, audio_signal, length): + return audio_signal, length + + +@pytest.mark.parametrize("unsupported", ["rote", "multilayer", "marker", "method"]) +def test_perception_packed_capability_rejects_each_unsupported_stack(unsupported): + perception = _make_perception() + if unsupported == "rote": + perception.rote = torch.nn.Identity() + elif unsupported == "multilayer": + wrapper = torch.nn.Module() + wrapper.encoder = perception.encoder + perception.encoder_multilayer = wrapper + elif unsupported == "marker": + perception.encoder.supports_sequence_packed_output = False + else: + perception._modules["encoder"] = _NoPackedEncoder() + + assert not perception.supports_sequence_packed_output + with pytest.raises(ValueError, match="Packed encoder sequences"): + perception.forward_sequence_packed( + input_signal=torch.randn(1, 8, 8), + input_signal_length=torch.tensor([8]), + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Perception packed-gradient parity requires CUDA") +def test_perception_packed_spec_augmentation_and_projection_gradients_match_legacy(): + torch.manual_seed(0) + legacy = _make_perception().cuda().train() + legacy.spec_augmentation = _SpecAugment() + packed = copy.deepcopy(legacy) + legacy_features = torch.randn(3, 8, 12, device="cuda", requires_grad=True) + packed_features = legacy_features.detach().clone().requires_grad_() + lengths = torch.tensor([12, 7, 4], device="cuda") + + legacy_output, output_lengths = legacy(input_signal=legacy_features, input_signal_length=lengths) + packed_output = packed.forward_sequence_packed(input_signal=packed_features, input_signal_length=lengths) + valid = torch.arange(legacy_output.shape[1], device="cuda")[None, :] < output_lengths[:, None] + legacy_output[valid].square().mean().backward() + packed_output.data.square().mean().backward() + + assert legacy.spec_augmentation.calls == packed.spec_augmentation.calls == 1 + torch.testing.assert_close(packed_features.grad, legacy_features.grad, rtol=1e-5, atol=1e-6) + torch.testing.assert_close(packed.proj.weight.grad, legacy.proj.weight.grad, rtol=1e-5, atol=1e-6) + + +def test_perception_packed_api_rejects_legacy_encoder_return_request(): + perception = _make_perception() + + with pytest.raises(TypeError, match="return_encoder_emb"): + perception.forward_sequence_packed( + input_signal=torch.randn(1, 8, 8), + input_signal_length=torch.tensor([8]), + return_encoder_emb=True, + ) diff --git a/tests/collections/speechlm2/test_perception_packed_sequence_fsdp2_distributed.py b/tests/collections/speechlm2/test_perception_packed_sequence_fsdp2_distributed.py new file mode 100644 index 000000000000..b1f35fd383c8 --- /dev/null +++ b/tests/collections/speechlm2/test_perception_packed_sequence_fsdp2_distributed.py @@ -0,0 +1,231 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from torch.distributed.device_mesh import init_device_mesh + +from nemo.collections.asr.modules.transformer_encoder import TransformerEncoder +from nemo.collections.asr.parts.packed_sequence import pack_encoder_output +from nemo.collections.speechlm2.models.salm_automodel import _fully_shard_perception +from nemo.collections.speechlm2.modules.perception import AudioPerceptionModule, IdentityConnector +from nemo.collections.speechlm2.parts.cp_helpers import encode_audio_with_cp_distribution +from tests.collections.asr.test_parallel_expert_encoder_two_branch import build_toy_packed_pe_encoder + + +class _FeaturePassthrough(torch.nn.Module): + def forward(self, input_signal, length): + return input_signal, length + + +class _RepeatFeaturePreprocessor(torch.nn.Module): + def forward(self, input_signal, length): + return input_signal.unsqueeze(1).expand(-1, 128, -1).contiguous(), length + + +class _WorldCpMesh: + def size(self): + return dist.get_world_size() + + def get_group(self): + return dist.group.WORLD + + +class _ScalePerception(torch.nn.Module): + supports_sequence_packed_output = True + + def __init__(self, device): + super().__init__() + self.scale = torch.nn.Parameter(torch.tensor([2.0], device=device)) + + def forward_sequence_packed(self, *, input_signal, input_signal_length, **kwargs): + return pack_encoder_output(input_signal.unsqueeze(-1) * self.scale, input_signal_length) + + +def _make_perception(device) -> AudioPerceptionModule: + encoder = TransformerEncoder( + feat_in=8, + d_model=32, + n_heads=2, + n_layers=2, + subsampling_factor=1, + drop_rate=0.0, + dropout_pre_encoder=0.0, + dropout_emb=0.0, + self_attention_model="rope", + qk_norm=True, + sync_max_audio_length=False, + ).to(device) + perception = AudioPerceptionModule.__new__(AudioPerceptionModule) + torch.nn.Module.__init__(perception) + perception.preprocessor = _FeaturePassthrough() + perception._modules["encoder"] = encoder + perception.modality_adapter = IdentityConnector() + perception.proj = torch.nn.Linear(32, 24, device=device) + perception.spec_augmentation = None + perception.rote = None + return perception.train() + + +def _make_pee_perception(device) -> AudioPerceptionModule: + encoder = build_toy_packed_pe_encoder().to(device=device, dtype=torch.bfloat16) + perception = AudioPerceptionModule.__new__(AudioPerceptionModule) + torch.nn.Module.__init__(perception) + perception.preprocessor = _RepeatFeaturePreprocessor() + perception._modules['encoder'] = encoder + perception.modality_adapter = IdentityConnector() + perception.proj = torch.nn.Linear(encoder.d_model, 24, device=device, dtype=torch.bfloat16) + perception.spec_augmentation = None + perception.rote = None + return perception.train() + + +def _run_fsdp2_packed_perception_test(rank: int, world_size: int, init_file: str): + torch.cuda.set_device(rank) + dist.init_process_group("nccl", init_method=f"file://{init_file}", rank=rank, world_size=world_size) + try: + device = torch.device("cuda", rank) + mesh = init_device_mesh("cuda", (world_size,), mesh_dim_names=("dp",)) + + # Exercise the production custom perception entry point. Rank 1 owns no + # valid tokens, while rank 0 has audio; both ranks must reach identical + # FSDP2 collectives and materialize gradients for every sharded parameter. + perception = _fully_shard_perception(_make_perception(device), mesh) + optimizer = torch.optim.AdamW(perception.parameters(), lr=1e-2) + optimizer_parameters = tuple(parameter for group in optimizer.param_groups for parameter in group["params"]) + projection_parameter = perception.proj.weight + projection_before = projection_parameter.detach().full_tensor().clone() + if rank == 0: + features = torch.randn(1, 8, 12, device=device, requires_grad=True) + lengths = torch.tensor([12], device=device) + else: + features = torch.empty(1, 8, 0, device=device, requires_grad=True) + lengths = torch.tensor([0], device=device) + packed = perception.forward_sequence_packed(input_signal=features, input_signal_length=lengths) + packed.data.float().sum().backward() + + assert features.grad is not None + assert all(parameter.grad is not None for parameter in optimizer_parameters) + optimizer.step() + projection_after = projection_parameter.detach().full_tensor() + assert not torch.equal(projection_before, projection_after) + assert optimizer.state[projection_parameter]["step"] == 1 + + # A second step gives every rank an all-empty batch, covering the zero-token + # collective case independently of the uneven-rank step above. + optimizer.zero_grad(set_to_none=True) + empty_features = torch.empty(1, 8, 0, device=device, requires_grad=True) + empty_lengths = torch.tensor([0], device=device) + empty = perception.forward_sequence_packed( + input_signal=empty_features, + input_signal_length=empty_lengths, + ) + empty.data.sum().backward() + assert empty_features.grad is not None + assert all(parameter.grad is not None for parameter in optimizer_parameters) + + # Exercise the CP distribution/gather path with an FSDP2-sharded custom + # packed method. B=1 forces one CP rank to encode a dummy row. + scaled = _fully_shard_perception(_ScalePerception(device), mesh) + audios = torch.tensor([[1.0, 2.0, 3.0]], device=device) + audio_lens = torch.tensor([3], device=device) + embeddings = encode_audio_with_cp_distribution( + scaled, + audios, + audio_lens, + chunk_size_seconds=None, + sampling_rate=1, + cp_mesh=_WorldCpMesh(), + fsdp_sync_group=dist.group.WORLD, + sequence_packed=True, + packed_cp_gather=True, + ) + assert len(embeddings) == 1 + torch.testing.assert_close(embeddings[0][:, 0], audios[0] * 2.0) + embeddings[0].sum().backward() + assert all(parameter.grad is not None for parameter in scaled.parameters()) + finally: + dist.destroy_process_group() + + +def _run_fsdp2_canonical_pee_test(rank: int, world_size: int, init_file: str): + torch.cuda.set_device(rank) + dist.init_process_group("nccl", init_method=f"file://{init_file}", rank=rank, world_size=world_size) + try: + torch.manual_seed(0) + device = torch.device("cuda", rank) + mesh = init_device_mesh("cuda", (world_size,), mesh_dim_names=("dp",)) + perception = _fully_shard_perception(_make_pee_perception(device), mesh) + + frames = 32 if rank == 0 else 16 + raw_audio = torch.randn(1, frames, device=device, dtype=torch.bfloat16, requires_grad=True) + lengths = torch.tensor([frames if rank == 0 else 0], device=device) + packed = perception.forward_sequence_packed(input_signal=raw_audio, input_signal_length=lengths) + packed.data.float().sum().backward() + + assert raw_audio.grad is not None and torch.isfinite(raw_audio.grad).all() + assert all( + parameter.grad is not None and torch.isfinite(parameter.grad).all() + for parameter in perception.parameters() + if parameter.requires_grad + ) + perception.zero_grad(set_to_none=True) + empty_audio = torch.randn(1, 16, device=device, dtype=torch.bfloat16, requires_grad=True) + empty = perception.forward_sequence_packed( + input_signal=empty_audio, + input_signal_length=torch.tensor([0], device=device), + ) + empty.data.sum().backward() + # With no valid samples, the packed feature stacker has no reason to retain + # an autograd edge to the waveform. The distributed invariant is that every + # trainable parameter still receives a zero gradient and reaches the same + # FSDP collectives on every rank. + assert all(parameter.grad is not None for parameter in perception.parameters() if parameter.requires_grad) + + perception.zero_grad(set_to_none=True) + cp_audio = torch.randn(1, 24, device=device, dtype=torch.bfloat16) + embeddings = encode_audio_with_cp_distribution( + perception, + cp_audio, + torch.tensor([24], device=device), + chunk_size_seconds=None, + sampling_rate=1, + cp_mesh=_WorldCpMesh(), + fsdp_sync_group=dist.group.WORLD, + sequence_packed=True, + packed_cp_gather=True, + ) + assert len(embeddings) == 1 and embeddings[0].shape[-1] == 24 + assert torch.isfinite(embeddings[0]).all() + embeddings[0].float().sum().backward() + assert all(parameter.grad is not None for parameter in perception.parameters() if parameter.requires_grad) + finally: + dist.destroy_process_group() + + +@pytest.mark.skipif(not torch.cuda.is_available() or torch.cuda.device_count() < 2, reason="Test requires 2 GPUs") +def test_canonical_pee_packed_perception_fsdp2_empty_rank_all_empty_and_cp(tmp_path): + mp.spawn( + _run_fsdp2_canonical_pee_test, + args=(2, str(tmp_path / "canonical_pee_fsdp2_init")), + nprocs=2, + join=True, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available() or torch.cuda.device_count() < 2, reason="Test requires 2 GPUs") +def test_packed_perception_fsdp2_custom_forward_empty_rank_and_cp_gather(tmp_path): + mp.spawn(_run_fsdp2_packed_perception_test, args=(2, str(tmp_path / "fsdp2_init")), nprocs=2, join=True) diff --git a/tests/collections/speechlm2/test_pretrained.py b/tests/collections/speechlm2/test_pretrained.py index 0c382b96d871..188c5045d236 100644 --- a/tests/collections/speechlm2/test_pretrained.py +++ b/tests/collections/speechlm2/test_pretrained.py @@ -58,12 +58,113 @@ def test_setup_speech_encoder_hydrates_missing_config_without_weights(): assert model.cfg.perception.modality_adapter.output_dim == 8 +@pytest.mark.parametrize( + ("chunk_size_seconds", "packed_encoder_sequences"), + [(None, False), (30, False), (30, True)], +) +def test_setup_parallel_expert_encoder_maps_shared_chunk_size(chunk_size_seconds, packed_encoder_sequences): + pe_encoder_overrides = { + "speaker_feature_config_version": 1, + "speaker_feature_mode": "continuous", + "speaker_activity_threshold": None, + "diar_normalize_type": "per_feature", + } + pe_encoder = SimpleNamespace( + d_model=4, + n_spk=8, + _feat_in=80, + freeze_asr=False, + freeze_diar=True, + spk_kernel_scale=1.0, + chunk_size_seconds=45.0, + _bundle_config=DictConfig({"chunk_size_seconds": 45.0}), + online_inference_enabled=False, + ) + model = SimpleNamespace( + cfg=DictConfig( + { + "pe_encoder_path": "/tmp/placeholderParallelExpertEncoder.nemo", + "pe_encoder_overrides": pe_encoder_overrides, + "encoder_chunk_size_seconds": chunk_size_seconds, + "packed_encoder_sequences": packed_encoder_sequences, + "perception": { + "preprocessor": {"features": 80, "normalize": "per_feature"}, + "modality_adapter": {"d_model": 4}, + }, + } + ), + perception=SimpleNamespace( + encoder=SimpleNamespace(d_model=4), + modality_adapter=object(), + proj=torch.nn.Linear(4, 8), + preprocessor=SimpleNamespace(featurizer=SimpleNamespace(normalize="per_feature")), + ), + ) + + with patch.object(pretrained.ParallelExpertEncoderPT, "load_from_nemo", return_value=pe_encoder) as load: + pretrained.setup_parallel_expert_encoder(model) + + load.assert_called_once_with( + "/tmp/placeholderParallelExpertEncoder.nemo", + map_location="cpu", + strict=True, + config_overrides=pe_encoder_overrides, + ) + assert pe_encoder.chunk_size_seconds == chunk_size_seconds + assert pe_encoder._bundle_config.chunk_size_seconds == chunk_size_seconds + assert model.perception.preprocessor.featurizer.normalize is None + assert model.cfg.perception.preprocessor.normalize is None + + +@pytest.mark.parametrize( + ("cfg_update", "match"), + [ + ( + { + "encoder_chunk_size_seconds": 30.0, + "packed_encoder_sequences": True, + "encoder_chunk_batch_size": 2, + }, + "encoder_chunk_batch_size is not supported", + ), + ( + {"encoder_chunk_size_seconds": -1.0, "packed_encoder_sequences": True}, + "encoder_chunk_size_seconds must be positive or null", + ), + ( + {"pe_asr_chunk_size_seconds": 30.0}, + "use model.encoder_chunk_size_seconds", + ), + ], +) +def test_setup_parallel_expert_encoder_validates_shared_chunking_config(cfg_update, match): + pe_encoder = SimpleNamespace(chunk_size_seconds=None) + cfg = { + "pe_encoder_path": "/tmp/placeholderParallelExpertEncoder.nemo", + "perception": {}, + } + cfg.update(cfg_update) + model = SimpleNamespace( + cfg=DictConfig(cfg), + perception=SimpleNamespace(encoder=object()), + ) + + with ( + patch.object(pretrained.ParallelExpertEncoderPT, "load_from_nemo", return_value=pe_encoder), + pytest.raises(ValueError, match=match), + ): + pretrained.setup_parallel_expert_encoder(model) + + def _mock_automodel_loader(config): automodel = SimpleNamespace(from_config=MagicMock(return_value=object()), from_pretrained=MagicMock()) return ( automodel, patch.object(pretrained.AutoConfig, "from_pretrained", return_value=config), - patch.dict("sys.modules", {"nemo_automodel": SimpleNamespace(NeMoAutoModelForCausalLM=automodel)}), + patch.dict( + "sys.modules", + {"nemo_automodel": SimpleNamespace(NeMoAutoModelForCausalLM=automodel)}, + ), patch("nemo.collections.speechlm2.parts.automodel_compat.remove_automodel_backend_for_hf_fallback"), ) @@ -76,7 +177,11 @@ def test_load_pretrained_automodel_llm_builds_missing_mtp_before_loading_weights config_patch, module_patch, compat_patch, - patch.object(pretrained, "_resolve_automodel_checkpoint_path", return_value="base-checkpoint"), + patch.object( + pretrained, + "_resolve_automodel_checkpoint_path", + return_value="base-checkpoint", + ), patch.object(pretrained, "_load_automodel_base_checkpoint_without_mtp", create=True) as base_load, ): result = pretrained.load_pretrained_automodel_llm( @@ -116,7 +221,11 @@ def test_load_pretrained_automodel_llm_preserves_native_mtp_config_by_default(): config_patch, module_patch, compat_patch, - patch.object(pretrained, "_resolve_automodel_checkpoint_path", return_value="native-mtp-checkpoint"), + patch.object( + pretrained, + "_resolve_automodel_checkpoint_path", + return_value="native-mtp-checkpoint", + ), patch.object(pretrained, "_load_automodel_base_checkpoint_without_mtp", create=True) as base_load, ): pretrained.load_pretrained_automodel_llm( @@ -150,7 +259,11 @@ def test_load_pretrained_automodel_llm_can_replace_native_mtp_config(): config_patch, module_patch, compat_patch, - patch.object(pretrained, "_resolve_automodel_checkpoint_path", return_value="native-mtp-checkpoint"), + patch.object( + pretrained, + "_resolve_automodel_checkpoint_path", + return_value="native-mtp-checkpoint", + ), patch.object(pretrained, "_load_automodel_base_checkpoint_without_mtp", create=True) as base_load, ): result = pretrained.load_pretrained_automodel_llm( @@ -188,7 +301,9 @@ def test_load_pretrained_automodel_llm_can_replace_native_mtp_config(): pytest.param(_REPEATED_MTP_OVERRIDES, id="fallback-config-present"), ], ) -def test_load_pretrained_automodel_llm_accepts_one_depth_native_head_as_repeated(mtp_config_overrides): +def test_load_pretrained_automodel_llm_accepts_one_depth_native_head_as_repeated( + mtp_config_overrides, +): config = SimpleNamespace( num_nextn_predict_layers=1, mtp_hybrid_override_pattern="*E", @@ -237,7 +352,9 @@ def test_load_pretrained_automodel_llm_accepts_one_depth_native_head_as_repeated pytest.param(_REPEATED_MTP_OVERRIDES, id="fallback-config-present"), ], ) -def test_load_pretrained_automodel_llm_rejects_multi_depth_native_head_as_repeated(mtp_config_overrides): +def test_load_pretrained_automodel_llm_rejects_multi_depth_native_head_as_repeated( + mtp_config_overrides, +): config = SimpleNamespace( num_nextn_predict_layers=4, mtp_hybrid_override_pattern="*E", @@ -276,7 +393,11 @@ def test_load_pretrained_automodel_llm_builds_repeated_head_for_checkpoint_witho config_patch, module_patch, compat_patch, - patch.object(pretrained, "_resolve_automodel_checkpoint_path", return_value="base-checkpoint"), + patch.object( + pretrained, + "_resolve_automodel_checkpoint_path", + return_value="base-checkpoint", + ), patch.object(pretrained, "_load_automodel_base_checkpoint_without_mtp", create=True) as base_load, ): result = pretrained.load_pretrained_automodel_llm( @@ -362,7 +483,11 @@ def test_load_pretrained_automodel_llm_rejects_repeated_mode_without_head_defini config_patch, module_patch, compat_patch, - patch.object(pretrained, "_resolve_automodel_checkpoint_path", return_value="base-checkpoint"), + patch.object( + pretrained, + "_resolve_automodel_checkpoint_path", + return_value="base-checkpoint", + ), pytest.raises(ValueError, match="requires either a checkpoint with a native MTP head"), ): pretrained.load_pretrained_automodel_llm( @@ -402,7 +527,10 @@ def test_load_pretrained_automodel_llm_forwards_hf_resolution_kwargs(): result = pretrained.load_pretrained_automodel_llm( "private-checkpoint", trust_remote_code=True, - mtp_config_overrides={"num_nextn_predict_layers": 1, "mtp_hybrid_override_pattern": "*"}, + mtp_config_overrides={ + "num_nextn_predict_layers": 1, + "mtp_hybrid_override_pattern": "*", + }, token="secret-token", revision="exact-revision", cache_dir="/cache", @@ -471,7 +599,11 @@ def test_resolve_automodel_checkpoint_path_uses_exact_snapshot(tmp_path, include def test_automodel_mtp_depth_supports_non_nemotron_config_fields(): assert ( pretrained._automodel_config_mtp_depth( - SimpleNamespace(num_nextn_predict_layers=2, mtp_hybrid_override_pattern=None, mtp_layers_block_type=None) + SimpleNamespace( + num_nextn_predict_layers=2, + mtp_hybrid_override_pattern=None, + mtp_layers_block_type=None, + ) ) == 2 ) diff --git a/tests/collections/speechlm2/test_salm_automodel.py b/tests/collections/speechlm2/test_salm_automodel.py index 927516e9d243..ad74f23844d8 100644 --- a/tests/collections/speechlm2/test_salm_automodel.py +++ b/tests/collections/speechlm2/test_salm_automodel.py @@ -14,11 +14,13 @@ # limitations under the License. import inspect import os +from contextlib import contextmanager import pytest import torch from lhotse import CutSet, SupervisionSegment from lhotse.testing.dummies import dummy_cut, dummy_recording +from lightning import LightningModule from transformers import GenerationConfig from nemo.collections.common.data.lhotse import NeMoMultimodalConversation @@ -206,6 +208,62 @@ def test_salm_automodel_training_step_uses_dataloader_iter_signature(): assert list(inspect.signature(SALMAutomodel.training_step).parameters) == ["self", "dataloader_iter"] +def test_salm_automodel_forward_enters_configured_te_fp8_context(): + events = [] + + class FakeFP8: + @contextmanager + def maybe_te_autocast(self): + events.append("enter") + yield + events.append("exit") + + class FakeLLM(torch.nn.Module): + def __init__(self): + super().__init__() + self.backend = type("Backend", (), {"te_fp8": FakeFP8()})() + + def forward(self, *args, inputs_embeds, **kwargs): + events.append("forward") + return {"logits": inputs_embeds} + + model = SALMAutomodel.__new__(SALMAutomodel) + LightningModule.__init__(model) + model.llm = FakeLLM() + model._fused_linear_cross_entropy = None + + outputs = model.forward(torch.randn(1, 2, 4)) + + assert outputs["logits"].shape == (1, 2, 4) + assert events == ["enter", "forward", "exit"] + + +def test_salm_automodel_backward_does_not_enter_te_fp8_context(monkeypatch): + events = [] + + class FakeFP8: + @contextmanager + def maybe_te_autocast(self): + events.append("enter") + yield + events.append("exit") + + class FakeLLM(torch.nn.Module): + def __init__(self): + super().__init__() + self.backend = type("Backend", (), {"te_fp8": FakeFP8()})() + + model = SALMAutomodel.__new__(SALMAutomodel) + LightningModule.__init__(model) + model.llm = FakeLLM() + monkeypatch.setattr(model, "_setup_moe_fsdp_sync", lambda: events.append("setup")) + monkeypatch.setattr(LightningModule, "backward", lambda *_args, **_kwargs: events.append("backward")) + + model.backward(torch.tensor(1.0)) + + assert events == ["setup", "backward"] + + def test_salm_automodel_pad_token_override_preserves_eot_labels(monkeypatch): seen = {} @@ -244,6 +302,87 @@ def add_special_tokens(self, _tokens): assert packed["target_ids"].tolist() == [-100, -100, 42, 11, -100] +def test_salm_automodel_fused_linear_forward_keeps_hidden_states_without_logits(): + class FakeLLM(torch.nn.Module): + def __init__(self): + super().__init__() + self.kwargs = None + + def forward(self, **kwargs): + self.kwargs = kwargs + hidden = kwargs["inputs_embeds"] + 3 + logits = hidden[..., :0] if kwargs.get("compute_logits") is False else hidden[..., :1] + return {"logits": logits, "hidden_states": (hidden,)} + + model = SALMAutomodel.__new__(SALMAutomodel) + torch.nn.Module.__init__(model) + model.cfg = {} + model._fused_linear_cross_entropy = object() + model.llm = FakeLLM() + model.train() + inputs = torch.randn(1, 5, 4) + + outputs = model.forward(inputs) + + assert model.llm.kwargs["compute_logits"] is False + assert model.llm.kwargs["output_hidden_states"] is True + assert "compute_mtp" not in model.llm.kwargs + torch.testing.assert_close(outputs["hidden_states"], inputs + 3) + assert outputs["logits"].shape == (1, 5, 0) + + +def test_salm_automodel_fused_linear_loss_consumes_hidden_states_and_lm_weight(): + calls = [] + + class FakeFusedLoss: + def __call__(self, hidden_states, target_ids, weight, grad_reduce_group): + calls.append((hidden_states, target_ids, weight, grad_reduce_group)) + return hidden_states.new_tensor(7.0) + + class FakeLLM(torch.nn.Module): + def __init__(self): + super().__init__() + self.lm_head = torch.nn.Linear(3, 5, bias=False) + + def get_output_embeddings(self): + return self.lm_head + + model = SALMAutomodel.__new__(SALMAutomodel) + torch.nn.Module.__init__(model) + model.llm = FakeLLM() + model._fused_linear_cross_entropy = FakeFusedLoss() + hidden = torch.randn(1, 4, 3) + targets = torch.tensor([[0, 1, -100, 2]]) + group = object() + + loss_sum, logits = model._compute_training_cross_entropy_sum( + {"hidden_states": hidden, "logits": torch.empty(0)}, targets, group + ) + + assert loss_sum.item() == 7.0 + assert logits is None + assert calls == [(hidden, targets, model.llm.lm_head.weight, group)] + + +def test_salm_automodel_notifies_garbage_collection_after_optimizer_step(monkeypatch): + calls = [] + + class FakeGarbageCollectionManager: + def on_optimizer_step(self): + calls.append("gc") + + model = SALMAutomodel.__new__(SALMAutomodel) + torch.nn.Module.__init__(model) + model._garbage_collection = FakeGarbageCollectionManager() + monkeypatch.setattr( + LightningModule, + "optimizer_step", + lambda *args, **kwargs: calls.append("optimizer"), + ) + model.optimizer_step(0, 0, object()) + assert calls == ["optimizer", "gc"] + + def test_salm_automodel_record_training_stats_uses_thd_metadata(): model = SALMAutomodel.__new__(SALMAutomodel) batch = {"input_ids": torch.zeros(3, 7, dtype=torch.long)} @@ -460,6 +599,65 @@ def test_salm_automodel_prepare_inputs_skips_chunking_when_size_is_null(device): assert torch.equal(input_signal_lens, torch.tensor([5], dtype=torch.long, device=device)) +@pytest.mark.parametrize("native_dataset_batch", [False, True]) +@pytest.mark.parametrize("device", chunking_test_devices()) +def test_salm_automodel_packed_no_chunking_embeds_only_real_tokens(monkeypatch, device, native_dataset_batch): + """Packed no-chunking batches compact IDs before the embedding lookup.""" + model = _make_chunking_test_model(encoder_chunk_size_seconds=None, sampling_rate=2, device=device) + model.cfg["packed_sequences"] = True + batch_size = 4 + sequence_length = 64 + row_lengths = [64, 8, 4, 2] + input_ids = torch.full((batch_size, sequence_length), model.text_pad_id, dtype=torch.long, device=device) + for row, length in enumerate(row_lengths): + tokens = torch.arange(10, 10 + length, dtype=torch.long, device=device) + tokens[0] = model.audio_locator_tag_id + input_ids[row, -length:] = tokens + loss_mask = input_ids != model.text_pad_id + loss_mask[input_ids == model.audio_locator_tag_id] = False + audios = torch.arange(1, batch_size * 3 + 1, dtype=torch.float32, device=device).reshape(batch_size, 3) + batch = { + "audio_lens": torch.full((batch_size,), 3, dtype=torch.long, device=device), + "input_ids": input_ids, + "loss_mask": loss_mask, + } + if native_dataset_batch: + batch.update( + { + "packed_audio_samples": audios.flatten(), + "audio_cu_seqlens": torch.arange(0, batch_size * 3 + 1, 3, dtype=torch.long, device=device), + "input_ids": torch.cat([row[-length:] for row, length in zip(input_ids, row_lengths)]), + "loss_mask": torch.cat([row[-length:] for row, length in zip(loss_mask, row_lengths)]), + "text_cu_seqlens": torch.tensor( + [0, *torch.tensor(row_lengths, device=device).cumsum(0).tolist()], + dtype=torch.long, + device=device, + ), + } + ) + else: + batch["audios"] = audios + original_embed_tokens = model._embed_tokens + embedded_shapes = [] + + def embed_tokens(flat_ids): + embedded_shapes.append(tuple(flat_ids.shape)) + return original_embed_tokens(flat_ids) + + monkeypatch.setattr(model, "_embed_tokens", embed_tokens) + + inputs = model.prepare_inputs(batch) + + real_token_count = sum(row_lengths) + assert embedded_shapes == [(real_token_count,)] + if not native_dataset_batch: + assert real_token_count < input_ids.numel() + assert inputs["input_embeds"].ndim == 2 + inputs["input_embeds"].sum().backward() + assert model.embed_tokens.weight.grad is not None + assert model.embed_tokens.weight.grad.abs().sum() > 0 + + @pytest.mark.parametrize("device", chunking_test_devices()) def test_salm_automodel_prepare_inputs_preserves_chunked_audio_order(device): model = _make_chunking_test_model(encoder_chunk_size_seconds=1.0, sampling_rate=2, device=device) @@ -512,6 +710,71 @@ def test_salm_automodel_generate_chunks_audio_before_llm(device): assert answer.shape == (1, 3) +@pytest.mark.parametrize("device", chunking_test_devices()) +def test_salm_automodel_limits_packed_encoder_opt_in_to_training(device): + model = _make_chunking_test_model(encoder_chunk_size_seconds=1.0, sampling_rate=2, device=device) + model.cfg["packed_encoder_sequences"] = True + batch = { + "audios": torch.tensor([[1.0, 2.0, 3.0, 4.0, 5.0]], device=device), + "audio_lens": torch.tensor([5], dtype=torch.long, device=device), + "input_ids": torch.tensor([[model.audio_locator_tag_id, 10]], dtype=torch.long, device=device), + "loss_mask": torch.tensor([[False, True]], dtype=torch.bool, device=device), + } + + inputs = model.prepare_inputs(batch) + + assert model.perception.sequence_packed_calls == 1 + assert torch.equal( + inputs["input_embeds"][0, :, 0], + torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0], device=device), + ) + + answer = model.generate( + prompts=torch.tensor([[model.audio_locator_tag_id, 10]], dtype=torch.long, device=device), + audios=torch.tensor([[6.0, 7.0, 8.0, 9.0, 10.0]], device=device), + audio_lens=torch.tensor([5], dtype=torch.long, device=device), + max_new_tokens=2, + ) + + assert model.perception.sequence_packed_calls == 1 + assert torch.equal( + model.llm.generate_kwargs["inputs_embeds"][0, :5, 0], + torch.tensor([6.0, 7.0, 8.0, 9.0, 10.0], device=device), + ) + assert answer.shape == (1, 2) + + +@pytest.mark.parametrize("device", chunking_test_devices()) +def test_salm_automodel_packed_audio_samples_match_padded_batch(device): + padded_model = _make_chunking_test_model(encoder_chunk_size_seconds=1.0, sampling_rate=2, device=device) + packed_model = _make_chunking_test_model(encoder_chunk_size_seconds=1.0, sampling_rate=2, device=device) + padded_model.cfg["packed_encoder_sequences"] = True + packed_model.cfg["packed_encoder_sequences"] = True + audios = torch.tensor([[1.0, 2.0, 3.0, 0.0, 0.0], [10.0, 11.0, 12.0, 13.0, 14.0]], device=device) + audio_lens = torch.tensor([3, 5], dtype=torch.long, device=device) + common = { + "audio_lens": audio_lens, + "input_ids": torch.tensor( + [[padded_model.audio_locator_tag_id, padded_model.audio_locator_tag_id, 10]], + dtype=torch.long, + device=device, + ), + "loss_mask": torch.tensor([[False, False, True]], dtype=torch.bool, device=device), + } + padded_batch = {**common, "audios": audios} + packed_batch = { + **common, + "packed_audio_samples": torch.cat([audios[0, :3], audios[1, :5]]), + "audio_cu_seqlens": torch.tensor([0, 3, 8], dtype=torch.long, device=device), + } + + expected = padded_model.prepare_inputs(padded_batch) + actual = packed_model.prepare_inputs(packed_batch) + + torch.testing.assert_close(actual["input_embeds"], expected["input_embeds"], rtol=0.0, atol=0.0) + assert torch.equal(actual["target_ids"], expected["target_ids"]) + + def _make_chunking_test_model(encoder_chunk_size_seconds, sampling_rate, device, hop_length=1): model = SALMAutomodel.__new__(SALMAutomodel) torch.nn.Module.__init__(model) diff --git a/tests/collections/speechlm2/test_salm_automodel_mtp.py b/tests/collections/speechlm2/test_salm_automodel_mtp.py index 1b1e83f6489a..f824ac125fae 100644 --- a/tests/collections/speechlm2/test_salm_automodel_mtp.py +++ b/tests/collections/speechlm2/test_salm_automodel_mtp.py @@ -16,6 +16,7 @@ import pytest import torch +from omegaconf import DictConfig import nemo.collections.speechlm2.models.salm_automodel as salm_module import nemo.collections.speechlm2.parts.mtp as mtp_module @@ -27,7 +28,8 @@ def _bare_model(): '''Create a SALMAutomodel instance without loading any weights.''' model = SALMAutomodel.__new__(SALMAutomodel) torch.nn.Module.__init__(model) - model.cfg = {"debug_log_training_batches": 0} + model.cfg = DictConfig({"debug_log_training_batches": 0}) + model._fused_linear_cross_entropy = None return model @@ -698,6 +700,82 @@ def _calculate_mtp_loss(*_args, **kwargs): assert captured_kwargs["cu_seqlens"] is None +def test_training_step_shares_one_materialized_lm_weight_with_main_and_mtp_losses(monkeypatch): + class _Perception(torch.nn.Module): + def __init__(self): + super().__init__() + self.preprocessor = torch.nn.Identity() + self.encoder = torch.nn.Identity() + + class _LLM(torch.nn.Module): + def __init__(self): + super().__init__() + self.lm_head = torch.nn.Linear(4, 8, bias=False) + + def get_output_embeddings(self): + return self.lm_head + + class _FusedLoss: + def __init__(self): + self.materialize_calls = [] + self.loss_weights = [] + self.shared_weight = torch.randn(8, 4, requires_grad=True) + + def materialize_lm_weight(self, weight, *, grad_reduce_group): + self.materialize_calls.append((weight, grad_reduce_group)) + return self.shared_weight + + def __call__(self, hidden_states, labels, weight, *, grad_reduce_group): + self.loss_weights.append(weight) + return hidden_states.sum() * 0 + 1.0 + + model = _bare_model() + model.perception = _Perception() + model.llm = _LLM() + model.lss_loss = None + fused_loss = _FusedLoss() + model._fused_linear_cross_entropy = fused_loss + model._mtp_loss_fn = fused_loss + model._mtp_loss_scaling_factor = 0.1 + model._trainer = None + model.tokenizer = type("Tokenizer", (), {"pad": -1, "unk_id": None})() + model._get_moe_dp_group = lambda: None + model.log = lambda *_args, **_kwargs: None + model.log_dict = lambda *_args, **_kwargs: None + model.maybe_log_moe_metrics = lambda _batch_idx: None + + inputs = { + "input_embeds": torch.zeros(5, 4), + "attention_mask": None, + "target_ids": torch.tensor([0, 1, 2, 3, 4]), + "llm_kwargs": {}, + "num_tokens": 5, + "num_examples": 1, + } + model.prepare_inputs = lambda _batch: inputs + model.forward = lambda *_args, **_kwargs: { + "logits": torch.empty(1, 5, 0), + "hidden_states": torch.zeros(1, 5, 4, requires_grad=True), + "mtp_per_depth_h": [torch.zeros(1, 5, 4, requires_grad=True)], + } + captured_kwargs = {} + + def _calculate_mtp_loss(*_args, **kwargs): + captured_kwargs.update(kwargs) + return SimpleNamespace( + loss=torch.tensor(0.25, requires_grad=True), + per_depth_losses=[torch.tensor(2.5, requires_grad=True)], + ) + + monkeypatch.setattr(salm_module, "calculate_mtp_loss_with_per_depth", _calculate_mtp_loss) + + model._training_step_batch({"input_ids": torch.tensor([[1, 2, 3, 4, 5]])}, batch_idx=0) + + assert fused_loss.materialize_calls == [(model.llm.lm_head.weight, None)] + assert fused_loss.loss_weights == [fused_loss.shared_weight] + assert captured_kwargs["lm_weight"] is fused_loss.shared_weight + + def test_mtp_validation_forward_uses_and_restores_native_gate(): llm = torch.nn.Module() llm.eval() diff --git a/tests/collections/speechlm2/test_salm_automodel_pee.py b/tests/collections/speechlm2/test_salm_automodel_pee.py index a4ef1465e344..f4c2b82cdef5 100644 --- a/tests/collections/speechlm2/test_salm_automodel_pee.py +++ b/tests/collections/speechlm2/test_salm_automodel_pee.py @@ -396,6 +396,46 @@ def test_pee_prepare_inputs_routes_spk_targets_as_spk_targets(dummy_pe_encoder): assert model.perception.spk_targets_calls[-1] is None +@pytest.mark.unit +@pytest.mark.parametrize( + ("packed_encoder_sequences", "expected_outer_chunk_size"), + [(False, 30.0), (True, None)], +) +def test_pee_prepare_inputs_routes_shared_chunk_size_at_the_correct_layer( + dummy_pe_encoder, monkeypatch, packed_encoder_sequences, expected_outer_chunk_size +): + from nemo.collections.speechlm2.parts import cp_helpers + + model = _make_pee_routing_test_model( + dummy_pe_encoder, + cfg={ + "encoder_chunk_size_seconds": 30.0, + "encoder_chunk_batch_size": None, + "packed_encoder_sequences": packed_encoder_sequences, + }, + ) + batch = { + "audios": torch.tensor([[1.0, 2.0, 3.0, 4.0, 5.0]]), + "audio_lens": torch.tensor([5], dtype=torch.long), + "input_ids": torch.tensor([[model.audio_locator_tag_id, 10]], dtype=torch.long), + "loss_mask": torch.tensor([[False, True]], dtype=torch.bool), + } + recorded = {} + original = cp_helpers.encode_audio_with_cp_distribution + + def record_chunk_routing(*args, **kwargs): + recorded["chunk_size_seconds"] = kwargs["chunk_size_seconds"] + # This test isolates routing. Other tests exercise the native packed path. + kwargs["sequence_packed"] = False + return original(*args, **kwargs) + + monkeypatch.setattr(cp_helpers, "encode_audio_with_cp_distribution", record_chunk_routing) + + model.prepare_inputs(batch) + + assert recorded["chunk_size_seconds"] == expected_outer_chunk_size + + @pytest.mark.unit def test_pee_generation_warns_that_outer_chunking_is_ignored(dummy_pe_encoder): model = _make_pee_routing_test_model( diff --git a/tests/collections/speechlm2/test_salm_cp_helpers.py b/tests/collections/speechlm2/test_salm_cp_helpers.py index 3cfaba3c1a62..47334b55bb7c 100644 --- a/tests/collections/speechlm2/test_salm_cp_helpers.py +++ b/tests/collections/speechlm2/test_salm_cp_helpers.py @@ -22,6 +22,7 @@ import pytest import torch +from nemo.collections.asr.parts.packed_sequence import pack_encoder_output from nemo.collections.speechlm2.parts.cp_helpers import encode_audio_with_cp_distribution, get_cp_mesh @@ -158,6 +159,8 @@ def __init__(self): self.last_input_signal_shape = None self.last_input_signal_length = None self.spk_targets_calls = [] + self.supports_sequence_packed_output = True + self.sequence_packed_calls = 0 def forward(self, *, input_signal, input_signal_length, spk_targets=None): self.num_calls += 1 @@ -169,6 +172,19 @@ def forward(self, *, input_signal, input_signal_length, spk_targets=None): lens = torch.full((B,), embs.shape[1], dtype=input_signal_length.dtype, device=input_signal_length.device) return embs, lens + def forward_sequence_packed(self, **kwargs): + self.sequence_packed_calls += 1 + cu_seqlens = kwargs.pop("input_signal_cu_seqlens", None) + if cu_seqlens is not None: + offsets = cu_seqlens.tolist() + rows = [ + kwargs["input_signal"][offsets[row] : offsets[row + 1]] + for row in range(kwargs["input_signal_length"].numel()) + ] + kwargs["input_signal"] = torch.nn.utils.rnn.pad_sequence(rows, batch_first=True) + embs, lens = self.forward(**kwargs) + return pack_encoder_output(embs, lens) + def test_encode_audio_cp_distribution_preserves_local_autograd(monkeypatch): perception = _TrainablePerceptionStub() @@ -304,3 +320,78 @@ def fake_all_reduce(tensor, op=None, group=None): assert perception.num_calls == 1 assert all_reduce_calls == [(1, "fake-fsdp-group")] assert dummy_audio_loss is None + + +def test_sequence_packed_cp_uses_flat_gather_and_preserves_local_autograd(monkeypatch): + perception = _TrainablePerceptionStub() + audios = torch.tensor([[1.0, 2.0, 0.0], [3.0, 4.0, 0.0]]) + audio_lens = torch.tensor([3, 3], dtype=torch.long) + gathered_shapes = [] + + def fake_all_gather(local_flat, group): + assert group == "fake-cp-group" + gathered_shapes.append(tuple(local_flat.shape)) + return (local_flat, torch.full_like(local_flat, 7.0)) + + def fake_lens_all_gather(gathered_lens, local_lens, group): + assert group == "fake-cp-group" + gathered_lens[0].copy_(local_lens) + gathered_lens[1].copy_(local_lens) + + monkeypatch.setattr("nemo.collections.speechlm2.parts.cp_helpers.dist.get_rank", lambda group: 0) + monkeypatch.setattr("nemo.collections.speechlm2.parts.cp_helpers.dist.all_reduce", lambda *args, **kwargs: None) + monkeypatch.setattr("nemo.collections.speechlm2.parts.cp_helpers.dist.all_gather", fake_lens_all_gather) + monkeypatch.setattr("nemo.collections.speechlm2.parts.cp_helpers.differentiable_all_gather", fake_all_gather) + + embs = encode_audio_with_cp_distribution( + perception, + audios, + audio_lens, + chunk_size_seconds=None, + sampling_rate=16000, + cp_mesh=_FakeCpMesh(), + sequence_packed=True, + packed_cp_gather=True, + ) + + assert perception.sequence_packed_calls == 1 + assert gathered_shapes == [(2, 1)] + assert len(embs) == 2 + torch.testing.assert_close(embs[1], torch.full((2, 1), 7.0)) + embs[0].sum().backward() + assert perception.scale.grad.item() == pytest.approx(3.0) + + +def test_packed_waveform_cp_slices_before_local_frontend(monkeypatch): + perception = _TrainablePerceptionStub() + audios = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0]) + audio_lens = torch.tensor([2, 5], dtype=torch.long) + audio_cu_seqlens = torch.tensor([0, 2, 7], dtype=torch.long) + + def fake_all_gather(local_flat, group): + return (local_flat, torch.zeros_like(local_flat)) + + def fake_lens_all_gather(gathered_lens, local_lens, group): + gathered_lens[0].copy_(local_lens) + gathered_lens[1].copy_(local_lens) + + monkeypatch.setattr("nemo.collections.speechlm2.parts.cp_helpers.dist.get_rank", lambda group: 0) + monkeypatch.setattr("nemo.collections.speechlm2.parts.cp_helpers.dist.all_reduce", lambda *args, **kwargs: None) + monkeypatch.setattr("nemo.collections.speechlm2.parts.cp_helpers.dist.all_gather", fake_lens_all_gather) + monkeypatch.setattr("nemo.collections.speechlm2.parts.cp_helpers.differentiable_all_gather", fake_all_gather) + + embs = encode_audio_with_cp_distribution( + perception, + audios, + audio_lens, + audio_cu_seqlens=audio_cu_seqlens, + chunk_size_seconds=None, + sampling_rate=16000, + cp_mesh=_FakeCpMesh(), + sequence_packed=True, + packed_cp_gather=True, + ) + + assert perception.last_input_signal_shape == (1, 2) + assert perception.last_input_signal_length == [2] + assert len(embs) == 2 diff --git a/tests/collections/speechlm2/test_salm_cp_helpers_audio_free_distributed.py b/tests/collections/speechlm2/test_salm_cp_helpers_audio_free_distributed.py new file mode 100644 index 000000000000..f4a838f992f1 --- /dev/null +++ b/tests/collections/speechlm2/test_salm_cp_helpers_audio_free_distributed.py @@ -0,0 +1,94 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from torch.distributed.fsdp import FullyShardedDataParallel as FSDP + +from nemo.collections.asr.parts.packed_sequence import pack_encoder_output +from nemo.collections.speechlm2.parts.cp_helpers import encode_audio_with_cp_distribution + + +class _Scale(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.full((2,), 2.0)) + + def forward(self, inputs): + return inputs.unsqueeze(-1) * self.weight.mean() + + +class _PackedPerception(torch.nn.Module): + supports_sequence_packed_output = True + + def __init__(self, device): + super().__init__() + self.core = FSDP(_Scale().to(device), device_id=device) + + def forward_sequence_packed(self, *, input_signal, input_signal_length, **kwargs): + return pack_encoder_output(self.core(input_signal), input_signal_length) + + +def _run_audio_free_chunk_test(rank: int, world_size: int, init_file: str): + torch.cuda.set_device(rank) + dist.init_process_group("nccl", init_method=f"file://{init_file}", rank=rank, world_size=world_size) + try: + device = torch.device("cuda", rank) + perception = _PackedPerception(device) + if rank == 0: + audios = torch.tensor([[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]], device=device) + lengths = torch.tensor([6], dtype=torch.long, device=device) + else: + audios = torch.empty(0, 6, device=device) + lengths = torch.empty(0, dtype=torch.long, device=device) + + embeddings, dummy_loss = encode_audio_with_cp_distribution( + perception, + audios, + lengths, + chunk_size_seconds=2.0, + chunk_batch_size=2, + sampling_rate=1, + cp_mesh=None, + fsdp_sync_group=dist.group.WORLD, + return_dummy_loss=True, + sequence_packed=True, + ) + + if rank == 0: + assert dummy_loss is None + assert len(embeddings) == 1 + torch.testing.assert_close(embeddings[0][:, 0], audios[0] * 2.0) + loss = embeddings[0].sum() + else: + assert embeddings == [] + assert dummy_loss is not None + loss = dummy_loss + loss.backward() + + grad = next(perception.core.parameters()).grad + assert grad is not None + gathered = [torch.zeros_like(grad) for _ in range(world_size)] + dist.all_gather(gathered, grad) + torch.testing.assert_close(gathered[0], gathered[1]) + assert torch.isfinite(gathered[0]).all() and gathered[0].abs().sum() > 0 + finally: + dist.destroy_process_group() + + +@pytest.mark.skipif(not torch.cuda.is_available() or torch.cuda.device_count() < 2, reason="Test requires 2 GPUs") +def test_packed_chunk_microbatches_keep_audio_free_fsdp_rank_collectives_and_backward_aligned(tmp_path): + mp.spawn(_run_audio_free_chunk_test, args=(2, str(tmp_path / "audio_free_init")), nprocs=2, join=True) diff --git a/tests/collections/speechlm2/test_salm_cp_helpers_distributed.py b/tests/collections/speechlm2/test_salm_cp_helpers_distributed.py new file mode 100644 index 000000000000..88f541f1d1c1 --- /dev/null +++ b/tests/collections/speechlm2/test_salm_cp_helpers_distributed.py @@ -0,0 +1,131 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from nemo.collections.asr.parts.packed_sequence import pack_encoder_output +from nemo.collections.speechlm2.parts.cp_helpers import encode_audio_with_cp_distribution + + +class _WorldCpMesh: + def size(self): + return dist.get_world_size() + + def get_group(self): + return dist.group.WORLD + + +class _PackedPerception(torch.nn.Module): + supports_sequence_packed_output = True + + def __init__(self): + super().__init__() + self.scale = torch.nn.Parameter(torch.tensor(2.0, device='cuda')) + + def forward_sequence_packed( + self, + *, + input_signal, + input_signal_length, + input_signal_cu_seqlens=None, + **kwargs, + ): + if input_signal_cu_seqlens is not None: + offsets = input_signal_cu_seqlens.tolist() + input_signal = torch.nn.utils.rnn.pad_sequence( + [input_signal[offsets[row] : offsets[row + 1]] for row in range(input_signal_length.numel())], + batch_first=True, + ) + padded = input_signal.unsqueeze(-1) * self.scale + return pack_encoder_output(padded, input_signal_length) + + +def _run_remote_gradient_test(rank: int, world_size: int, init_file: str): + torch.cuda.set_device(rank) + dist.init_process_group('nccl', init_method=f'file://{init_file}', rank=rank, world_size=world_size) + try: + cases = [ + ( + [[1.0, 2.0], [3.0, 4.0]], + [2, 2], + lambda outputs: outputs[1 - rank].sum(), + [3.0, 7.0], + ), + ( + [[1.0, 2.0, 3.0], [4.0, 0.0, 0.0], [5.0, 6.0, 0.0]], + [3, 1, 2], + lambda outputs: outputs[2].sum() if rank == 0 else outputs[0].sum() + outputs[1].sum(), + [10.0, 11.0], + ), + ( + [[2.0, 3.0]], + [2], + lambda outputs: outputs[0].sum(), + [10.0, 0.0], + ), + ( + [[0.0, 0.0, 0.0], [4.0, 5.0, 6.0]], + [0, 3], + lambda outputs: sum(output.sum() for output in outputs), + [0.0, 30.0], + ), + ( + [[0.0, 0.0, 0.0]], + [0], + lambda outputs: outputs[0].sum(), + [0.0, 0.0], + ), + ] + for audio_rows, lengths, make_loss, expected_grads in cases: + padded_audios = torch.tensor(audio_rows, device='cuda') + audio_lens = torch.tensor(lengths, dtype=torch.long, device='cuda') + for pack_waveforms in (False, True): + perception = _PackedPerception() + if pack_waveforms: + audios = torch.cat([audio[:length] for audio, length in zip(padded_audios, lengths)]) + audio_cu_seqlens = torch.cat([audio_lens.new_zeros(1), audio_lens.cumsum(dim=0)]) + else: + audios = padded_audios + audio_cu_seqlens = None + embeddings = encode_audio_with_cp_distribution( + perception, + audios, + audio_lens, + audio_cu_seqlens=audio_cu_seqlens, + chunk_size_seconds=None, + sampling_rate=16_000, + cp_mesh=_WorldCpMesh(), + sequence_packed=True, + packed_cp_gather=True, + ) + + assert [embedding.shape[0] for embedding in embeddings] == lengths + for embedding, audio, length in zip(embeddings, padded_audios, lengths): + torch.testing.assert_close(embedding[:, 0], audio[:length] * 2.0) + make_loss(embeddings).backward() + + grad = perception.scale.grad.detach() + gathered_grads = [torch.zeros_like(grad) for _ in range(world_size)] + dist.all_gather(gathered_grads, grad) + torch.testing.assert_close(torch.stack(gathered_grads).cpu(), torch.tensor(expected_grads)) + finally: + dist.destroy_process_group() + + +@pytest.mark.skipif(not torch.cuda.is_available() or torch.cuda.device_count() < 2, reason='Test requires 2 GPUs') +def test_packed_cp_gather_handles_uneven_batches_dummies_and_remote_gradients(tmp_path): + mp.spawn(_run_remote_gradient_test, args=(2, str(tmp_path / 'cp_init')), nprocs=2, join=True) diff --git a/tests/collections/speechlm2/test_salm_dataset.py b/tests/collections/speechlm2/test_salm_dataset.py index 60b61ec5ac7b..83151be0ec6a 100644 --- a/tests/collections/speechlm2/test_salm_dataset.py +++ b/tests/collections/speechlm2/test_salm_dataset.py @@ -13,9 +13,13 @@ # See the License for the specific language governing permissions and # limitations under the License. +import io + +import numpy as np import pytest +import soundfile as sf import torch -from lhotse import CutSet, SupervisionSegment +from lhotse import CutSet, Recording, SupervisionSegment, fastcopy from lhotse.testing.dummies import dummy_cut, dummy_recording import nemo.collections.speechlm2.data.salm_dataset as salm_dataset_module @@ -40,6 +44,233 @@ def test_multispeaker_config_requires_explicit_speaker_count(): salm_dataset_module.MultiSpeakerConfig.from_dict({}) +@pytest.mark.unit +def test_multispeaker_targets_precede_in_memory_audio_drop(monkeypatch): + audio = io.BytesIO() + sf.write(audio, np.zeros((640, 2), dtype=np.float32), 16000, format="WAV") + cut = Recording.from_bytes(audio.getvalue(), recording_id="native-wds").to_cut() + cut = fastcopy( + cut, + custom={ + "_source_read_key": "/data/shard-0.tar#sample-0", + "_source_range_bytes": len(audio.getvalue()), + }, + ) + conversation = NeMoMultimodalConversation( + id="example-0", + turns=[ + AudioTurn(role="user", cut=cut, audio_locator_tag="<|audio|>"), + TextTurn(role="assistant", value="hello"), + ], + token_equivalent_duration=0.01, + ) + conversation.input_ids = torch.tensor([7, 8], dtype=torch.long) + conversation.mask = torch.tensor([False, True]) + conversations = CutSet([conversation]) + + def fake_audio_collate(conversations_arg, *args, **kwargs): + return torch.zeros(1, 640), torch.tensor([640]), conversations_arg + + def fake_speaker_activity_from_cut(materialized_cut, **kwargs): + assert materialized_cut.load_audio().shape == (1, 640) + return torch.zeros(4, 2) + + monkeypatch.setattr( + salm_dataset_module, + "collate_conversation_audio_fault_tolerant", + fake_audio_collate, + ) + monkeypatch.setattr( + salm_dataset_module, + "speaker_activity_from_cut", + fake_speaker_activity_from_cut, + ) + dataset = salm_dataset_module.SALMDataset( + tokenizer=_Tokenizer(), + multispeaker_cfg={ + "num_speakers": 2, + "sample_rate": 16000, + "window_stride": 0.01, + "subsampling_factor": 1, + }, + ) + + batch = dataset[conversations] + + assert torch.all(batch["spk_targets"] == -1.0) + (returned_cut,) = next(iter(batch["conversations"])).list_cuts() + assert returned_cut.recording.sources[0].type == "shar" + assert returned_cut.custom["_source_read_key"] == "/data/shard-0.tar#sample-0" + + +@pytest.mark.unit +def test_salm_dataset_can_return_packed_audio_samples(monkeypatch): + cut = dummy_cut(0, duration=0.03, recording=dummy_recording(0, duration=0.03, with_data=True)) + conversation = NeMoMultimodalConversation( + id="example-0", + turns=[ + AudioTurn(role="user", cut=cut, audio_locator_tag="<|audio|>"), + TextTurn(role="assistant", value="hello"), + ], + token_equivalent_duration=0.01, + ) + conversation.input_ids = torch.tensor([7, 8], dtype=torch.long) + conversation.mask = torch.tensor([False, True]) + conversations = CutSet([conversation]) + expected_samples = torch.tensor([1.0, 2.0, 3.0]) + expected_cu_seqlens = torch.tensor([0, 3], dtype=torch.long) + expected_lens = torch.tensor([3], dtype=torch.long) + + def fake_packed_audio_collate(conversations_arg, *args, **kwargs): + assert conversations_arg is conversations + return expected_samples, expected_cu_seqlens, expected_lens, conversations_arg + + monkeypatch.setattr( + salm_dataset_module, + "collate_conversation_audio_packed_fault_tolerant", + fake_packed_audio_collate, + ) + dataset = salm_dataset_module.SALMDataset(tokenizer=_Tokenizer(), pack_audio=True) + + batch = dataset[conversations] + + assert "audios" not in batch + assert batch["packed_audio_samples"] is expected_samples + assert batch["audio_cu_seqlens"] is expected_cu_seqlens + assert batch["audio_lens"] is expected_lens + assert batch["input_ids"].shape == (1, 2) + + +@pytest.mark.unit +def test_salm_dataset_packed_sequences_never_pads_variable_length_tensors(monkeypatch): + conversations = [] + for idx, (input_ids, mask) in enumerate( + ( + ([7, 8, 9], [False, True, True]), + ([10, 11], [False, True]), + ) + ): + cut = dummy_cut(idx, duration=0.03, recording=dummy_recording(idx, duration=0.03, with_data=True)) + conversation = NeMoMultimodalConversation( + id=f"example-{idx}", + turns=[ + AudioTurn(role="user", cut=cut, audio_locator_tag="<|audio|>"), + TextTurn(role="assistant", value="hello"), + ], + token_equivalent_duration=0.01, + ) + conversation.input_ids = torch.tensor(input_ids, dtype=torch.long) + conversation.mask = torch.tensor(mask) + conversations.append(conversation) + conversations = CutSet(conversations) + + def fake_packed_audio_collate(conversations_arg, *args, **kwargs): + return ( + torch.tensor([1.0, 2.0, 10.0, 11.0, 12.0]), + torch.tensor([0, 2, 5], dtype=torch.long), + torch.tensor([2, 3], dtype=torch.long), + conversations_arg, + ) + + monkeypatch.setattr( + salm_dataset_module, + "collate_conversation_audio_packed_fault_tolerant", + fake_packed_audio_collate, + ) + dataset = salm_dataset_module.SALMDataset(tokenizer=_Tokenizer(), pack_sequences=True) + + batch = dataset[conversations] + + assert dataset.pack_audio is True + assert "audios" not in batch + assert batch["packed_audio_samples"].shape == (5,) + assert torch.equal(batch["audio_cu_seqlens"], torch.tensor([0, 2, 5])) + assert torch.equal(batch["input_ids"], torch.tensor([7, 8, 9, 10, 11])) + assert torch.equal(batch["loss_mask"], torch.tensor([False, True, True, False, True])) + assert torch.equal(batch["text_cu_seqlens"], torch.tensor([0, 3, 5])) + + +@pytest.mark.unit +def test_salm_dataset_packs_multispeaker_targets_without_time_padding(monkeypatch): + conversations = [] + for idx in range(2): + cut = dummy_cut(idx, duration=0.03, recording=dummy_recording(idx, duration=0.03, with_data=True)) + conversation = NeMoMultimodalConversation( + id=f"example-{idx}", + turns=[ + AudioTurn(role="user", cut=cut, audio_locator_tag="<|audio|>"), + TextTurn(role="assistant", value="hello"), + ], + token_equivalent_duration=0.01, + ) + conversation.input_ids = torch.tensor([7 + idx, 9], dtype=torch.long) + conversation.mask = torch.tensor([False, True]) + conversations.append(conversation) + conversations = CutSet(conversations) + + def fake_packed_audio_collate(conversations_arg, *args, **kwargs): + return ( + torch.arange(6, dtype=torch.float32), + torch.tensor([0, 3, 6], dtype=torch.long), + torch.tensor([3, 3], dtype=torch.long), + conversations_arg, + ) + + activity_lengths = iter((3, 1)) + + def fake_speaker_activity_from_cut(cut, **kwargs): + return torch.ones(next(activity_lengths), 1) + + monkeypatch.setattr( + salm_dataset_module, + "collate_conversation_audio_packed_fault_tolerant", + fake_packed_audio_collate, + ) + monkeypatch.setattr(salm_dataset_module, "speaker_activity_from_cut", fake_speaker_activity_from_cut) + dataset = salm_dataset_module.SALMDataset( + tokenizer=_Tokenizer(), + pack_sequences=True, + multispeaker_cfg={"num_speakers": 2}, + ) + + batch = dataset[conversations] + + assert batch["spk_targets"].shape == (4, 2) + assert torch.all(batch["spk_targets"] == -1.0) + assert torch.equal(batch["spk_target_length"], torch.tensor([3, 1])) + assert torch.equal(batch["spk_target_cu_seqlens"], torch.tensor([0, 3, 4])) + + +@pytest.mark.unit +def test_salm_dataset_reports_exact_packing_efficiency(monkeypatch): + conversations = [] + for idx, num_tokens in enumerate((3072, 7168)): + cut = dummy_cut(idx, duration=0.03, recording=dummy_recording(idx, duration=0.03, with_data=True)) + conversation = NeMoMultimodalConversation( + id=f"example-{idx}", + turns=[ + AudioTurn(role="user", cut=cut, audio_locator_tag="<|audio|>"), + TextTurn(role="assistant", value="hello"), + ], + token_equivalent_duration=0.01, + ) + conversation.input_ids = torch.tensor([7, 8], dtype=torch.long) + conversation.mask = torch.tensor([False, True]) + conversation.num_tokens = num_tokens + conversations.append(conversation) + conversations = CutSet(conversations) + + def fake_audio_collate(conversations_arg, *args, **kwargs): + return torch.zeros(2, 480), torch.tensor([480, 480]), conversations_arg + + monkeypatch.setattr(salm_dataset_module, "collate_conversation_audio_fault_tolerant", fake_audio_collate) + dataset = salm_dataset_module.SALMDataset(tokenizer=_Tokenizer(), batch_tokens=12288) + + batch = dataset[conversations] + + assert batch["packing_efficiency"].item() == pytest.approx(10240 / 12288) + + @pytest.mark.unit @pytest.mark.parametrize( ("rttm_filepath", "expected_targets"), diff --git a/tests/collections/speechlm2/test_salm_packed_sequences.py b/tests/collections/speechlm2/test_salm_packed_sequences.py index 6bf57213765d..6b3161fd559e 100644 --- a/tests/collections/speechlm2/test_salm_packed_sequences.py +++ b/tests/collections/speechlm2/test_salm_packed_sequences.py @@ -79,6 +79,91 @@ def _basic_batch(): return input_ids, embeds, target_ids, replacements +def test_prepare_packed_llm_inputs_compacts_ids_before_embedding_and_preserves_backward(): + """Embedding work scales with real tokens, not the dense B*S rectangle.""" + batch_size = 8 + sequence_length = 257 + hidden_size = 4 + vocab_size = 64 + row_lengths = [257, 8, 7, 6, 5, 4, 3, 2] + input_ids = torch.full((batch_size, sequence_length), PAD, dtype=torch.long) + for row, length in enumerate(row_lengths): + input_ids[row, -length:] = torch.arange(1, length + 1).remainder(vocab_size - 1).add(1) + target_ids = input_ids.where(input_ids != PAD, -100) + embedding = torch.nn.Embedding(vocab_size, hidden_size) + embedded_shapes = [] + + def embed_tokens(flat_ids): + embedded_shapes.append(tuple(flat_ids.shape)) + return embedding(flat_ids) + + actual = prepare_packed_llm_inputs( + input_ids=input_ids, + text_embs=None, + audio_embs=[], + target_ids=target_ids, + padding_id=PAD, + placeholder_id=AUDIO, + embed_tokens=embed_tokens, + ) + + real_token_count = sum(row_lengths) + assert embedded_shapes == [(real_token_count,)] + assert real_token_count < input_ids.numel() + + # The compact lookup must remain numerically identical to the historical + # dense-then-unpad path. + dense_embeds = torch.nn.functional.embedding(input_ids, embedding.weight.detach()) + expected = prepare_packed_llm_inputs( + input_ids=input_ids, + text_embs=dense_embeds, + audio_embs=[], + target_ids=target_ids, + padding_id=PAD, + placeholder_id=AUDIO, + ) + torch.testing.assert_close(actual["input_embeds"], expected["input_embeds"]) + assert torch.equal(actual["target_ids"], expected["target_ids"]) + assert torch.equal(actual["llm_kwargs"]["cu_seqlens"], expected["llm_kwargs"]["cu_seqlens"]) + + actual["input_embeds"].sum().backward() + flat_real_ids = torch.cat([row[-length:] for row, length in zip(input_ids, row_lengths)]) + expected_counts = torch.bincount(flat_real_ids, minlength=vocab_size).to(embedding.weight.dtype) + expected_grad = expected_counts[:, None].expand(-1, hidden_size) + torch.testing.assert_close(embedding.weight.grad, expected_grad) + + +def test_prepare_packed_llm_inputs_accepts_native_flat_text_with_padded_parity(): + input_ids, _, target_ids, replacements = _basic_batch() + embedding = torch.nn.Embedding(128, 2) + expected = prepare_packed_llm_inputs( + input_ids=input_ids, + text_embs=None, + audio_embs=replacements, + target_ids=target_ids, + padding_id=PAD, + placeholder_id=AUDIO, + embed_tokens=embedding, + ) + flat_input_ids = torch.cat([input_ids[0], input_ids[1, 2:]]) + flat_target_ids = torch.cat([target_ids[0], target_ids[1, 2:]]) + actual = prepare_packed_llm_inputs( + input_ids=flat_input_ids, + text_embs=None, + audio_embs=replacements, + target_ids=flat_target_ids, + padding_id=PAD, + placeholder_id=AUDIO, + embed_tokens=embedding, + text_cu_seqlens=torch.tensor([0, 6, 10]), + ) + + torch.testing.assert_close(actual["input_embeds"], expected["input_embeds"]) + assert torch.equal(actual["target_ids"], expected["target_ids"]) + assert torch.equal(actual["llm_kwargs"]["cu_seqlens"], expected["llm_kwargs"]["cu_seqlens"]) + assert actual["num_examples"].item() == 2 + + def test_basic_pack_shapes_and_cu_seqlens(): input_ids, embeds, target_ids, replacements = _basic_batch() out = pack_audio_into_text_embeds( @@ -107,6 +192,64 @@ def test_basic_pack_shapes_and_cu_seqlens(): assert out["position_ids"].shape == (T_total,) +def test_fp8_token_alignment_pads_only_trailing_slots(): + input_ids = torch.tensor([[1, 2, 3, 4, 5]]) + embeds = torch.arange(10, dtype=torch.float32).reshape(1, 5, 2) + target_ids = input_ids.clone() + out = pack_audio_into_text_embeds( + input_ids=input_ids, + embeds=embeds, + target_ids=target_ids, + replacements=[], + padding_id=PAD, + placeholder_id=AUDIO, + token_alignment=8, + ) + + assert out["seq_lens"].squeeze(-1).tolist() == [5] + assert out["seq_lens_padded"].squeeze(-1).tolist() == [8] + assert out["cu_seqlens"].tolist() == [0, 8] + assert out["max_seqlen"].item() == 8 + torch.testing.assert_close(out["inputs_embeds"][:5], embeds[0]) + assert torch.count_nonzero(out["inputs_embeds"][5:]).item() == 0 + assert out["labels"].tolist() == [2, 3, 4, 5, -100, -100, -100, -100] + + +def test_fp8_token_alignment_accounts_for_context_parallel_sharding(): + input_ids = torch.tensor([[1, 2, 3, 4, 5]]) + embeds = torch.arange(10, dtype=torch.float32).reshape(1, 5, 2) + target_ids = input_ids.clone() + packed = pack_audio_into_text_embeds( + input_ids=input_ids, + embeds=embeds, + target_ids=target_ids, + replacements=[], + padding_id=PAD, + placeholder_id=AUDIO, + cp_size=2, + tp_size=3, + token_alignment=8, + ) + + assert packed["seq_lens"].squeeze(-1).tolist() == [5] + assert packed["seq_lens_padded"].squeeze(-1).tolist() == [48] + assert packed["inputs_embeds"].shape[0] == 48 + assert packed["inputs_embeds"].shape[0] % (2 * 8) == 0 + + +def test_token_alignment_must_be_positive(): + with pytest.raises(ValueError, match="token_alignment must be a positive integer"): + pack_audio_into_text_embeds( + input_ids=torch.tensor([[1]]), + embeds=torch.zeros(1, 1, 2), + target_ids=torch.tensor([[1]]), + replacements=[], + padding_id=PAD, + placeholder_id=AUDIO, + token_alignment=0, + ) + + def test_mtp_inputs_are_shifted_before_te_context_parallel_partition(monkeypatch): """Each CP rank receives globally shifted MTP inputs and targets.""" labels = torch.tensor([10, 11, 12, 13, 20, 21, 22, 23]) diff --git a/tests/collections/speechlm2/test_salm_train.py b/tests/collections/speechlm2/test_salm_train.py index 15b123d2033d..195853cb7b70 100644 --- a/tests/collections/speechlm2/test_salm_train.py +++ b/tests/collections/speechlm2/test_salm_train.py @@ -20,7 +20,6 @@ import pytest - _SALM_TRAIN_PATH = Path(__file__).parents[3] / "examples" / "speechlm2" / "salm_train.py" _SPEC = importlib.util.spec_from_file_location("salm_train_for_test", _SALM_TRAIN_PATH) _SALM_TRAIN = importlib.util.module_from_spec(_SPEC) @@ -42,6 +41,27 @@ def __init__(self, tokenizer): assert dataset.tokenizer is tokenizer +@pytest.mark.unit +def test_create_salm_dataset_does_not_forward_prompt_format_options(monkeypatch): + class LegacySALMDataset: + def __init__(self, tokenizer): + self.tokenizer = tokenizer + + tokenizer = object() + monkeypatch.setattr(_SALM_TRAIN, "SALMDataset", LegacySALMDataset) + data_cfg = { + "train_ds": { + "prompt_format": "nemotron-nano-v3", + "audio_locator_tag": "<|audio|>", + "token_equivalent_duration": 0.08, + } + } + + dataset = _SALM_TRAIN._create_salm_dataset(tokenizer, data_cfg) + + assert dataset.tokenizer is tokenizer + + @pytest.mark.unit def test_create_salm_dataset_forwards_configured_multispeaker_config(monkeypatch): multispeaker_cfg = {"num_speakers": 2} @@ -61,7 +81,37 @@ def __init__(self, tokenizer, multispeaker_cfg=None): @pytest.mark.unit -def test_train_uses_compatible_dataset_factory(monkeypatch, tmp_path): +def test_create_salm_dataset_enables_packed_audio_without_a_second_config(monkeypatch): + class PackedSALMDataset: + def __init__(self, tokenizer, pack_audio=False): + self.tokenizer = tokenizer + self.pack_audio = pack_audio + + tokenizer = object() + monkeypatch.setattr(_SALM_TRAIN, "SALMDataset", PackedSALMDataset) + + dataset = _SALM_TRAIN._create_salm_dataset(tokenizer, {}, pack_audio=True) + + assert dataset.tokenizer is tokenizer + assert dataset.pack_audio is True + + +@pytest.mark.unit +@pytest.mark.parametrize( + ("model_cfg", "expected_pack_audio", "expected_pack_sequences"), + [ + ({}, False, False), + ({"use_nemo_automodel": True, "packed_encoder_sequences": True}, True, False), + ({"use_nemo_automodel": True, "packed_sequences": True}, False, True), + ], +) +def test_train_uses_compatible_dataset_factory( + monkeypatch, + tmp_path, + model_cfg, + expected_pack_audio, + expected_pack_sequences, +): tokenizer = object() dataset = object() calls = [] @@ -84,11 +134,15 @@ class FakeDataModule: def __init__(self, data_cfg, tokenizer, dataset): pass - def create_salm_dataset(tokenizer_arg, data_cfg): - calls.append((tokenizer_arg, data_cfg)) + def create_salm_dataset(tokenizer_arg, data_cfg, *, pack_audio=False, pack_sequences=False): + calls.append((tokenizer_arg, data_cfg, pack_audio, pack_sequences)) return dataset monkeypatch.setattr(_SALM_TRAIN, "SALM", FakeSALM) + if model_cfg.get("use_nemo_automodel", False): + import nemo.collections.speechlm2 + + monkeypatch.setattr(nemo.collections.speechlm2, "SALMAutomodel", FakeSALM) monkeypatch.setattr(_SALM_TRAIN, "Trainer", FakeTrainer) monkeypatch.setattr(_SALM_TRAIN, "DataModule", FakeDataModule) monkeypatch.setattr(_SALM_TRAIN, "_create_salm_dataset", create_salm_dataset) @@ -101,10 +155,10 @@ def create_salm_dataset(tokenizer_arg, data_cfg): cfg = _SALM_TRAIN.OmegaConf.create( { "data": {"train_ds": {"seed": 0}}, - "model": {}, + "model": model_cfg, "trainer": {}, } ) _SALM_TRAIN.train.__wrapped__(cfg) - assert calls == [(tokenizer, cfg.data)] + assert calls == [(tokenizer, cfg.data, expected_pack_audio, expected_pack_sequences)] diff --git a/uv.lock b/uv.lock index 6bc53d6c3919..c35bec333214 100644 --- a/uv.lock +++ b/uv.lock @@ -1693,6 +1693,7 @@ nvvm = [ { name = "nvidia-nvvm", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-12-nemo-toolkit-compiled' and extra == 'extra-12-nemo-toolkit-compiled-a100') or (sys_platform == 'linux' and extra == 'extra-12-nemo-toolkit-compiled' and extra == 'extra-12-nemo-toolkit-compiled-a100') or (sys_platform == 'win32' and extra == 'extra-12-nemo-toolkit-compiled' and extra == 'extra-12-nemo-toolkit-compiled-a100')" }, ] + [[package]] name = "cycler" version = "0.12.1"