Training in progress - step 500

asr_config.py

@@ -21,7 +21,7 @@ class ASRConfig(transformers.PretrainedConfig):
         audio_sample_rate: int = 16000,
         projector_init_std: float = 0.02,
         projector_pool_stride: int = 4,
-        downsample_rate: int = 16,
+        downsample_rate: int = 5,  # Granite default
         projector_hidden_dim: Optional[int] = None,
         projector_type: str = "moe",  # "moe", "swiglu", "residual", "shared_moe", "mlp", "qformer"
         projector_num_layers: int = 2,  # Number of layers (for residual projector)
@@ -32,11 +32,11 @@ class ASRConfig(transformers.PretrainedConfig):
         num_experts_per_tok: int = 2,  # Top-k experts per token
         router_aux_loss_coef: float = 0.01,  # Auxiliary loss coefficient for load balancing
         use_specaugment: bool = True,  # Apply SpecAugment during training
-        # QFormer-specific configuration
-        qformer_window_size: int = 100,  # Window size for QFormer processing
+        # QFormer-specific configuration (Granite defaults)
+        qformer_window_size: int = 15,  # Window size for QFormer processing
         qformer_hidden_size: Optional[int] = None,  # QFormer hidden size (defaults to encoder_dim)
         qformer_num_layers: int = 2,  # Number of QFormer transformer layers
-        qformer_num_heads: int = 8,  # Number of attention heads in QFormer (must divide hidden size)
+        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_diversity_penalty: float = 0.0,

asr_modeling.py

@@ -316,10 +316,20 @@ class ASRModel(PreTrainedModel, GenerationMixin):
 
         return input_features
 
-    def _encode_audio(self, audio_features: torch.Tensor) -> torch.Tensor:
+    def _encode_audio(
+        self,
+        audio_features: torch.Tensor,
+        audio_attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
         """Encode audio and project to LLM embedding space.
 
-        Returns flattened audio embeddings of shape (total_audio_tokens, hidden_dim).
+        Args:
+            audio_features: Mel spectrogram features (batch, n_mels, mel_len)
+            audio_attention_mask: Mask indicating real vs padded mel frames (batch, mel_len)
+                If provided, encoder output is truncated to actual audio length.
+
+        Returns:
+            Flattened audio embeddings of shape (total_audio_tokens, hidden_dim).
         """
         # Apply SpecAugment during training (before encoding)
         audio_features = self._apply_specaugment(audio_features)
@@ -328,6 +338,14 @@ class ASRModel(PreTrainedModel, GenerationMixin):
             encoder_out = self.audio_tower(input_features=audio_features)
             hidden_states = encoder_out.last_hidden_state
 
+        # Truncate to actual audio length if attention mask provided
+        if audio_attention_mask is not None:
+            # mel_frames -> encoder_frames (stride-2 conv in Whisper encoder)
+            real_encoder_len = audio_attention_mask.sum(dim=-1) // 2
+            # For batched inputs, truncate to the max real length in the batch
+            max_real_len = real_encoder_len.max().item()
+            hidden_states = hidden_states[:, :max_real_len]
+
         audio_embeds = self.projector(hidden_states)
 
         # Flatten: (batch, seq, hidden) -> (batch * seq, hidden)
@@ -338,6 +356,7 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         self,
         input_ids: Optional[torch.Tensor] = None,
         input_features: Optional[torch.Tensor] = None,
+        audio_attention_mask: Optional[torch.Tensor] = None,
         attention_mask: Optional[torch.Tensor] = None,
         position_ids: Optional[torch.Tensor] = None,
         past_key_values: Optional[torch.Tensor] = None,
@@ -347,14 +366,19 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         cache_position: Optional[torch.Tensor] = None,
         **kwargs,
     ) -> CausalLMOutputWithPast:
-        """Forward pass for training and inference."""
+        """Forward pass for training and inference.
+
+        Args:
+            audio_attention_mask: Mask for audio features indicating real vs padded frames.
+                If provided, encoder output is truncated to actual audio length.
+        """
         # Get text embeddings if not provided
         if inputs_embeds is None:
             inputs_embeds = self.language_model.get_input_embeddings()(input_ids)
 
         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_embeds = self._encode_audio(input_features, audio_attention_mask)
 
             # Replace <audio> token placeholders with audio embeddings using masked_scatter
             audio_token_mask = (input_ids == self.audio_token_id).unsqueeze(-1)
@@ -396,14 +420,24 @@ class ASRModel(PreTrainedModel, GenerationMixin):
 
         return model_inputs
 
-    def _get_num_audio_tokens(self, input_features: torch.Tensor) -> int:
+    def _get_num_audio_tokens(
+        self,
+        input_features: torch.Tensor,
+        audio_attention_mask: Optional[torch.Tensor] = None,
+    ) -> int:
         """Calculate number of audio tokens based on input shape and projector.
 
         Whisper: input_features shape is (batch, n_mels, mel_len)
         Encoder output is mel_len // 2 due to stride-2 conv
         Projector then applies its own downsampling.
+
+        If audio_attention_mask is provided, uses actual audio length instead of padded length.
         """
-        mel_len = input_features.shape[-1]
+        if audio_attention_mask is not None:
+            # Use actual audio length (max in batch for batched inputs)
+            mel_len = audio_attention_mask.sum(dim=-1).max().item()
+        else:
+            mel_len = input_features.shape[-1]
         # Whisper encoder halves the sequence length
         encoder_output_len = mel_len // 2
         # Use projector's method to get final token count
