Training in progress - step 500

asr_config.py

@@ -22,10 +22,15 @@ class ASRConfig(transformers.PretrainedConfig):
         projector_init_std: float = 0.02,
         projector_pool_stride: int = 2,
         projector_hidden_dim: Optional[int] = None,
-        projector_type: str = "moe",  # "moe", "swiglu", or "residual"
+        projector_type: str = "moe",  # "moe", "swiglu", "residual", "shared_moe"
         projector_num_layers: int = 2,  # Number of layers (for residual projector)
         projector_dropout: float = 0.05,  # Dropout rate for projector layers
         projector_input_noise: float = 0.02,  # Input noise for projector
+        # MoE-specific configuration
+        num_experts: int = 4,  # Number of experts in MoE projectors
+        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
         label_smoothing: float = 0.0,  # Label smoothing for cross-entropy loss
         inference_diversity_penalty: float = 0.0,
         inference_warmup_tokens: int = 10,
@@ -45,9 +50,10 @@ class ASRConfig(transformers.PretrainedConfig):
         # Set default generation parameters
         generation_defaults = {
             "num_beams": 1,
-            "max_new_tokens": 64,
-            "min_new_tokens": 1,
+            "max_new_tokens": 256,
+            "min_new_tokens": 0,
             "do_sample": False,
+            "temperature": 0.1,
             "repetition_penalty": 1.0,
             "length_penalty": 1.0,
             "no_repeat_ngram_size": 0,
@@ -73,9 +79,29 @@ class ASRConfig(transformers.PretrainedConfig):
         self.projector_num_layers = projector_num_layers
         self.projector_dropout = projector_dropout
         self.projector_input_noise = projector_input_noise
+        # MoE-specific configuration
+        self.num_experts = num_experts
+        self.num_experts_per_tok = num_experts_per_tok
+        self.router_aux_loss_coef = router_aux_loss_coef
+        self.use_specaugment = use_specaugment
         self.label_smoothing = label_smoothing
         self.inference_diversity_penalty = inference_diversity_penalty
         self.inference_warmup_tokens = inference_warmup_tokens
+
+        # Generation parameters (use explicit value if provided, else use default)
+        self.num_beams = num_beams if num_beams is not None else generation_defaults["num_beams"]
+        self.max_new_tokens = max_new_tokens if max_new_tokens is not None else generation_defaults["max_new_tokens"]
+        self.min_new_tokens = min_new_tokens if min_new_tokens is not None else generation_defaults["min_new_tokens"]
+        self.do_sample = do_sample if do_sample is not None else generation_defaults["do_sample"]
+        self.repetition_penalty = repetition_penalty if repetition_penalty is not None else generation_defaults["repetition_penalty"]
+        self.length_penalty = length_penalty if length_penalty is not None else generation_defaults["length_penalty"]
+        self.no_repeat_ngram_size = no_repeat_ngram_size if no_repeat_ngram_size is not None else generation_defaults["no_repeat_ngram_size"]
+        self.use_cache = use_cache if use_cache is not None else generation_defaults["use_cache"]
+        self.temperature = temperature if temperature is not None else generation_defaults["temperature"]
+        self.top_k = top_k
+        self.top_p = top_p
+        self.early_stopping = early_stopping
+
         if "audio_config" not in kwargs:
             self.audio_config = transformers.AutoConfig.from_pretrained(audio_model_id)
             # Override dtype to match model_dtype

asr_modeling.py

@@ -16,10 +16,26 @@ from transformers.models.whisper.modeling_whisper import (
     _compute_mask_indices,
 )
 
-from .asr_config import ASRConfig
-from .moe_projector import MoEAudioProjector
-from .residual_projector import ResidualAudioProjector
-from .swiglu_projector import AudioProjector
+try:
+    from .asr_config import ASRConfig
+    from .moe_projector import MoEAudioProjector
+    from .residual_projector import ResidualAudioProjector
+    from .swiglu_projector import AudioProjector
+    from .shared_moe_projector import SharedMoEAudioProjector
+except ImportError:
+    from asr_config import ASRConfig  # type: ignore[no-redef]
+    from moe_projector import MoEAudioProjector  # type: ignore[no-redef]
+    from residual_projector import ResidualAudioProjector  # type: ignore[no-redef]
+    from swiglu_projector import AudioProjector  # type: ignore[no-redef]
+    from shared_moe_projector import SharedMoEAudioProjector  # type: ignore[no-redef]
+
+# Map projector type names to classes
+PROJECTOR_CLASSES = {
+    "swiglu": AudioProjector,
+    "residual": ResidualAudioProjector,
+    "moe": MoEAudioProjector,
+    "shared_moe": SharedMoEAudioProjector,
+}
 
 
 class ASRModel(PreTrainedModel):
@@ -173,12 +189,13 @@ class ASRModel(PreTrainedModel):
 
         # Select projector type based on config
         projector_type = getattr(config, "projector_type", "moe")
-        if projector_type == "swiglu":
-            projector = AudioProjector(config)
-        elif projector_type == "residual":
-            projector = ResidualAudioProjector(config)
-        else:  # default to "moe"
-            projector = MoEAudioProjector(config)
+        projector_class = PROJECTOR_CLASSES.get(projector_type)
+        if projector_class is None:
+            raise ValueError(
+                f"Unknown projector_type: {projector_type}. "
+                f"Valid options: {list(PROJECTOR_CLASSES.keys())}"
+            )
+        projector = projector_class(config)
 
         # Move projector to same device as language model (important when using quantization)
         device = next(self.language_model.parameters()).device
@@ -234,7 +251,10 @@ class ASRModel(PreTrainedModel):
 
     def get_processor(self):
         """Get the processor for this model."""
-        from .asr_processing import ASRProcessor
+        try:
+            from .asr_processing import ASRProcessor
+        except ImportError:
+            from asr_processing import ASRProcessor  # type: ignore[no-redef]
 
         return ASRProcessor(feature_extractor=self.feature_extractor, tokenizer=self.tokenizer)
 
@@ -247,15 +267,7 @@ class ASRModel(PreTrainedModel):
         input_features: torch.Tensor,
         attention_mask: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
-        """
-        Apply SpecAugment masking to input features.
 
-        Uses Whisper's default parameters:
-        - mask_time_prob: 0.05 (5% of time steps)
-        - mask_time_length: 10 frames
-        - mask_feature_prob: 0.0 (disabled by default)
-        - mask_feature_length: 10 features
-        """
         if not getattr(self.config, "use_specaugment", False):
             return input_features
 
@@ -459,6 +471,12 @@ class ASRModel(PreTrainedModel):
                 label_smoothing=getattr(self.config, "label_smoothing", 0.0),
             )
 
