Training in progress - step 500

asr_config.py

@@ -54,7 +54,7 @@ class ASRConfig(transformers.PretrainedConfig):
         lora_rank: int = 8,  # SALMONN default
         lora_alpha: int = 32,  # SALMONN default (scaling factor 4.0)
         lora_dropout: float = 0.0,
-        lora_target_modules: Optional[list] = None,  # Default: ["q_proj", "v_proj"]
+        lora_target_modules: Optional[list] = None,  # Default: all linear layers
         freeze_projector: bool = False,  # True for Stage 2 (LoRA-only training)
         max_new_tokens: Optional[int] = None,
         min_new_tokens: Optional[int] = None,
@@ -121,7 +121,9 @@ class ASRConfig(transformers.PretrainedConfig):
         self.lora_rank = lora_rank
         self.lora_alpha = lora_alpha
         self.lora_dropout = lora_dropout
-        self.lora_target_modules = lora_target_modules or ["q_proj", "v_proj"]
+        self.lora_target_modules = lora_target_modules or [
+            "q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"
+        ]
         self.freeze_projector = freeze_projector
 
         # Generation parameters (use explicit value if provided, else use default)

asr_modeling.py

@@ -190,27 +190,30 @@ class ASRModel(PreTrainedModel, GenerationMixin):
                 state_dict = load_file(model_file)
                 model.load_state_dict(state_dict, strict=False)
 
-            # Load LoRA adapters if they exist and use_lora is enabled
+            # Load LoRA adapters if use_lora is enabled
             if getattr(config, "use_lora", False):
-                adapter_file = cached_file(
+                # Check for adapter_config.json (required by PEFT to load adapters)
+                adapter_config_file = cached_file(
                     pretrained_model_name_or_path,
-                    "adapter_model.safetensors",
+                    "adapter_config.json",
                     _raise_exceptions_for_missing_entries=False,
                     **cache_kwargs,
                 )
-                if adapter_file is not None:
+                if adapter_config_file is not None:
+                    # Load saved adapter weights using the original repo_id/path
+                    # PEFT handles Hub downloads and caching internally
                     from peft import PeftModel
 
-                    # Get the directory containing the adapter
-                    import os
-
-                    adapter_dir = os.path.dirname(adapter_file)
-                    # Load adapter weights into the already-wrapped PeftModel
+                    # language_model is bare (not PEFT-wrapped) since we skipped _setup_lora
                     model.language_model = PeftModel.from_pretrained(
-                        model.language_model.base_model,
-                        adapter_dir,
+                        model.language_model,
+                        pretrained_model_name_or_path,  # Use original repo_id, not cache path
                         is_trainable=True,
+                        **cache_kwargs,
                     )
+                else:
+                    # No saved adapters - initialize fresh LoRA for training
+                    model._setup_lora(config)
 
             return model
         finally:
@@ -259,7 +262,10 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         self.projector = self._create_projector(config, target_dtype)
 
         # Setup LoRA if enabled (Stage 2 fine-tuning)
-        if getattr(config, "use_lora", False):
+        # Skip if loading from pretrained - from_pretrained will handle adapter loading
+        if getattr(config, "use_lora", False) and not getattr(
+            self.__class__, "_is_loading_from_pretrained", False
+        ):
             self._setup_lora(config)
 
         # Freeze projector if specified (for Stage 2 LoRA-only training)
@@ -821,8 +827,10 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         self.feature_extractor.save_pretrained(save_dir)
 
         # Save LoRA adapters if present (creates adapter_model.safetensors and adapter_config.json)
+        # Don't save embedding layers - the <audio> token embedding is never used
+        # (it's replaced with projected audio embeddings before the LLM sees it)
         if hasattr(self.language_model, "peft_config"):
-            self.language_model.save_pretrained(save_dir)
+            self.language_model.save_pretrained(save_dir, save_embedding_layers=False)
 
         # Add processor auto_map to preprocessor_config.json
         config_path = save_dir / "preprocessor_config.json"

asr_pipeline.py

@@ -504,7 +504,9 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
             if repeat_count >= 1:
                 words = words[: idx + n]
                 text = " ".join(words)
-                print(f"[DEBUG] Truncated repetition: {original_len} -> {len(words)} words (n={n}, repeats={repeat_count})")
+                print(
+                    f"[DEBUG] Truncated repetition: {original_len} -> {len(words)} words (n={n}, repeats={repeat_count})"
+                )
                 break
 
         # 3. COMBINE ACRONYMS

projectors.py

@@ -135,7 +135,9 @@ class MOSAProjector(nn.Module):
 
         # --- 1. Router Branch ---
         # Mean pool encoder outputs for routing decisions
-        x_pooled = x.reshape(batch_size, out_len, self.k, self.encoder_dim).mean(dim=2)  # (B, out_len, D)
+        x_pooled = x.reshape(batch_size, out_len, self.k, self.encoder_dim).mean(
+            dim=2
+        )  # (B, out_len, D)
 
         # Router logits and softmax gating (dense MoE)
         routing_weights = F.softmax(self.router(x_pooled), dim=-1)  # (B, out_len, num_experts)