@@ -414,6 +448,7 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         self,
         input_ids: Optional[torch.Tensor] = None,
         input_features: Optional[torch.Tensor] = None,
+        audio_attention_mask: Optional[torch.Tensor] = None,
         attention_mask: Optional[torch.Tensor] = None,
         system_prompt: Optional[str] = None,
         **generate_kwargs,
@@ -423,6 +458,10 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         Can be called in two ways:
         1. With input_ids containing <audio> tokens (from processor)
         2. With just audio, and we build the prompt internally
+
+        Args:
+            audio_attention_mask: Mask for audio features indicating real vs padded frames.
+                If provided, encoder output is truncated to actual audio length.
         """
         if input_features is None:
             raise ValueError("input_features required for generation")
@@ -431,11 +470,11 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         batch_size = input_features.shape[0]
 
         # Encode audio -> flattened embeddings
-        audio_embeds = self._encode_audio(input_features)
+        audio_embeds = self._encode_audio(input_features, audio_attention_mask)
 
         # If input_ids not provided, build prompt with correct number of audio tokens
         if input_ids is None:
-            num_audio_tokens = self._get_num_audio_tokens(input_features)
+            num_audio_tokens = self._get_num_audio_tokens(input_features, audio_attention_mask)
             audio_placeholder = "<audio>" * num_audio_tokens
 
             system_prompt = system_prompt or self.system_prompt

asr_pipeline.py

@@ -442,13 +442,18 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
 
         if isinstance(model_inputs, dict):
             input_features = model_inputs.get("input_features")
+            audio_attention_mask = model_inputs.get("attention_mask")
             if input_features is not None:
                 input_features = input_features.to(self.model.device)
+            if audio_attention_mask is not None:
+                audio_attention_mask = audio_attention_mask.to(self.model.device)
         else:
             input_features = model_inputs.to(self.model.device)
+            audio_attention_mask = None
 
         generated_ids = self.model.generate(
             input_features=input_features,
+            audio_attention_mask=audio_attention_mask,
             **generate_kwargs,
         )
 

asr_processing.py

@@ -51,13 +51,16 @@ class ASRProcessor(ProcessorMixin):
             audio_inputs = self.feature_extractor(
                 audio,
                 sampling_rate=getattr(self.feature_extractor, "sampling_rate", 16000),
+                return_attention_mask=True,
                 return_tensors=return_tensors,
                 **kwargs,
             )
             result["input_features"] = audio_inputs["input_features"]
-            # Whisper encoder output length = mel_len // 2 (stride-2 conv)
-            mel_len = audio_inputs["input_features"].shape[-1]
-            encoder_output_len = mel_len // 2
+            result["audio_attention_mask"] = audio_inputs["attention_mask"]
+
+            # Use actual audio length (from attention mask) for token count
+            real_mel_len = audio_inputs["attention_mask"].sum(dim=-1).max().item()
+            encoder_output_len = real_mel_len // 2
             num_audio_tokens = self.projector.get_output_length(encoder_output_len)
         else:
             num_audio_tokens = 0

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:8c73e72c37d7af2b5891b83b2e3ecb6da7e7a8946b6e6a3322039ff1b53ab86d
+oid sha256:d6202489b6684a47e4df573ca7d40197cc26cd5347efccc3543770f625514f84
 size 110222312

projectors.py

@@ -609,15 +609,14 @@ class QFormerAudioProjector(nn.Module):
         encoder_dim = config.encoder_dim
         llm_dim = config.llm_dim
 
-        # Window and downsampling parameters
-        self.window_size = getattr(config, "qformer_window_size", 100)
-        self.downsample_rate = getattr(config, "downsample_rate", 16)
+        # Window and downsampling parameters (Granite defaults: window=15, downsample=5)
+        self.window_size = getattr(config, "qformer_window_size", 15)
+        self.downsample_rate = getattr(config, "downsample_rate", 5)
         self.num_queries = self.window_size // self.downsample_rate
 
         # QFormer hidden size (matches encoder for cross-attention)
         qformer_hidden = getattr(config, "qformer_hidden_size", None) or encoder_dim
         qformer_num_layers = getattr(config, "qformer_num_layers", 2)
-        # Default heads must divide hidden size evenly (1280 / 16 = 80)
         qformer_num_heads = getattr(config, "qformer_num_heads", 16)
         qformer_intermediate = getattr(config, "qformer_intermediate_size", None) or (qformer_hidden * 4)
 
@@ -631,7 +630,7 @@ class QFormerAudioProjector(nn.Module):
         else:
             self.encoder_proj = None
 
-        # Configure QFormer
+        # Configure QFormer to match Granite's exact config
         qformer_config = Blip2QFormerConfig(
             hidden_size=qformer_hidden,
             num_hidden_layers=qformer_num_layers,
@@ -639,6 +638,12 @@ class QFormerAudioProjector(nn.Module):
             intermediate_size=qformer_intermediate,
             encoder_hidden_size=qformer_hidden,
             cross_attention_frequency=1,
+            # Granite-specific settings
+            hidden_act="gelu",
+            attention_probs_dropout_prob=0.1,
+            hidden_dropout_prob=0.1,
+            layer_norm_eps=1e-12,
+            initializer_range=0.02,
         )
         self.qformer = AutoModel.from_config(qformer_config)