+            # Add auxiliary loss from MoE projectors if available
+            if hasattr(self.projector, "get_aux_loss"):
+                aux_loss = self.projector.get_aux_loss()
+                if aux_loss is not None and aux_loss.numel() > 0:
+                    loss = loss + aux_loss.to(loss.device)
+
         return CausalLMOutputWithPast(
             loss=loss,
             logits=outputs.logits,
@@ -521,18 +539,38 @@ class ASRModel(PreTrainedModel):
             prompt_ids, audio_embeds, audio_mask=audio_mask
         )
 
-        # Set generation defaults
-        generate_kwargs.setdefault("max_new_tokens", getattr(self.config, "max_new_tokens", 128))
-        generate_kwargs.setdefault("use_cache", True)
+        # Set generation defaults from config
+        generate_kwargs.setdefault("max_new_tokens", self.config.max_new_tokens)
+        generate_kwargs.setdefault("num_beams", self.config.num_beams)
+        generate_kwargs.setdefault("do_sample", self.config.do_sample)
+        generate_kwargs.setdefault("use_cache", self.config.use_cache)
+        generate_kwargs.setdefault("length_penalty", self.config.length_penalty)
+        generate_kwargs.setdefault("repetition_penalty", self.config.repetition_penalty)
+        generate_kwargs.setdefault("no_repeat_ngram_size", self.config.no_repeat_ngram_size)
+        generate_kwargs.setdefault("temperature", self.config.temperature)
+        if self.config.top_k is not None:
+            generate_kwargs.setdefault("top_k", self.config.top_k)
+        if self.config.top_p is not None:
+            generate_kwargs.setdefault("top_p", self.config.top_p)
         generate_kwargs.setdefault(
             "eos_token_id", self.tokenizer.convert_tokens_to_ids("<|im_end|>")
         )
         generate_kwargs.setdefault("pad_token_id", self.tokenizer.pad_token_id)
 
