Training in progress - step 1000

asr_config.py

@@ -6,6 +6,19 @@ import transformers
 DEFAULT_ENCODER_CONV_LAYERS = [(1, 3, 1), (1, 3, 2)]
 
 
+def compute_encoder_output_length(mel_length, conv_layers=None):
+    """Apply encoder conv layer formulas to compute output length.
+
+    Works with both Python ints and torch tensors of mel lengths; the formula
+    `(L + 2*p - (k-1) - 1) // s + 1` per layer is identical for both.
+    """
+    layers = conv_layers if conv_layers is not None else DEFAULT_ENCODER_CONV_LAYERS
+    length = mel_length
+    for padding, kernel_size, stride in layers:
+        length = (length + 2 * padding - (kernel_size - 1) - 1) // stride + 1
+    return length
+
+
 class ASRConfig(transformers.PretrainedConfig):
     """Configuration class for the ASR model.
 
@@ -14,7 +27,7 @@ class ASRConfig(transformers.PretrainedConfig):
     - Text decoder (Qwen)
     - Projector (MLP, MOSA, MoE, QFormer)
     - Generation parameters
-    - Training options (SpecAugment, LoRA)
+    - Training options (LoRA)
     """
 
     model_type = "asr_model"
@@ -38,9 +51,6 @@ class ASRConfig(transformers.PretrainedConfig):
         downsample_rate: int = 5,  # Granite default
         projector_hidden_dim: Optional[int] = None,
         projector_type: str = "mlp",  # "mlp", "mosa", "moe", "qformer"
-        projector_num_layers: int = 2,  # Number of layers in MLP projector
-        projector_init_std: float = 0.02,  # Weight initialization std
-        projector_dropout: float = 0.0,  # Dropout rate for projector layers
         # MoE-specific configuration
         num_experts: int = 4,  # Number of experts in MoE projectors
         num_experts_per_tok: int = 2,  # Top-k experts per token
@@ -51,14 +61,6 @@ class ASRConfig(transformers.PretrainedConfig):
         qformer_num_layers: int = 2,  # Number of QFormer transformer layers
         qformer_num_heads: int = 16,  # Number of attention heads in QFormer
         qformer_intermediate_size: Optional[int] = None,  # FFN size (defaults to 4x hidden)
-        label_smoothing: float = 0.0,  # Label smoothing for cross-entropy loss
-        inference_warmup_tokens: int = 10,
-        # SpecAugment settings
-        use_specaugment: bool = False,
-        num_time_masks: int = 2,
-        time_mask_length: int = 10,
-        num_freq_masks: int = 0,
-        freq_mask_length: int = 10,
         # LoRA configuration (for Stage 2 fine-tuning)
         use_lora: bool = False,
         lora_rank: int = 8,  # SALMONN default
@@ -88,22 +90,20 @@ class ASRConfig(transformers.PretrainedConfig):
             model_dtype: Model dtype ("bfloat16", "float16", "float32")
             projector_type: Projector architecture ("mlp", "mosa", "moe", "qformer")
             use_lora: Enable LoRA adapters for Stage 2 fine-tuning
