Upload genesis_architecture.py with huggingface_hub

genesis_architecture.py

@@ -10,42 +10,61 @@ class GenesisConfig:
     embed_dim: int = 768
     num_layers: int = 12
     num_heads: int = 12
-    kv_heads: int = 3
     mlp_dim: int = 2560
+    n_experts: int = 4
+    top_k: int = 2
     tie_embeddings: bool = True
 
-class _GQA(nn.Module):
+class _Attn(nn.Module):
     def __init__(self, cfg):
         super().__init__()
-        self.nh = cfg.num_heads; self.kvh = cfg.kv_heads
+        self.nh = cfg.num_heads
         self.hd = cfg.embed_dim // cfg.num_heads
-        self.g  = cfg.num_heads // cfg.kv_heads
-        self.q  = nn.Linear(cfg.embed_dim, cfg.num_heads * self.hd, bias=False)
-        self.k  = nn.Linear(cfg.embed_dim, cfg.kv_heads * self.hd, bias=False)
-        self.v  = nn.Linear(cfg.embed_dim, cfg.kv_heads * self.hd, bias=False)
-        self.o  = nn.Linear(cfg.embed_dim, cfg.embed_dim, bias=False)
+        self.qkv  = nn.Linear(cfg.embed_dim, 3 * cfg.embed_dim, bias=False)
+        self.proj = nn.Linear(cfg.embed_dim, cfg.embed_dim, bias=False)
     def forward(self, x):
-        B, T, C = x.shape; hd = self.hd
-        q = self.q(x).view(B, T, self.nh,  hd).transpose(1, 2)
-        k = self.k(x).view(B, T, self.kvh, hd).transpose(1, 2).repeat_interleave(self.g, 1)
-        v = self.v(x).view(B, T, self.kvh, hd).transpose(1, 2).repeat_interleave(self.g, 1)
-        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
-        return self.o(y.transpose(1, 2).contiguous().view(B, T, C))
+        B, T, C = x.shape
+        q, k, v = self.qkv(x).split(C, dim=-1)
+        def _s(t): return t.view(B, T, self.nh, self.hd).transpose(1, 2)
+        y = F.scaled_dot_product_attention(_s(q), _s(k), _s(v), is_causal=True)
+        return self.proj(y.transpose(1, 2).contiguous().view(B, T, C))
 
-class _MLP(nn.Module):
-    def __init__(self, cfg):
+class _Expert(nn.Module):
+    def __init__(self, d, ffn):
         super().__init__()
-        self.fc1 = nn.Linear(cfg.embed_dim, cfg.mlp_dim, bias=False)
-        self.fc2 = nn.Linear(cfg.mlp_dim, cfg.embed_dim, bias=False)
+        self.fc1 = nn.Linear(d, ffn, bias=False)
+        self.fc2 = nn.Linear(ffn, d, bias=False)
     def forward(self, x): return self.fc2(F.gelu(self.fc1(x), approximate="tanh"))
 
+class _MoE(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.ne = cfg.n_experts; self.tk = cfg.top_k
+        self.experts = nn.ModuleList([_Expert(cfg.embed_dim, cfg.mlp_dim) for _ in range(cfg.n_experts)])
+        self.gate    = nn.Linear(cfg.embed_dim, cfg.n_experts, bias=False)
+    def forward(self, x):
+        B, T, C = x.shape; flat = x.view(-1, C)
+        probs = torch.softmax(self.gate(flat), dim=-1)
+        top_v, top_i = torch.topk(probs, self.tk, dim=-1)
+        top_v = top_v / (top_v.sum(-1, keepdim=True) + 1e-6)
+        out = torch.zeros_like(flat)
+        for i, expert in enumerate(self.experts):
+            mask = (top_i == i).any(-1)
+            if mask.any():
+                w = torch.where(top_i == i, top_v, torch.zeros_like(top_v)).sum(-1)
+                out[mask] += w[mask].unsqueeze(-1) * expert(flat[mask])
+        aux = self.ne * (probs.mean(0) ** 2).sum()
+        return out.view(B, T, C), aux
+
 class _Block(nn.Module):
     def __init__(self, cfg):
         super().__init__()
-        self.ln1 = nn.LayerNorm(cfg.embed_dim); self.attn = _GQA(cfg)
-        self.ln2 = nn.LayerNorm(cfg.embed_dim); self.mlp  = _MLP(cfg)
+        self.ln1 = nn.LayerNorm(cfg.embed_dim); self.attn = _Attn(cfg)
+        self.ln2 = nn.LayerNorm(cfg.embed_dim); self.moe  = _MoE(cfg)
     def forward(self, x):
-        x = x + self.attn(self.ln1(x)); return x + self.mlp(self.ln2(x))
+        x = x + self.attn(self.ln1(x))
+        moe_out, aux = self.moe(self.ln2(x))
+        return x + moe_out, aux
 
 class GenesisLM(nn.Module):
     def __init__(self, cfg):
@@ -65,5 +84,8 @@ class GenesisLM(nn.Module):
     def forward(self, ids):
         B, T = ids.shape
         x = self.embed(ids) + self.pos_emb(torch.arange(T, device=ids.device))
-        for blk in self.layers: x = blk(x)
-        return self.lm_head(self.norm(x)), None
+        total_aux = torch.zeros(1, device=ids.device)
+        for blk in self.layers:
+            x, aux = blk(x)
+            total_aux = total_aux + aux
+        return self.lm_head(self.norm(x)), total_aux / self.cfg.num_layers