+        # Create dummy input_ids matching inputs_embeds length for repetition penalty tracking
+        # Use pad_token_id as placeholder since the actual tokens don't matter for penalty calc
+        dummy_input_ids = torch.full(
+            (inputs_embeds.shape[0], inputs_embeds.shape[1]),
+            self.tokenizer.pad_token_id,
+            dtype=torch.long,
+            device=device,
+        )
+
         # Generate (type ignore needed as generate() has complex return type)
         # Note: When using inputs_embeds, generate() returns only new tokens
         # (no placeholder positions for input embeddings), so no stripping needed
         output = self.language_model.generate(  # type: ignore[operator]
+            input_ids=dummy_input_ids,
             inputs_embeds=inputs_embeds,
             attention_mask=attention_mask,
             **generate_kwargs,
@@ -576,6 +614,17 @@ class ASRModel(PreTrainedModel):
         src_dir = PathlibPath(__file__).parent
         for asr_file in src_dir.glob("asr_*.py"):
             shutil.copy(asr_file, save_dir / asr_file.name)
+        # Copy projector files
+        projector_files = [
+            "moe_projector.py",
+            "residual_projector.py",
+            "swiglu_projector.py",
+            "shared_moe_projector.py",
+        ]
+        for projector_file in projector_files:
+            src_path = src_dir / projector_file
+            if src_path.exists():
+                shutil.copy(src_path, save_dir / projector_file)
 
 
 # Register with transformers Auto classes

asr_pipeline.py

@@ -4,9 +4,11 @@ from typing import Any
 
 import torch
 import transformers
-from truecase import get_true_case
 
-from .asr_modeling import ASRModel
+try:
+    from .asr_modeling import ASRModel
+except ImportError:
+    from asr_modeling import ASRModel  # type: ignore[no-redef]
 
 
 class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
@@ -27,11 +29,6 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
             model=model, feature_extractor=feature_extractor, tokenizer=tokenizer, **kwargs
         )
 
-        # Initialize text normalizer (WhisperTokenizer has the normalize method we need)
-        from transformers import WhisperTokenizer
-
-        self.text_normalizer = WhisperTokenizer.from_pretrained("openai/whisper-tiny")
-
     def __call__(self, inputs, **kwargs):
         generate_kwargs = {}
         generate_keys = [
@@ -89,8 +86,6 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
                 all_tokens.extend(tokens.tolist() if torch.is_tensor(tokens) else tokens)
 
         text = self.tokenizer.decode(all_tokens, skip_special_tokens=True).strip()
-        text = self.text_normalizer.normalize(text)
-        text = get_true_case(text)
 
         return {"text": text}
 
@@ -105,7 +100,13 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
             if "bytes" in inputs:
                 inputs = self._decode_audio_bytes(inputs["bytes"])
             elif "array" in inputs:
-                inputs = {"raw": inputs["array"], "sampling_rate": inputs["sampling_rate"]}
+                inputs = {
+                    "raw": inputs["array"],
+                    "sampling_rate": inputs.get("sampling_rate", self.feature_extractor.sampling_rate),
+                }
+            elif "path" in inputs and "array" not in inputs:
+                # Lazy-loaded audio - load from path
+                inputs = self._decode_audio_bytes(Path(inputs["path"]).read_bytes())
         elif hasattr(inputs, "array") and hasattr(inputs, "sampling_rate"):
             inputs = {"raw": inputs.array, "sampling_rate": inputs.sampling_rate}
         elif hasattr(inputs, "__array__") and not isinstance(inputs, (dict, bytes, str)):
@@ -116,6 +117,23 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
                 "sampling_rate": self.model.config.audio_sample_rate,
             }
 