-            use_specaugment: Enable SpecAugment data augmentation
         """
-        # Set default generation parameters (greedy decoding only)
+        # Set default generation parameters (greedy decoding only).
+        # Applied via setattr below — keeping these out of kwargs so they
+        # don't get re-overwritten by super().__init__(**kwargs) at the end.
         generation_defaults = {
             "num_beams": 1,
             "max_new_tokens": 128,
             "min_new_tokens": 0,
             "repetition_penalty": 1.0,
             "length_penalty": 1.0,
-            "no_repeat_ngram_size": 0,  # Prevent repeating 3-grams like "so so so"
+            "no_repeat_ngram_size": 0,
             "use_cache": True,
         }
 
-        # Apply defaults (config.json values take precedence)
-        kwargs = {**generation_defaults, **kwargs}
-
         self.audio_model_id = audio_model_id
         self.text_model_id = text_model_id
         self.attn_implementation = attn_implementation
@@ -113,13 +113,10 @@ class ASRConfig(transformers.PretrainedConfig):
         self.llm_dim = llm_dim
         self.encoder_conv_layers = encoder_conv_layers or DEFAULT_ENCODER_CONV_LAYERS
         self.audio_sample_rate = audio_sample_rate
-        self.projector_init_std = projector_init_std
         self.projector_pool_stride = projector_pool_stride
         self.downsample_rate = downsample_rate
         self.projector_hidden_dim = projector_hidden_dim
         self.projector_type = projector_type
-        self.projector_num_layers = projector_num_layers
-        self.projector_dropout = projector_dropout
         # MoE-specific configuration
         self.num_experts = num_experts
         self.num_experts_per_tok = num_experts_per_tok
@@ -130,14 +127,6 @@ class ASRConfig(transformers.PretrainedConfig):
         self.qformer_num_layers = qformer_num_layers
         self.qformer_num_heads = qformer_num_heads
         self.qformer_intermediate_size = qformer_intermediate_size
-        self.label_smoothing = label_smoothing
-        self.inference_warmup_tokens = inference_warmup_tokens
-        # SpecAugment configuration
-        self.use_specaugment = use_specaugment
-        self.num_time_masks = num_time_masks
-        self.time_mask_length = time_mask_length
-        self.num_freq_masks = num_freq_masks
-        self.freq_mask_length = freq_mask_length
         # LoRA configuration
         self.use_lora = use_lora
         self.lora_rank = lora_rank

asr_modeling.py

@@ -17,16 +17,13 @@ from transformers.generation import GenerationMixin
 from transformers.modeling_outputs import CausalLMOutputWithPast
 
 try:
-    from .asr_config import ASRConfig
+    from .asr_config import ASRConfig, compute_encoder_output_length
     from .projectors import PROJECTOR_CLASSES
 except ImportError:
-    from asr_config import ASRConfig  # type: ignore[no-redef]
+    from asr_config import ASRConfig, compute_encoder_output_length  # type: ignore[no-redef]
     from projectors import PROJECTOR_CLASSES  # type: ignore[no-redef]
 
 
-from torchaudio.transforms import SpecAugment
-
-
 def _gather_audio_embeds(audio_embeds: torch.Tensor, token_counts: torch.Tensor) -> torch.Tensor:
     """Flatten per-sample audio embeddings into a packed tensor.
 
@@ -56,7 +53,6 @@ class ASRModel(PreTrainedModel, GenerationMixin):
     _supports_flash_attn_2 = True
     supports_gradient_checkpointing = True
     _is_loading_from_pretrained: bool = False
-    _pretrained_model_path: Optional[str] = None
 
     TRANSCRIBE_PROMPT = "Transcribe the speech to text"
 
@@ -72,7 +68,6 @@ class ASRModel(PreTrainedModel, GenerationMixin):
 
         # Set flag to avoid device_map="auto" in sub-model loaders
         cls._is_loading_from_pretrained = True
-        cls._pretrained_model_path = pretrained_model_name_or_path
 
         try:
             model = cls(config, **kwargs)
@@ -134,7 +129,6 @@ class ASRModel(PreTrainedModel, GenerationMixin):
             return model
         finally:
             cls._is_loading_from_pretrained = False
-            cls._pretrained_model_path = None
 
     def __init__(self, config: ASRConfig, **kwargs) -> None:
         super().__init__(config)
@@ -190,17 +184,6 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         if getattr(config, "freeze_projector", False):
             self.projector.requires_grad_(False)
 
-        # SpecAugment for data augmentation during training
-        if getattr(config, "use_specaugment", False):
-            self.spec_augment = SpecAugment(
-                n_time_masks=config.num_time_masks,
-                time_mask_param=config.time_mask_length,
-                n_freq_masks=config.num_freq_masks,
-                freq_mask_param=config.freq_mask_length,
-            )
-        else:
-            self.spec_augment = None
-
         # For model parallelism
         self._no_split_modules = getattr(self.language_model, "_no_split_modules", [])
 
@@ -340,7 +323,13 @@ class ASRModel(PreTrainedModel, GenerationMixin):
             self.tokenizer.add_special_tokens(
                 {"additional_special_tokens": existing_special + ["<audio>"]}
             )
-            self.language_model.resize_token_embeddings(len(self.tokenizer), mean_resizing=False)
+            # mean_resizing=True initializes the new <audio> row at the mean of
+            # existing rows so its scale matches the pretrained distribution. The
+            # input-side <audio> embedding is overwritten via masked_scatter and
+            # never seen by the LM, but with tied embeddings (Qwen3-0.6B) this
+            # same row is the lm_head column for predicting <audio>; a Gaussian
+            # draw at config.initializer_range was visible in early-step logits.
+            self.language_model.resize_token_embeddings(len(self.tokenizer), mean_resizing=True)
 
         self.audio_token_id = self.tokenizer.convert_tokens_to_ids("<audio>")
         self.tokenizer.padding_side = "right"
