Update modeling_mythos.py

modeling_mythos.py

@@ -9,7 +9,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 from dataclasses import dataclass
 
-from mamba_ssm.modules.mamba3 import Mamba3
+from fla.layers.mamba3 import Mamba3
 from gdn2 import GatedDeltaNet2
 
 
@@ -28,7 +28,6 @@ class MythosConfig:
     coda_attn_layers: int = 2
     max_loop_iters: int = 32
     act_threshold: float = 0.99
-    hyper_n_streams: int = 4
     audio_vocab: int = 1024
     audio_n_codebooks: int = 8
     image_vocab: int = 8192
@@ -36,6 +35,14 @@ class MythosConfig:
     vision_encoder_dim: int = 1152
     rope_theta: float = 500000.0
     lora_rank: int = 64
+
+    # MoE upscale at the mla_exit FFN site only. Disabled by default --
+    # dense backbone trains with plain FFN everywhere. Flip moe_enabled
+    # to True to swap in MoEFFN at that one site.
+    moe_enabled: bool = False
+    moe_n_experts: int = 8
+    moe_expert_dim: int = 1024
+    moe_top_k: int = 2
     ffn_mult: float = 1.333
     tie_embeddings: bool = True
     mamba_headdim: int = 64
@@ -64,10 +71,57 @@ class FFN(nn.Module):
         return self.down(F.silu(self.gate(x)) * self.up(x))
 
 
+class MoEFFN(nn.Module):
+    """MoE upscale for the mla_exit FFN site. Disabled by default
+    (cfg.moe_enabled=False) -- the dense backbone trains with plain FFN.
+    When enabled, swaps in N experts (each a standard SwiGLU FFN at
+    expert_dim) with top-k routing.
+    """
+    def __init__(self, cfg):
+        super().__init__()
+        self.n_experts    = cfg.moe_n_experts
+        self.top_k        = cfg.moe_top_k
+        self.router       = nn.Linear(cfg.dim, cfg.moe_n_experts, bias=False)
+        self.experts      = nn.ModuleList([
+            FFN(cfg.dim, cfg.moe_expert_dim) for _ in range(cfg.moe_n_experts)
+        ])
+
+    def forward(self, x):
+        B, T, D = x.shape
+        x_flat = x.view(B * T, D)
+
+        router_logits  = self.router(x_flat)
+        router_weights = F.softmax(router_logits, dim=-1, dtype=torch.float32)
+        top_weights, top_idx = router_weights.topk(self.top_k, dim=-1)
+        top_weights = (top_weights / top_weights.sum(dim=-1, keepdim=True)).to(x.dtype)
+
+        out = torch.zeros_like(x_flat)
+        for k in range(self.top_k):
+            idx = top_idx[:, k]
+            w   = top_weights[:, k].unsqueeze(-1)
+            for e in range(self.n_experts):
+                mask = idx == e
+                if mask.any():
+                    out[mask] += w[mask] * self.experts[e](x_flat[mask])
+
+        return out.view(B, T, D)
+
+
 def make_ffn(cfg):
     return FFN(cfg.dim, int(cfg.dim * cfg.ffn_mult))
 
 
+def make_ffn_mla(cfg):
+    """FFN factory specifically for the mla_exit site -- the one place
+    the original design called out for MoE upscaling. Everywhere else
+    (M3/GDN2 blocks) always uses plain make_ffn(), regardless of
+    cfg.moe_enabled.
+    """
+    if cfg.moe_enabled:
+        return MoEFFN(cfg)
+    return make_ffn(cfg)
+
+
 def precompute_freqs_cis(head_dim, max_len, theta):
     inv = 1.0 / (theta ** (torch.arange(0, head_dim, 2).float() / head_dim))
     t = torch.arange(max_len).float()
@@ -162,7 +216,7 @@ class MLADenseBlock(nn.Module):
         self.norm1 = RMSNorm(cfg.dim)
         self.attn  = MLAttention(cfg)
         self.norm2 = RMSNorm(cfg.dim)
-        self.ffn   = make_ffn(cfg)
+        self.ffn   = make_ffn_mla(cfg)
 
     def forward(self, x, freqs_cis):
         x = x + self.attn(self.norm1(x), freqs_cis)
@@ -175,11 +229,11 @@ class Mamba3Block(nn.Module):
         super().__init__()
         self.norm1 = RMSNorm(cfg.dim)
         self.mixer = Mamba3(
-            d_model=cfg.dim,
-            d_state=cfg.mamba_d_state,
-            headdim=cfg.mamba_headdim,
+            hidden_size=cfg.dim,
+            state_size=cfg.mamba_d_state,
+            head_dim=cfg.mamba_headdim,
             expand=1,
-            ngroups=1,
+            n_groups=1,
             is_mimo=False,
             chunk_size=64,
         )
@@ -187,7 +241,7 @@ class Mamba3Block(nn.Module):
         self.ffn   = make_ffn(cfg)
 
     def forward(self, x):
-        x = x + self.mixer(self.norm1(x))
+        x = x + self.mixer(self.norm1(x))[0]
         x = x + self.ffn(self.norm2(x))
         return x
 