+        # Resample to target sample rate if needed (workaround for transformers bug)
+        # See: https://github.com/huggingface/transformers/pull/41298
+        if isinstance(inputs, dict) and "sampling_rate" in inputs:
+            in_sr = inputs["sampling_rate"]
+            target_sr = self.feature_extractor.sampling_rate
+            if in_sr != target_sr:
+                import librosa
+                import numpy as np
+
+                audio = inputs["raw"]
+                if hasattr(audio, "numpy"):
+                    audio = audio.numpy()
+                resampled = librosa.resample(
+                    np.asarray(audio, dtype=np.float32), orig_sr=in_sr, target_sr=target_sr
+                )
+                inputs = {"raw": resampled, "sampling_rate": target_sr}
+
         return super().preprocess(inputs, **preprocess_params)
 
     def _decode_audio_bytes(self, wav_bytes: bytes) -> dict[str, Any]:
@@ -225,7 +243,5 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
             tokens = tokens[0]
 
         text = self.tokenizer.decode(tokens, skip_special_tokens=True).strip()
-        text = self.text_normalizer.normalize(text)
-        text = get_true_case(text)
 
         return {"text": text}

asr_processing.py

@@ -4,7 +4,10 @@ from pathlib import Path
 import transformers
 from transformers import AutoTokenizer, ProcessorMixin
 
-from .asr_config import ASRConfig
+try:
+    from .asr_config import ASRConfig
+except ImportError:
+    from asr_config import ASRConfig  # type: ignore[no-redef]
 
 
 class ASRProcessor(ProcessorMixin):

moe_projector.py