@@ -352,9 +341,20 @@ class ASRModel(PreTrainedModel, GenerationMixin):
                 cfg.eos_token_id = self.tokenizer.eos_token_id
                 cfg.bos_token_id = self.tokenizer.bos_token_id
 
-    def _init_weights(self, _module):
-        """Weight initialization (projector weights are initialized in MoEAudioProjector)."""
-        pass
+    def train(self, mode: bool = True):
+        """Set train/eval mode, but keep frozen submodules out of train mode.
+
+        HF Trainer calls `model.train()` at the top of every training step, which
+        recursively switches every submodule into train mode — re-enabling dropout
+        on modules with `requires_grad_(False)`. The frozen encoder (and the LM
+        when `freeze_language_model=True`) should always run deterministically;
+        train-mode dropout only adds noise that can't improve a frozen network.
+        """
+        super().train(mode)
+        self.audio_tower.train(False)
+        if getattr(self.config, "freeze_language_model", True):
+            self.language_model.train(False)
+        return self
 
     def _set_gradient_checkpointing(self, enable: bool = True, gradient_checkpointing_func=None):
         """Enable/disable gradient checkpointing for the language model."""
@@ -396,34 +396,40 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         )
 
     def state_dict(self, *args, **kwargs) -> dict[str, torch.Tensor]:
-        """Save trainable weights: projector, plus the language model when fine-tuned."""
+        """Save trainable weights: projector, plus the language model when fine-tuned.
+
+        With LoRA attached, the language_model entries are flattened to plain
+        (non-PEFT) HF naming so model.safetensors round-trips through
+        ASRModel.from_pretrained — which builds a vanilla base LM, overlays
+        these weights, and only then re-attaches PEFT. lora_*/adapter weights
+        are skipped here; PEFT serializes them separately as
+        adapter_model.safetensors via the save_pretrained path below.
+        """
         sd = {f"projector.{k}": v for k, v in self.projector.state_dict().items()}
         if not getattr(self.config, "freeze_language_model", True):
-            sd.update(
-                {f"language_model.{k}": v for k, v in self.language_model.state_dict().items()}
-            )
+            lm = self.language_model
+            if hasattr(lm, "peft_config"):
+                for k, v in lm.state_dict().items():
+                    if "lora_" in k:
+                        continue
+                    if k.startswith("base_model.model."):
+                        k = k[len("base_model.model.") :]
+                    # LoRA layers wrap the original Linear as `<name>.base_layer.<weight|bias>`.
+                    k = k.replace(".base_layer.", ".")
+                    sd[f"language_model.{k}"] = v
+            else:
+                sd.update({f"language_model.{k}": v for k, v in lm.state_dict().items()})
         return sd
 
     def _compute_encoder_output_lengths(
         self,
         audio_attention_mask: torch.Tensor,
     ) -> torch.Tensor:
