Remove src/models/moe_swiglu_v2.py

src/models/moe_swiglu_v2.py

@@ -1,182 +0,0 @@
-"""
-Streamlined 16-Expert SwiGLU Mixture-of-Experts (MoE) Layer for Vision Transformers.
-Uses pure SwiGLU FFN experts, Gating Residuals (cross-layer routing memory),
-and Context-Aware Global Projections.
-"""
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-class SwiGLU(nn.Module):
-    """3-Matrix SwiGLU Feed-Forward Network: SiLU(W1 * x) * (W3 * x) @ W2"""
-    def __init__(self, d_model: int, d_ff: int):
-        super().__init__()
-        self.w1 = nn.Linear(d_model, d_ff, bias=False)
-        self.w2 = nn.Linear(d_ff, d_model, bias=False)
-        self.w3 = nn.Linear(d_model, d_ff, bias=False)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return self.w2(F.silu(self.w1(x)) * self.w3(x))
-
-class MoESwiGLUv2(nn.Module):
-    """
-    Pure SwiGLU MoE FFN layer (Feb 2026 specialization-ready).
-    - num_ffn: Number of SwiGLU experts (default 16, use 20 for 3k or 32 for fine granularity)
-    - top_k: Active experts per token (default 2)
-    - gating_residual: Linear projection of previous-layer gate logits
-    - gate_ctx: Global CLS context projection into gate logits
-    - shared_experts: Dense shared expert(s) always active (DeepSeekMoE-style, 0 = pure sparse)
-    """
-    def __init__(
-        self,
-        d_model: int,
-        d_ff: int = 512,
-        num_ffn: int = 16,
-        top_k: int = 2,
-        gating_residual: bool = True,
-        gate_ctx: bool = True,
-        shared_experts: int = 0,
-    ):
-        super().__init__()
-        self.d_model = d_model
-        self.d_ff = d_ff
-        self.num_ffn = num_ffn
-        self.top_k = min(top_k, num_ffn)
-        self.gate_ctx = gate_ctx
-        self.shared_experts = int(shared_experts)
-
-        # Router linear layer: [d_model -> num_ffn] (routed only — shared are dense)
-        self.router = nn.Linear(d_model, num_ffn, bias=False)
-        self.experts = nn.ModuleList([SwiGLU(d_model, d_ff) for _ in range(num_ffn)])
-        # Shared dense experts (always on, no routing) — DeepSeekMoE-style
-        self.shared_mlp = nn.ModuleList([SwiGLU(d_model, d_ff) for _ in range(self.shared_experts)]) if self.shared_experts > 0 else None
-
-        # Cross-layer gating residual projection
-        self.gating_residual = None
-        if gating_residual:
-            self.gating_residual = nn.Linear(num_ffn, num_ffn, bias=False)
-            nn.init.zeros_(self.gating_residual.weight)
-
-        # Context-aware global projection
-        self.ctx_proj = None
-        if gate_ctx:
-            self.ctx_proj = nn.Linear(d_model, num_ffn, bias=False)
-            nn.init.zeros_(self.ctx_proj.weight)
-
-    def forward(self, x: torch.Tensor, prev_gate: torch.Tensor = None, ctx: torch.Tensor = None):
-        """
-        x: [B, T, D] Token embeddings
-        prev_gate: [B, T, num_ffn] Gate logits from previous layer (optional)
-        ctx: [B, D] Global CLS context embedding (optional)
-        Returns:
-          output: [B, T, D]
-          aux_loss: Scalar load balancing loss
-          gate_logits: [B, T, num_ffn]
-        """
-        B, T, D = x.shape
-        x_2d = x.reshape(-1, D)
-        N = x_2d.size(0)
-
-        # 1. Compute Gate Logits
-        gate_logits = self.router(x_2d) # [N, num_ffn]
-
-        if self.ctx_proj is not None and ctx is not None:
-            # Broadcast global CLS context across all T tokens
-            gate_logits = gate_logits + self.ctx_proj(ctx).repeat_interleave(T, dim=0)
-
-        if self.gating_residual is not None and prev_gate is not None:
-            prev_flat = prev_gate.reshape(-1, self.num_ffn)
-            if prev_flat.size(0) == N:
-                gate_logits = gate_logits + self.gating_residual(prev_flat)
-
-        # 2. Softmax & Top-K Selection
-        gate_weights = F.softmax(gate_logits, dim=-1)
-        top_weights, top_indices = gate_weights.topk(self.top_k, dim=-1)
-        top_weights = top_weights / top_weights.sum(dim=-1, keepdim=True).clamp(min=1e-9)
-        route_weights = top_weights.to(dtype=x_2d.dtype)
-
-        # 3. Flatten tokens and sort by expert index for batch GEMM efficiency
-        flat_expert = top_indices.reshape(-1)
-        flat_token = (
-            torch.arange(N, device=x.device)
-            .unsqueeze(1)
-            .expand(-1, self.top_k)
-            .reshape(-1)
-        )
-        flat_weight = route_weights.reshape(-1, 1)
-
-        order = torch.argsort(flat_expert, stable=True)
-        sorted_tokens = flat_token[order]
-        sorted_weights = flat_weight[order]
-        sorted_x = x_2d.index_select(0, sorted_tokens)
-
-        counts = torch.bincount(flat_expert, minlength=self.num_ffn).to(torch.int64)
-        offsets = torch.cat(
-            [torch.zeros(1, dtype=torch.int64, device=x.device),
-             torch.cumsum(counts, 0)[:-1]]
-        )
-
-        output = torch.zeros_like(x_2d)
-
-        # 4. Dispatch each expert — also collect per-expert mean activation for intra-loss
-        expert_mean_acts: list[torch.Tensor] = []
-        for expert_idx in range(self.num_ffn):
-            lo = int(offsets[expert_idx])
-            hi = lo + int(counts[expert_idx])
-            if lo == hi:
-                continue
-            tids = sorted_tokens[lo:hi]
-            w = sorted_weights[lo:hi]
-            # Keep raw expert act before weighting for specialization loss
-            raw_out = self.experts[expert_idx](sorted_x[lo:hi])
-            # For intra loss we store L2-normalized mean act per expert on this batch
-            if counts[expert_idx] > 0:
-                with torch.no_grad():
-                    # detached mean to avoid backprop through loss shaping the act itself too strongly
-                    pass
-                mean_act = raw_out.mean(dim=0)
-                expert_mean_acts.append((expert_idx, mean_act.float()))
-            expert_out = raw_out * w
-            output.index_add_(0, tids, expert_out.to(dtype=output.dtype))
-
-        # 4b. Shared dense experts (always active, DeepSeekMoE-style) — add on top of sparse routed
-        if self.shared_mlp is not None:
-            shared_out = 0
-            for mlp in self.shared_mlp:
-                shared_out = shared_out + mlp(x_2d)
-            # Average shared contribution (keeps scale comparable to sparse)
-            output = output + (shared_out / max(1, len(self.shared_mlp))).to(dtype=output.dtype)
-
-        # 5. Load-Balancing Auxiliary Loss
-        fractions = counts.float() / counts.sum().clamp(min=1e-12).float()
-        avg_prob = gate_weights.mean(dim=0)
-        aux_loss = self.num_ffn * (fractions.to(dtype=gate_weights.dtype) * avg_prob).sum()
-
-        # 6. Intra-layer SwiGLU specialization loss (Feb 2026)
-        # Penalize cosine similarity between distinct routed experts' mean activations on same batch.
-        # Encourages complementarity vs redundant Eucalyptus-venation experts. Cheap: [E, D] matmul.
-        intra_loss = torch.tensor(0.0, device=x.device, dtype=x.dtype)
-        if len(expert_mean_acts) >= 2:
-            # Stack available experts: [E_active, D]
-            idxs, acts = zip(*expert_mean_acts)
-            acts_t = torch.stack(list(acts), dim=0)  # [Ea, D]
-            acts_n = F.normalize(acts_t, dim=1, p=2, eps=1e-6)
-            cos = acts_n @ acts_n.T  # [Ea, Ea]
-            # Mask diagonal, penalize mean absolute cosine (want orthogonal / complementary)
-            mask = torch.eye(cos.size(0), device=cos.device, dtype=torch.bool)
-            # Use upper triangle only
-            cos_off = cos.masked_select(~mask)
-            if cos_off.numel() > 0:
-                # Hinge: only penalize >0.1 similarity, allow slight negative correlation
-                intra_loss = F.relu(cos_off.abs() - 0.1).mean()
-                # Keep in FP32 for stability, then cast back
-                intra_loss = intra_loss.to(dtype=x.dtype)
-
-        # Return intra_loss as 4th element (backward compat: caller can ignore)
-        return (
-            output.view(B, T, D),
-            aux_loss,
-            gate_logits.view(B, T, self.num_ffn),
-            intra_loss,
-            gate_weights.view(B, T, self.num_ffn),
-        )