@@ -211,40 +265,6 @@ class GDN2Block(nn.Module):
         return x
 
 
-class HyperConnections(nn.Module):
-    def __init__(self, dim, n_streams):
-        super().__init__()
-        self.n = n_streams
-        self.static_beta   = nn.Parameter(torch.ones(n_streams) / n_streams)
-        self.static_alpha  = nn.Parameter(torch.eye(n_streams))
-        self.static_alpha0 = nn.Parameter(torch.ones(n_streams) / n_streams)
-        self.norm       = RMSNorm(dim)
-        self.dyn_beta   = nn.Linear(dim, 1, bias=False)
-        self.dyn_alpha  = nn.Linear(dim, n_streams, bias=False)
-        self.dyn_alpha0 = nn.Linear(dim, 1, bias=False)
-        nn.init.zeros_(self.dyn_beta.weight)
-        nn.init.zeros_(self.dyn_alpha.weight)
-        nn.init.zeros_(self.dyn_alpha0.weight)
-
-    def expand(self, x):
-        return x.unsqueeze(-2).expand(*x.shape[:-1], self.n, x.shape[-1]).contiguous()
-
-    def width(self, streams):
-        ns = self.norm(streams)
-        beta = self.static_beta + torch.tanh(self.dyn_beta(ns)).squeeze(-1)
-        return (beta.unsqueeze(-1) * streams).sum(dim=-2)
-
-    def depth(self, streams, block_out):
-        ns = self.norm(streams)
-        alpha  = self.static_alpha  + torch.tanh(self.dyn_alpha(ns))
-        alpha0 = self.static_alpha0 + torch.tanh(self.dyn_alpha0(ns)).squeeze(-1)
-        mixed  = torch.einsum("btij,btjd->btid", alpha, streams)
-        return mixed + alpha0.unsqueeze(-1) * block_out.unsqueeze(-2)
-
-    def collapse(self, streams):
-        return streams.sum(dim=-2)
-
-
 class LTIInjection(nn.Module):
     def __init__(self, dim):
         super().__init__()
@@ -295,24 +315,11 @@ class RecurrentDenseBlock(nn.Module):
         super().__init__()
         self.cfg = cfg
 
-        # Pattern: (M3 -> GDN2 -> MLA) x 4
+        # Pattern: (M3 -> GDN2 -> MLA) x 1, looped max_loop_iters times via forward()
         self.m3_0  = Mamba3Block(cfg)
         self.gdn_0 = GDN2Block(cfg)
         self.mla_0 = MLADenseBlock(cfg)
 
-        self.m3_1  = Mamba3Block(cfg)
-        self.gdn_1 = GDN2Block(cfg)
-        self.mla_1 = MLADenseBlock(cfg)
-
-        self.m3_2  = Mamba3Block(cfg)
-        self.gdn_2 = GDN2Block(cfg)
-        self.mla_2 = MLADenseBlock(cfg)
-
-        self.m3_3  = Mamba3Block(cfg)
-        self.gdn_3 = GDN2Block(cfg)
-        self.mla_3 = MLADenseBlock(cfg)
-
-        self.hyper     = HyperConnections(cfg.dim, cfg.hyper_n_streams)
         self.injection = LTIInjection(cfg.dim)
         self.act       = ACTHalting(cfg.dim)
         self.lora      = LoRADepthAdapter(cfg.dim, cfg.lora_rank, cfg.max_loop_iters)
@@ -324,38 +331,25 @@ class RecurrentDenseBlock(nn.Module):
         x = self.gdn_0(x)
         x = self.mla_0(x, freqs_cis)
 
-        x = self.m3_1(x)
-        x = self.gdn_1(x)
-        x = self.mla_1(x, freqs_cis)
-
-        x = self.m3_2(x)
-        x = self.gdn_2(x)
-        x = self.mla_2(x, freqs_cis)
-
-        x = self.m3_3(x)
-        x = self.gdn_3(x)
-        x = self.mla_3(x, freqs_cis)
-
         return x + self.lora(x, t)
 
     def forward(self, h, e, freqs_cis, n_loops=None):
         n_loops = n_loops or self.cfg.max_loop_iters
         B, T, D = h.shape
 
-        streams      = self.hyper.expand(h)
+        x            = h
         halted       = torch.zeros(B, T, device=h.device, dtype=torch.bool)
         cumulative_p = torch.zeros(B, T, device=h.device, dtype=torch.float32)
         h_out        = torch.zeros(B, T, D, device=h.device, dtype=torch.float32)
 
         for t in range(n_loops):
-            x = self.hyper.width(streams)
             x = self.loop_emb(x, t)
             x = self.injection(x, e, torch.zeros_like(x))
             x = self.norm(x)
             block_out = self.inner(x, freqs_cis, t)
-            streams   = self.hyper.depth(streams, block_out)
+            x         = x + block_out
 
-            h_cur = self.hyper.collapse(streams)
+            h_cur = x
             p     = self.act(h_cur)
 
             still     = (~halted).float()