-        """Compute per-sample encoder output lengths using conv layer formulas.
-
-        Args:
-            audio_attention_mask: Mask indicating real vs padded mel frames (batch, mel_len)
-
-        Returns:
-            Tensor of encoder output lengths per sample (batch,)
-        """
-        # Get mel frame lengths from attention mask
-        lengths = audio_attention_mask.sum(dim=-1)
-
-        # Apply conv layer formulas: output = (input + 2*pad - (kernel-1) - 1) // stride + 1
-        for padding, kernel_size, stride in self.config.encoder_conv_layers:
-            lengths = (lengths + 2 * padding - (kernel_size - 1) - 1) // stride + 1
-
-        return lengths
+        """Compute per-sample encoder output lengths using conv layer formulas."""
+        return compute_encoder_output_length(
+            audio_attention_mask.sum(dim=-1),
+            self.config.encoder_conv_layers,
+        )
 
     def _encode_audio(
         self,
@@ -468,9 +474,6 @@ class ASRModel(PreTrainedModel, GenerationMixin):
             inputs_embeds = self.language_model.get_input_embeddings()(input_ids)
 
         if input_features is not None and input_ids is not None:
-            if self.training and self.spec_augment is not None:
-                input_features = self.spec_augment(input_features)
-
             is_audio_token = input_ids == self.audio_token_id
             if audio_token_counts is None:
                 audio_token_counts = is_audio_token.sum(dim=-1)
@@ -556,13 +559,9 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         device = input_features.device
         batch_size = input_features.shape[0]
 
-        # Encode audio -> flattened embeddings
+        # Encode audio -> flattened embeddings (no per-sample host sync)
         encoder_lengths = self._compute_encoder_output_lengths(audio_attention_mask)
-        token_counts = torch.tensor(
-            [self.projector.get_output_length(int(length.item())) for length in encoder_lengths],
-            device=input_features.device,
-            dtype=torch.long,
-        )
+        token_counts = self.projector.get_output_length(encoder_lengths).to(torch.long)
         audio_embeds = self._encode_audio(input_features, token_counts)
 
         # If input_ids not provided, build prompt with correct number of audio tokens
@@ -646,13 +645,9 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         device = input_features.device
         batch_size = input_features.shape[0]
 
-        # Encode audio -> flattened embeddings
+        # Encode audio -> flattened embeddings (no per-sample host sync)
         encoder_lengths = self._compute_encoder_output_lengths(audio_attention_mask)
-        token_counts = torch.tensor(
-            [self.projector.get_output_length(int(length.item())) for length in encoder_lengths],
-            device=input_features.device,
-            dtype=torch.long,
-        )
+        token_counts = self.projector.get_output_length(encoder_lengths).to(torch.long)
         audio_embeds = self._encode_audio(input_features, token_counts)
 
         # Build prompt with correct number of audio tokens

asr_pipeline.py

@@ -23,9 +23,9 @@ __all__ = ["ForcedAligner", "SpeakerDiarizer", "ASRPipeline"]
 
 _THINK_TAG_RE = re.compile(r"<think>.*?</think>\s*", flags=re.DOTALL)
 _DEFAULT_MIN_REPEATS = 3
-_TRAILING_CHAR_RE = re.compile(r"(.)\1{" + str(_DEFAULT_MIN_REPEATS - 1) + r",}$")
+_TRAILING_CHAR_RE = re.compile(rf"(.)\1{{{_DEFAULT_MIN_REPEATS - 1},}}$")
 _TRAILING_WORD_RE = re.compile(
-    r"\b(\w+)(?:\s+\1){" + str(_DEFAULT_MIN_REPEATS - 1) + r",}\s*$", re.IGNORECASE
+    rf"\b(\w+)(?:\s+\1){{{_DEFAULT_MIN_REPEATS - 1},}}\s*$", re.IGNORECASE
 )
 
 
@@ -291,10 +291,8 @@ def _truncate_repetitions(text: str, min_repeats: int = 3) -> str:
         char_pattern = _TRAILING_CHAR_RE
         word_pattern = _TRAILING_WORD_RE
     else:
-        char_pattern = re.compile(r"(.)\1{" + str(min_repeats - 1) + r",}$")
-        word_pattern = re.compile(
-            r"\b(\w+)(?:\s+\1){" + str(min_repeats - 1) + r",}\s*$", re.IGNORECASE
-        )
+        char_pattern = re.compile(rf"(.)\1{{{min_repeats - 1},}}$")
+        word_pattern = re.compile(rf"\b(\w+)(?:\s+\1){{{min_repeats - 1},}}\s*$", re.IGNORECASE)
 
     text = char_pattern.sub(r"\1", text)
     while word_pattern.search(text):
@@ -303,28 +301,24 @@ def _truncate_repetitions(text: str, min_repeats: int = 3) -> str:
     # 3. Truncate repeated phrases (2-20 words) at end
     # e.g., "i am sorry i am sorry i am sorry" -> "i am sorry"
     words = text.split()
-    if len(words) >= min_repeats * 2:
-        # Try phrase lengths from 2 to 20 words
-        for phrase_len in range(2, min(21, len(words) // min_repeats + 1)):
-            # Check if the last phrase_len words repeat
-            phrase = " ".join(words[-phrase_len:])
-            # Build pattern to match repeated phrases at end
-            phrase_escaped = re.escape(phrase)
-            phrase_pattern = re.compile(
-                r"(^|.*?\s)("
-                + phrase_escaped
-                + r")(?:\s+"
-                + phrase_escaped
-                + r"){"
-                + str(min_repeats - 1)
-                + r",}\s*$",
-                re.IGNORECASE,
-            )
-            match = phrase_pattern.match(text)
-            if match:
-                # Keep prefix + one instance of the phrase
-                text = (match.group(1) + match.group(2)).strip()
-                words = text.split()
-                break
+    if len(words) < min_repeats * 2:
+        return text
+
+    # Cheap pre-check: trailing window must contain duplicates for any phrase repeat
+    # to be possible. set(window) == window means all unique → no repetition.
+    window = words[-min_repeats * 2 :]
+    if len(set(window)) == len(window):
+        return text
+
+    for phrase_len in range(2, min(21, len(words) // min_repeats + 1)):
+        phrase_escaped = re.escape(" ".join(words[-phrase_len:]))
+        phrase_pattern = re.compile(
+            rf"(^|.*?\s)({phrase_escaped})(?:\s+{phrase_escaped}){{{min_repeats - 1},}}\s*$",
+            re.IGNORECASE,
+        )
+        match = phrase_pattern.match(text)
+        if match:
+            text = (match.group(1) + match.group(2)).strip()
+            break
 
     return text

asr_processing.py

@@ -5,9 +5,13 @@ import transformers
 from transformers import ProcessorMixin
 
 try:
-    from .asr_config import DEFAULT_ENCODER_CONV_LAYERS, ASRConfig
+    from .asr_config import DEFAULT_ENCODER_CONV_LAYERS, ASRConfig, compute_encoder_output_length
 except ImportError:
-    from asr_config import DEFAULT_ENCODER_CONV_LAYERS, ASRConfig  # type: ignore[no-redef]
+    from asr_config import (  # type: ignore[no-redef]
+        DEFAULT_ENCODER_CONV_LAYERS,
+        ASRConfig,
+        compute_encoder_output_length,
+    )
 
 
 class ASRProcessor(ProcessorMixin):
@@ -42,10 +46,7 @@ class ASRProcessor(ProcessorMixin):
 
     def _compute_encoder_output_length(self, mel_length: int) -> int:
         """Compute encoder output length using conv layer formulas."""
-        length = mel_length
-        for padding, kernel_size, stride in self.encoder_conv_layers:
-            length = (length + 2 * padding - (kernel_size - 1) - 1) // stride + 1
-        return length
+        return compute_encoder_output_length(mel_length, self.encoder_conv_layers)
 
     def __call__(
         self,

projectors.py

@@ -43,6 +43,11 @@ class MLPAudioProjector(nn.Module):
         self.norm = LlamaRMSNorm(hidden_dim, eps=1e-6)
         self.act = nn.GELU()
         self.linear_2 = nn.Linear(hidden_dim, llm_dim, bias=False)
+        # Output norm aligns the projector's RMS with the LM's embed_tokens
+        # distribution. Without it, linear_2's Kaiming-uniform init produces
+        # outputs ~30× quieter than embed rows, which saturates softmax at
+        # audio positions and starves them of gradient.
+        self.norm_2 = LlamaRMSNorm(llm_dim, eps=1e-6)
 
     def get_output_length(self, input_length: int) -> int:
         """Calculate output sequence length given input length (matches GLM-ASR)."""
@@ -62,7 +67,8 @@ class MLPAudioProjector(nn.Module):
         x = self.linear_1(x)
         x = self.norm(x)
         x = self.act(x)
-        return self.linear_2(x)
+        x = self.linear_2(x)
+        return self.norm_2(x)
 
 
 # =============================================================================
@@ -102,6 +108,12 @@ class MOSAProjector(nn.Module):
     Uses Conv1d for downsampling (2 layers, stride 2 each = 4x total).
     """
 
+    ADAPTER_HIDDEN_DIM = 4096
+    ROUTER_HIDDEN_DIM = 512
+    CONV_KERNEL = 3
+    CONV_STRIDE = 2
+    CONV_PADDING = 1
+
     def __init__(self, config):
         """Initialize MOSA projector.
 
@@ -112,31 +124,28 @@ class MOSAProjector(nn.Module):
         self.encoder_dim = getattr(config, "encoder_dim", None) or 1280
         self.llm_dim = getattr(config, "llm_dim", None) or 2048
         self.num_experts = getattr(config, "num_experts", None) or 4  # MOSA-Base uses 4
-        adapter_hidden = getattr(config, "adapter_hidden_dim", None) or 4096
-        router_hidden = getattr(config, "router_hidden_dim", None) or 512
 
-        # --- 1. Conv1d Downsampler (4x reduction) ---
-        # 2 layers of stride-2 convolution
+        conv_kwargs = {
+            "kernel_size": self.CONV_KERNEL,
+            "stride": self.CONV_STRIDE,
+            "padding": self.CONV_PADDING,
+        }
         self.downsampler = nn.Sequential(
-            nn.Conv1d(self.encoder_dim, self.encoder_dim, kernel_size=3, stride=2, padding=1),
+            nn.Conv1d(self.encoder_dim, self.encoder_dim, **conv_kwargs),
             nn.GELU(),
-            nn.Conv1d(self.encoder_dim, self.llm_dim, kernel_size=3, stride=2, padding=1),
+            nn.Conv1d(self.encoder_dim, self.llm_dim, **conv_kwargs),
             nn.GELU(),
         )
 
-        # --- 2. Simple Router (MOSA-Base: 2 layers with ReLU) ---
-        # Takes downsampled features (llm_dim) -> 512 -> num_experts
         self.router = nn.Sequential(
-            nn.Linear(self.llm_dim, router_hidden),
+            nn.Linear(self.llm_dim, self.ROUTER_HIDDEN_DIM),
             nn.ReLU(),
-            nn.Linear(router_hidden, self.num_experts),
+            nn.Linear(self.ROUTER_HIDDEN_DIM, self.num_experts),
         )
 
-        # --- 3. Experts (Simple 2-layer GELU adapters) ---
-        # Each expert: llm_dim -> hidden -> llm_dim (much smaller than frame-stacking)
         self.experts = nn.ModuleList(
             [
-                SimpleAdapter(self.llm_dim, adapter_hidden, self.llm_dim)
+                SimpleAdapter(self.llm_dim, self.ADAPTER_HIDDEN_DIM, self.llm_dim)
                 for _ in range(self.num_experts)
             ]
         )
@@ -150,26 +159,22 @@ class MOSAProjector(nn.Module):
         Returns:
             Projected features of shape [batch, out_len, llm_dim]
         """
-        # --- 1. Conv1d Downsampling ---
-        # Permute for Conv1d: [B, S, D] -> [B, D, S]
-        x = x.transpose(1, 2)
-        x = self.downsampler(x)
-        # Permute back: [B, D, S] -> [B, S, D]
-        x = x.transpose(1, 2)
-
-        # --- 2. Routing ---
+        x = self.downsampler(x.transpose(1, 2)).transpose(1, 2)
+
         routing_weights = F.softmax(self.router(x), dim=-1)  # (B, out_len, num_experts)
 
-        # --- 3. Expert Mixture (Dense Execution) ---
-        expert_outputs = torch.stack([expert(x) for expert in self.experts])  # (E, B, out_len, D)
-        return torch.einsum("ebsd, bse -> bsd", expert_outputs, routing_weights)
+        # Accumulate weighted expert outputs without materializing all experts at once.
+        output = self.experts[0](x) * routing_weights[..., 0:1]
+        for i, expert in enumerate(self.experts[1:], start=1):
+            output = output + expert(x) * routing_weights[..., i : i + 1]
+        return output
 
     def get_output_length(self, input_length: int) -> int:
         """Calculate output sequence length after Conv1d downsampling (4x reduction)."""
-        # Conv1d with stride 2, kernel 3, padding 1: out = (in + 2*1 - 3) // 2 + 1 = (in - 1) // 2 + 1
-        # Applied twice for 4x total reduction
-        after_conv1 = (input_length + 2 * 1 - 3) // 2 + 1
-        return (after_conv1 + 2 * 1 - 3) // 2 + 1
+        length = input_length
+        for _ in range(2):
+            length = (length + 2 * self.CONV_PADDING - self.CONV_KERNEL) // self.CONV_STRIDE + 1
+        return length
 
 
 # =============================================================================
@@ -414,10 +419,13 @@ class QFormerAudioProjector(nn.Module):
         # Final projection to LLM dimension (Granite uses bias=True)
         self.linear = nn.Linear(qformer_hidden, llm_dim)
 
-    def get_output_length(self, input_length: int) -> int:
-        """Calculate output sequence length given input length."""
-        # QFormer uses window-based processing with num_queries per window
-        nblocks = math.ceil(input_length / self.window_size)
+    def get_output_length(self, input_length):
+        """Calculate output sequence length given input length.
+
+        Accepts either Python ints or torch tensors; uses ceiling division so
+        the formula is identical for both — math.ceil would block tensors.
+        """
+        nblocks = (input_length + self.window_size - 1) // self.window_size
         return nblocks * self.num_queries
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: