Delete moe_projector.py

moe_projector.py

@@ -1,162 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F  # noqa: N812
-from transformers.models.llama.modeling_llama import LlamaRMSNorm
-
-
-class SimpleAdapter(nn.Module):
-    """
-    MOSA Section III-B:
-    "consists of two linear layers with a ReLU activation in between,
-    projecting the hidden dimension from 3072 to 4096 and back to 3072."
-    """
-
-    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):
-        x = self.fc1(x)
-        x = self.relu(x)
-        x = self.dropout(x)
-        return self.fc2(x)
-
-
-class MoEAudioProjector(nn.Module):
-    """
-    MOSA-style projector: Mixture of Simple Adapters.
-
-    From paper (arXiv:2508.18998):
-    - Dense mixture (softmax over ALL experts) instead of sparse Top-K
-    - Simple Linear->ReLU->Linear adapters (3072->4096->3072)
-    - No auxiliary losses - just cross-entropy on transcripts
-    - Conv downsampling: stride 4 total (two conv layers, stride 2 each)
-    """
-
-    def __init__(self, config):
-        super().__init__()
-
-        # Dimensions:
-        # Whisper-large-v3 encoder_dim = 1280
-        # SmolLM3-3B hidden_size = 2048
-        self.encoder_dim = config.encoder_dim  # 1280
-        self.llm_dim = config.llm_dim  # 2048
-
-        # Number of experts: Base=4, Large=8
-        self.num_experts = getattr(config, "num_experts", 4)
-
-        # 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
-        self.conv = nn.Sequential(
-            nn.Conv1d(self.encoder_dim, self.llm_dim, kernel_size=3, stride=2, padding=1),
-            nn.ReLU(),
-            nn.Conv1d(self.llm_dim, self.llm_dim, kernel_size=3, stride=2, padding=1),
-            nn.ReLU(),
-        )
-
-        # --- Router (Section III-B) ---
-        # Base: "two linear layers... mapping from 1280 to 512 and finally to 4"
-        router_hidden = 512
-        self.router = nn.Sequential(
-            nn.Linear(self.encoder_dim, router_hidden),
-            nn.ReLU(),
-            nn.Linear(router_hidden, self.num_experts),
-        )
-
-        # --- Experts / Adapters (Section III-B) ---
-        # "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, dropout=self.dropout_rate)
-                for _ in range(self.num_experts)
-            ]
-        )
-
-        # Normalization for stability (not in original MOSA but prevents FPE)
-        self.ln_post = LlamaRMSNorm(self.llm_dim, eps=1e-6)
-
-        # Initialize weights
-        self._init_weights()
-
-    def _init_weights(self):
-        """Initialize weights for stable training."""
-        std = 0.02
-        with torch.no_grad():
-            # Conv layers
-            for module in self.conv:
-                if isinstance(module, nn.Conv1d):
-                    nn.init.normal_(module.weight, mean=0.0, std=std)
-                    if module.bias is not None:
-                        nn.init.zeros_(module.bias)
-
-            # Router
-            for module in self.router:
-                if isinstance(module, nn.Linear):
-                    nn.init.normal_(module.weight, mean=0.0, std=std)
-                    if module.bias is not None:
-                        nn.init.zeros_(module.bias)
-
-            # Experts
-            for expert in self.experts:
-                nn.init.normal_(expert.fc1.weight, mean=0.0, std=std)
-                nn.init.normal_(expert.fc2.weight, mean=0.0, std=std)
-                if expert.fc1.bias is not None:
-                    nn.init.zeros_(expert.fc1.bias)
-                if expert.fc2.bias is not None:
-                    nn.init.zeros_(expert.fc2.bias)
-
-            # LayerNorm
-            self.ln_post.weight.data.fill_(1.0)
-
-    def forward(self, x):
-        """
-        Args:
-            x: [batch_size, seq_len, encoder_dim] from Whisper encoder (1280)
-
-        Returns:
-            output: [batch_size, seq_len // 4, llm_dim] (3072)
-        """
-        batch_size, seq_len, _ = x.shape
-
-        # Pad to be divisible by stride (4)
-        pad_amt = (4 - (seq_len % 4)) % 4
-        if pad_amt > 0:
-            x = F.pad(x, (0, 0, 0, pad_amt))
-            seq_len = x.shape[1]
-
-        # 1. Convolutional Downsampling
-        # (B, T, C) -> (B, C, T) -> conv -> (B, C, T//4) -> (B, T//4, C)
-        h_conv = self.conv(x.permute(0, 2, 1)).permute(0, 2, 1)
-
-        # 2. Router on high-res input, then downsample weights
-        router_logits = self.router(x)  # [B, T, num_experts]
-        # Average over stride window to match conv output
-        router_logits = router_logits.view(batch_size, seq_len // 4, 4, self.num_experts).mean(
-            dim=2
-        )
-        # Dense softmax
-        routing_weights = F.softmax(router_logits, dim=-1)  # [B, T//4, num_experts]
-
-        # 3. Weighted sum of expert outputs (Eq. 2: y = sum(w_i * E_i(x)))
-        # Use in-place add to reduce memory allocations
-        final_out = torch.zeros_like(h_conv)
-        for i, expert in enumerate(self.experts):
-            expert_out = expert(h_conv)
-            expert_weight = routing_weights[:, :, i : i + 1]
-            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)