Training in progress - step 500

asr_modeling.py

@@ -269,11 +269,7 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         """Only save trainable projector weights."""
         return {f"projector.{k}": v for k, v in self.projector.state_dict().items()}
 
-    def _apply_specaugment(
-        self,
-        input_features: torch.Tensor,
-        attention_mask: Optional[torch.Tensor] = None,
-    ) -> torch.Tensor:
+    def _apply_specaugment(self, input_features: torch.Tensor) -> torch.Tensor:
         if not getattr(self.config, "use_specaugment", False):
             return input_features
 
@@ -294,7 +290,6 @@ class ASRModel(PreTrainedModel, GenerationMixin):
                 (batch_size, sequence_length),
                 mask_prob=mask_time_prob,
                 mask_length=mask_time_length,
-                attention_mask=attention_mask,
                 min_masks=2,
             )
             mask_time_indices = torch.tensor(
@@ -321,22 +316,16 @@ class ASRModel(PreTrainedModel, GenerationMixin):
 
         return input_features
 
-    def _encode_audio(
-        self,
-        audio_features: torch.Tensor,
-        audio_attention_mask: Optional[torch.Tensor] = None,
-    ) -> torch.Tensor:
+    def _encode_audio(self, audio_features: torch.Tensor) -> torch.Tensor:
         """Encode audio and project to LLM embedding space.
 
         Returns flattened audio embeddings of shape (total_audio_tokens, hidden_dim).
         """
         # Apply SpecAugment during training (before encoding)
-        audio_features = self._apply_specaugment(audio_features, audio_attention_mask)
+        audio_features = self._apply_specaugment(audio_features)
 
         with torch.no_grad():
-            encoder_out = self.audio_tower(
-                input_features=audio_features, attention_mask=audio_attention_mask
-            )
+            encoder_out = self.audio_tower(input_features=audio_features)
             hidden_states = encoder_out.last_hidden_state
 
         audio_embeds = self.projector(hidden_states)
@@ -356,7 +345,6 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         labels: Optional[torch.Tensor] = None,
         use_cache: Optional[bool] = None,
         cache_position: Optional[torch.Tensor] = None,
-        audio_attention_mask: Optional[torch.Tensor] = None,
         **kwargs,
     ) -> CausalLMOutputWithPast:
         """Forward pass for training and inference."""
@@ -366,7 +354,7 @@ class ASRModel(PreTrainedModel, GenerationMixin):
 
         if input_features is not None and input_ids is not None:
             # Encode audio -> flattened (total_audio_tokens, hidden_dim)
-            audio_embeds = self._encode_audio(input_features, audio_attention_mask)
+            audio_embeds = self._encode_audio(input_features)
 
             # Replace <audio> token placeholders with audio embeddings using masked_scatter
             audio_token_mask = (input_ids == self.audio_token_id).unsqueeze(-1)
@@ -427,7 +415,6 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         input_ids: Optional[torch.Tensor] = None,
         input_features: Optional[torch.Tensor] = None,
         attention_mask: Optional[torch.Tensor] = None,
-        audio_attention_mask: Optional[torch.Tensor] = None,
         system_prompt: Optional[str] = None,
         **generate_kwargs,
     ) -> torch.Tensor:
@@ -444,7 +431,7 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         batch_size = input_features.shape[0]
 
         # Encode audio -> flattened embeddings
-        audio_embeds = self._encode_audio(input_features, audio_attention_mask)
+        audio_embeds = self._encode_audio(input_features)
 
         # If input_ids not provided, build prompt with correct number of audio tokens
         if input_ids is None:

projectors.py

@@ -680,14 +680,13 @@ class QFormerAudioProjector(nn.Module):
         effective_batch = batch_size * nblocks
         hidden_states = hidden_states.view(effective_batch, self.window_size, -1)
 
-        # Expand queries to match batch size (Granite relies on broadcast, but CUDA has issues)
+        # Expand queries to match batch size
         query_embeds = self.query.expand(effective_batch, -1, -1)
 
         # QFormer cross-attention
         query_output = self.qformer(
             query_embeds=query_embeds,
             encoder_hidden_states=hidden_states,
-            encoder_attention_mask=None,
             return_dict=True,
         )