@@ -11,14 +11,19 @@ class SimpleAdapter(nn.Module):
     projecting the hidden dimension from 3072 to 4096 and back to 3072."
     """
 
-    def __init__(self, in_features, hidden_features, out_features):
+    def __init__(self, in_features, hidden_features, out_features, dropout=0.0):
         super().__init__()
         self.fc1 = nn.Linear(in_features, hidden_features)
         self.relu = nn.ReLU()
+        self.dropout = nn.Dropout(dropout)
         self.fc2 = nn.Linear(hidden_features, out_features)
 
     def forward(self, x):
-        return self.fc2(self.relu(self.fc1(x)))
+        x = self.fc1(x)
+        x = self.relu(x)
+        x = self.dropout(x)
+        x = self.fc2(x)
+        return x
 
 
 class MoEAudioProjector(nn.Module):
@@ -47,6 +52,9 @@ class MoEAudioProjector(nn.Module):
         # Adapter hidden dim: paper uses 4096
         adapter_hidden = getattr(config, "projector_hidden_dim", None) or 4096
 
+        # Dropout rate for experts (not applied to router)
+        self.dropout_rate = getattr(config, "projector_dropout", 0.1)
+
         # --- Convolutional Subsampling (Section III-B) ---
         # "two convolutional layers, each with a kernel size of 3 and a stride of 2"
         # Maps encoder_dim (1280) -> llm_dim (3072), total stride=4
@@ -70,7 +78,7 @@ class MoEAudioProjector(nn.Module):
         # "projecting the hidden dimension from 3072 to 4096 and back to 3072"
         self.experts = nn.ModuleList(
             [
-                SimpleAdapter(self.llm_dim, adapter_hidden, self.llm_dim)
+                SimpleAdapter(self.llm_dim, adapter_hidden, self.llm_dim, dropout=self.dropout_rate)
                 for _ in range(self.num_experts)
             ]
         )
@@ -149,3 +157,7 @@ class MoEAudioProjector(nn.Module):
             final_out.add_(expert_out * expert_weight)
 
         return self.ln_post(final_out)
+
+    def get_aux_loss(self) -> torch.Tensor:
+        """Return auxiliary loss (none for dense MoE - all experts always used)."""
+        return torch.tensor(0.0)

shared_moe_projector.py

@@ -0,0 +1,334 @@
+"""Shared MoE Audio Projector.
+
+A simplified MoE projector combining the best ideas:
+- Shared expert: Always-on baseline processing (from GLM4)
+- Zero-initialized router: Learns specialization naturally (from Qwen3)
+- Simple top-k softmax: No grouping complexity (from Mixtral)
+- Renormalized weights: Top-k weights sum to 1
+
+Architecture:
+    Output = SharedExpert(x) + TopKRoutedExperts(x)
+
+The shared expert ensures every audio token gets consistent baseline
+processing, while routed experts can specialize for different patterns
+(e.g., vowels vs consonants, silence vs speech).
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F  # noqa: N812
+
+
+class SharedExpert(nn.Module):
+    """Shared expert MLP that processes all tokens."""
+
+    def __init__(self, input_dim: int, hidden_dim: int, output_dim: int):
+        super().__init__()
+        self.gate_proj = nn.Linear(input_dim, hidden_dim, bias=False)
+        self.up_proj = nn.Linear(input_dim, hidden_dim, bias=False)
+        self.down_proj = nn.Linear(hidden_dim, output_dim, bias=False)
+        self.act = nn.SiLU()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.down_proj(self.act(self.gate_proj(x)) * self.up_proj(x))
+
+
+class SwiGLUExpert(nn.Module):
+    """Single SwiGLU expert MLP."""
+
+    def __init__(self, input_dim: int, hidden_dim: int, output_dim: int):
+        super().__init__()
+        self.gate_proj = nn.Linear(input_dim, hidden_dim, bias=False)
+        self.up_proj = nn.Linear(input_dim, hidden_dim, bias=False)
+        self.down_proj = nn.Linear(hidden_dim, output_dim, bias=False)
+        self.act = nn.SiLU()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.down_proj(self.act(self.gate_proj(x)) * self.up_proj(x))
+
+
+class RoutedExperts(nn.Module):
+    """
+    Sparse routed experts using token dispatch.
+
+    For each expert, gathers assigned tokens, processes them, then scatters back.
+    Memory-efficient: O(num_tokens * hidden_dim) instead of
+    O(num_tokens * num_experts * hidden_dim * input_dim).
+    """
+
+    def __init__(
+        self, num_experts: int, top_k: int, input_dim: int, hidden_dim: int, output_dim: int
+    ):
+        super().__init__()
+        self.num_experts = num_experts
+        self.top_k = top_k
+        self.output_dim = output_dim
+
+        # ModuleList of expert MLPs
+        self.experts = nn.ModuleList([
+            SwiGLUExpert(input_dim, hidden_dim, output_dim)
+            for _ in range(num_experts)
+        ])
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        top_k_indices: torch.Tensor,
+        top_k_weights: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Token dispatch approach - memory efficient.
+
+        Args:
+            hidden_states: [num_tokens, input_dim]
+            top_k_indices: [num_tokens, top_k]
+            top_k_weights: [num_tokens, top_k]
+
+        Returns:
+            output: [num_tokens, output_dim]
+        """
+        num_tokens = hidden_states.shape[0]
+        device = hidden_states.device
+        dtype = hidden_states.dtype
+
+        # Output accumulator
+        output = torch.zeros(num_tokens, self.output_dim, device=device, dtype=dtype)
+
+        # Process each expert
+        for expert_idx, expert in enumerate(self.experts):
+            # Find which (token, slot) pairs use this expert
+            # top_k_indices: [N, K], we want all positions where value == expert_idx
+            expert_mask = top_k_indices == expert_idx  # [N, K]
+
+            if not expert_mask.any():
+                continue
+
+            # Get token indices and slot indices where this expert is selected
+            token_indices, slot_indices = torch.where(expert_mask)
+
+            # Gather the tokens for this expert
+            expert_input = hidden_states[token_indices]  # [num_selected, input_dim]
+
+            # Process through expert
+            expert_output = expert(expert_input)  # [num_selected, output_dim]
+
+            # Get weights for these tokens at these slots
+            weights = top_k_weights[token_indices, slot_indices]  # [num_selected]
+
+            # Weighted output
+            weighted_output = expert_output * weights.unsqueeze(-1)
+
+            # Scatter-add back to output
+            output.index_add_(0, token_indices, weighted_output)
+
+        return output
+
+
+class SharedMoEBlock(nn.Module):
+    """MoE block with shared expert + sparse routed experts."""
+
+    def __init__(
+        self,
+        input_dim: int,
+        hidden_dim: int,
+        output_dim: int,
+        num_experts: int = 4,
+        top_k: int = 2,
+    ):
+        super().__init__()
+        self.num_experts = num_experts
+        self.top_k = top_k
+
+        # Router: zero-initialized for natural learning
+        self.router = nn.Linear(input_dim, num_experts, bias=False)
+        nn.init.zeros_(self.router.weight)
+
+        # Shared expert (always active)
+        self.shared_expert = SharedExpert(input_dim, hidden_dim, output_dim)
+
+        # Routed experts (sparse)
+        self.routed_experts = RoutedExperts(
+            num_experts, self.top_k, input_dim, hidden_dim, output_dim
+        )
+
+
+        # For auxiliary loss
+        self.last_router_logits = None
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        batch_size, seq_len, dim = hidden_states.shape
+
+        # Shared expert output (all tokens)
+        shared_out = self.shared_expert(hidden_states)
+
+        # Routing
+        flat_hidden = hidden_states.view(-1, dim)
+        router_logits = self.router(flat_hidden)
+        self.last_router_logits = router_logits
+
+        # Softmax -> top-k -> renormalize
+        router_probs = F.softmax(router_logits.float(), dim=-1)
+        top_k_weights, top_k_indices = torch.topk(router_probs, self.top_k, dim=-1)
+        top_k_weights = top_k_weights / top_k_weights.sum(dim=-1, keepdim=True)
+        top_k_weights = top_k_weights.to(hidden_states.dtype)
+
+        # Routed expert output
+        routed_out = self.routed_experts(flat_hidden, top_k_indices, top_k_weights)
+        routed_out = routed_out.view(batch_size, seq_len, -1)
+
+        # Combine: shared expert baseline + routed experts (grow in via zero-init down_proj)
+        return shared_out + routed_out
+
+
+def load_balancing_loss(router_logits: torch.Tensor, num_experts: int, top_k: int) -> torch.Tensor:
+    """Auxiliary loss to encourage balanced expert usage."""
+    if router_logits is None:
+        return torch.tensor(0.0)
+
+    probs = F.softmax(router_logits.float(), dim=-1)
+    _, selected = torch.topk(probs, top_k, dim=-1)
+
+    # Fraction of tokens per expert
+    expert_mask = F.one_hot(selected, num_experts).float()
+    tokens_per_expert = expert_mask.mean(dim=(0, 1))
+
+    # Average probability per expert
+    prob_per_expert = probs.mean(dim=0)
+
+    # Balance loss
+    return (tokens_per_expert * prob_per_expert).sum() * num_experts
+
+
+def z_loss(router_logits: torch.Tensor) -> torch.Tensor:
+    """Z-loss to prevent router logits from growing too large.
+
+    From DeepSeek/Switch Transformer: penalizes large logits to keep
+    softmax in its "soft" regime where gradients flow properly.
+    """
+    if router_logits is None:
+        return torch.tensor(0.0)
+
+    # logsumexp ≈ max(logits), squaring penalizes large values
+    return torch.logsumexp(router_logits.float(), dim=-1).square().mean()
+
+
+class SharedMoEAudioProjector(nn.Module):
+    """Shared MoE Audio Projector.
+
+    Combines a shared expert (always-on) with sparse routed experts.
+    Uses zero-initialized router for natural specialization learning.
+
+    Config options:
+        - num_experts: Number of routed experts (default: 4)
+        - num_experts_per_tok: Top-k routing (default: 2)
+        - router_aux_loss_coef: Load balancing loss weight (default: 0.01)
+        - router_z_loss_coef: Z-loss weight to prevent large logits (default: 0.001)
+    """
+
+    def __init__(self, config):
+        super().__init__()
+
+        # Temporal downsampling
+        self.k = getattr(config, "projector_pool_stride", 4)
+
+        # Dimensions
+        self.encoder_dim = config.encoder_dim
+        in_dim = self.encoder_dim * self.k
+        out_dim = config.llm_dim
+        # No expansion - keep hidden dim same as input dim
+        hidden_dim = getattr(config, "projector_hidden_dim", None) or in_dim
+
+        # MoE config
+        self.num_experts = getattr(config, "num_experts", 4)
+        self.top_k = getattr(config, "num_experts_per_tok", 2)
+        self.aux_loss_coef = getattr(config, "router_aux_loss_coef", 0.01)
+        self.z_loss_coef = getattr(config, "router_z_loss_coef", 0.001)
+
+        # Layers
+        self.moe = SharedMoEBlock(in_dim, hidden_dim, out_dim, self.num_experts, self.top_k)
+
+        # Init
+        self._init_weights()
+
+    def _init_weights(self):
+        with torch.no_grad():
+            # Xavier init: std = 1/sqrt(fan_in)
+            in_dim = self.encoder_dim * self.k
+            std = 1.0 / (in_dim ** 0.5)
+
+            # Use a smaller std for the final projection in the shared expert's residual path
+            down_proj_std = std / 2.0
+
+            # Shared expert
+            nn.init.normal_(self.moe.shared_expert.gate_proj.weight, std=std)
+            nn.init.normal_(self.moe.shared_expert.up_proj.weight, std=std)
+            nn.init.normal_(self.moe.shared_expert.down_proj.weight, std=down_proj_std)
+
+            # Routed experts - zero init down_proj so they "grow in" from zero
+            for expert in self.moe.routed_experts.experts:
+                nn.init.normal_(expert.gate_proj.weight, std=std)
+                nn.init.normal_(expert.up_proj.weight, std=std)
+                nn.init.zeros_(expert.down_proj.weight)
+
+            # Router stays zero-initialized
+
+    def forward(self, x: torch.Tensor, attention_mask: torch.Tensor | None = None) -> torch.Tensor:
+        batch_size, seq_len, dim = x.size()
+
+        # Dtype
+        target_dtype = self.moe.shared_expert.gate_proj.weight.dtype
+        if x.dtype != target_dtype:
+            x = x.to(target_dtype)
+
+        # Pad for pooling
+        if seq_len % self.k:
+            x = F.pad(x, (0, 0, 0, self.k - seq_len % self.k))
+            if attention_mask is not None:
+                attention_mask = F.pad(attention_mask, (0, self.k - seq_len % self.k), value=0)
+
+        # Store pooled attention mask for aux loss
+        if attention_mask is not None:
+            # Max-pool the attention mask
+            pooled_mask = F.max_pool1d(attention_mask.float().unsqueeze(1), self.k, self.k)
+            self.last_attention_mask = pooled_mask.squeeze(1).bool()
+        else:
+            self.last_attention_mask = None
+
+        # Temporal pooling
+        x = x.view(batch_size, -1, dim * self.k)
+
+        # Forward
+        x = self.moe(x)
+
+        return x
+
+    def get_aux_loss(self) -> torch.Tensor:
+        """Get auxiliary losses (call after forward).
+
+        Combines:
+        - Load balancing loss: encourages balanced expert usage
+        - Z-loss: prevents router logits from growing too large
+        """
+        router_logits = self.moe.last_router_logits
+        if router_logits is None:
+            return torch.tensor(0.0, device=self.moe.router.weight.device)
+
+        # Retrieve the attention mask stored during the forward pass
+        attention_mask = getattr(self, "last_attention_mask", None)
+
+        # If a mask exists, filter the logits to only include un-padded tokens
+        if attention_mask is not None:
+            flat_mask = attention_mask.view(-1)
+            # Ensure the mask is not all False, which would create an empty tensor
+            if flat_mask.any():
+                active_logits = router_logits[flat_mask]
+            else:
+                # If the mask is all False, there are no tokens to compute loss on
+                return torch.tensor(0.0, device=router_logits.device)
+        else:
+            active_logits = router_logits
+
+        balance_loss = load_balancing_loss(active_logits, self.num_experts, self.top_k)
+        z = z_loss(active_logits)
+
+        return self.aux_loss_coef * balance_loss + self.z_loss_coef * z

swiglu_projector.py

@@ -25,34 +25,34 @@ class SwiGLU(nn.Module):
 class AudioProjector(nn.Module):
     def __init__(self, config):
         super().__init__()
-        self.k = getattr(config, "projector_pool_stride", 4)  # Downsampling rate
+        self.k = getattr(config, "projector_pool_stride", 4)
         in_dim = config.encoder_dim * self.k
         out_dim = config.llm_dim
         hidden_dim = config.projector_hidden_dim
         if hidden_dim is None:
-            hidden_dim = config.encoder_dim * 4
+            hidden_dim = config.encoder_dim * 2
 
         dropout_rate = getattr(config, "projector_dropout", 0.0)
 
-        from transformers.models.llama.modeling_llama import LlamaRMSNorm
-
-        self.ln_pre = LlamaRMSNorm(in_dim, eps=1e-6)
-        self.proj = SwiGLU(in_dim, hidden_dim, out_dim, dropout=dropout_rate)
-        self.ln_post = LlamaRMSNorm(out_dim, eps=1e-6)
+        self.proj1 = SwiGLU(in_dim, hidden_dim, hidden_dim, dropout=dropout_rate)
+        self.proj2 = SwiGLU(hidden_dim, hidden_dim, out_dim, dropout=dropout_rate)
         self.output_dropout = nn.Dropout(dropout_rate)
 
         with torch.no_grad():
             std = getattr(config, "projector_init_std", 0.02)
-            self.ln_pre.weight.data.fill_(1.0)
-            self.ln_post.weight.data.fill_(1.0)
-            nn.init.normal_(self.proj.w1.weight, mean=0.0, std=std)
-            nn.init.normal_(self.proj.w2.weight, mean=0.0, std=std)
-            nn.init.normal_(self.proj.w3.weight, mean=0.0, std=std)
+            # Initialize first layer
+            nn.init.normal_(self.proj1.w1.weight, mean=0.0, std=std)
+            nn.init.normal_(self.proj1.w2.weight, mean=0.0, std=std)
+            nn.init.normal_(self.proj1.w3.weight, mean=0.0, std=std)
+            # Initialize second layer
+            nn.init.normal_(self.proj2.w1.weight, mean=0.0, std=std)
+            nn.init.normal_(self.proj2.w2.weight, mean=0.0, std=std)
+            nn.init.normal_(self.proj2.w3.weight, mean=0.0, std=std)
 
     def forward(self, x):
         batch_size, seq_len, dim = x.size()
 
-        target_dtype = self.proj.w1.weight.dtype
+        target_dtype = self.proj1.w1.weight.dtype
         if x.dtype != target_dtype:
             x = x.to(target_dtype)
 
@@ -62,8 +62,7 @@ class AudioProjector(nn.Module):
             x = F.pad(x, (0, 0, 0, pad_len))
 
         x = x.contiguous().view(batch_size, -1, dim * self.k)
-        x = self.ln_pre(x)
-        x = self.proj(x)
-        x = self.ln_post(x)
+        x = self.proj1(x)
+        x = self.proj2(x)
 
         return self.output_dropout(x)

tokenizer.json

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:d4aeaf198f783cbf58d8cd59812baac429ffe49147bf9648f6618de20b8d4a4c
-size 17209003
+oid sha256:64999f2f5e05d34613701df1999669c5dce7e3891e1628a002518ee68a8626d1
+size 17209101