Upload model.py with huggingface_hub

model.py

@@ -0,0 +1,262 @@
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint as ckpt
+
+from config import MoEConfig
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        var = x.pow(2).mean(-1, keepdim=True)
+        x = x * torch.rsqrt(var + self.eps)
+        return x * self.weight
+
+
+def precompute_rope(head_dim: int, max_seq_len: int, theta: float, device, dtype):
+    inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
+    t = torch.arange(max_seq_len, device=device).float()
+    freqs = torch.outer(t, inv_freq)
+    cos = freqs.cos().to(dtype)
+    sin = freqs.sin().to(dtype)
+    return cos, sin
+
+
+def rotate_half(x: torch.Tensor) -> torch.Tensor:
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    # x: (batch, heads, seq, head_dim); cos/sin: (seq, head_dim/2)
+    cos = torch.cat([cos, cos], dim=-1)[None, None, :, :]
+    sin = torch.cat([sin, sin], dim=-1)[None, None, :, :]
+    return x * cos + rotate_half(x) * sin
+
+
+class GQAAttention(nn.Module):
+    def __init__(self, cfg: MoEConfig):
+        super().__init__()
+        self.n_heads = cfg.num_attn_heads
+        self.n_kv_heads = cfg.num_kv_heads
+        self.head_dim = cfg.head_dim
+        self.q_proj = nn.Linear(cfg.hidden_size, self.n_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(cfg.hidden_size, self.n_kv_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(cfg.hidden_size, self.n_kv_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(self.n_heads * self.head_dim, cfg.hidden_size, bias=False)
+
+    def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+        b, t, _ = x.shape
+        q = self.q_proj(x).view(b, t, self.n_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(b, t, self.n_kv_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(b, t, self.n_kv_heads, self.head_dim).transpose(1, 2)
+
+        q = apply_rope(q, cos, sin)
+        k = apply_rope(k, cos, sin)
+
+        rep = self.n_heads // self.n_kv_heads
+        k = k.repeat_interleave(rep, dim=1)
+        v = v.repeat_interleave(rep, dim=1)
+
+        out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        out = out.transpose(1, 2).contiguous().view(b, t, -1)
+        return self.o_proj(out)
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, hidden_size: int, intermediate_size: int):
+        super().__init__()
+        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class DeepSeekMoE(nn.Module):
+    """Fine-grained MoE: many small routed experts + always-on shared experts.
+
+    Standard emergent top-k routing (no domain seeding). Load-balancing is
+    encouraged via an auxiliary loss on router probabilities, not by
+    constraining which tokens go where.
+    """
+
+    def __init__(self, cfg: MoEConfig):
+        super().__init__()
+        self.num_routed = cfg.num_routed_experts
+        self.top_k = cfg.num_experts_per_tok
+        self.aux_loss_coef = cfg.router_aux_loss_coef
+
+        self.shared_experts = nn.ModuleList(
+            [SwiGLU(cfg.hidden_size, cfg.moe_intermediate_size) for _ in range(cfg.num_shared_experts)]
+        )
+        self.routed_experts = nn.ModuleList(
+            [SwiGLU(cfg.hidden_size, cfg.moe_intermediate_size) for _ in range(cfg.num_routed_experts)]
+        )
+        self.gate = nn.Linear(cfg.hidden_size, cfg.num_routed_experts, bias=False)
+
+        self.last_aux_loss: torch.Tensor | None = None
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        b, t, h = x.shape
+        x_flat = x.view(-1, h)  # (N, h)
+
+        out = torch.zeros_like(x_flat)
+        for expert in self.shared_experts:
+            out = out + expert(x_flat)
+
+        router_logits = self.gate(x_flat)  # (N, num_routed)
+        router_probs = F.softmax(router_logits, dim=-1, dtype=torch.float32)
+
+        top_w, top_idx = torch.topk(router_probs, self.top_k, dim=-1)
+        top_w = (top_w / top_w.sum(dim=-1, keepdim=True)).to(x.dtype)  # renormalize
+
+        # Vectorized dispatch: sort the N*top_k (token, expert) assignments by expert
+        # id so each expert's tokens are contiguous, then run each expert once on its
+        # whole contiguous slice. This forces exactly one GPU->CPU sync per layer
+        # (for `counts.tolist()`) instead of one per expert -- 24x fewer syncs than a
+        # per-expert masked-boolean loop, and no wasted matmuls on empty selections.
+        flat_expert_ids = top_idx.reshape(-1)  # (N*top_k,)
+        flat_weights = top_w.reshape(-1)  # (N*top_k,)
+        flat_token_ids = (
+            torch.arange(x_flat.shape[0], device=x.device).unsqueeze(1).expand(-1, self.top_k).reshape(-1)
+        )
+
+        sort_order = torch.argsort(flat_expert_ids)
+        sorted_expert_ids = flat_expert_ids[sort_order]
+        sorted_token_ids = flat_token_ids[sort_order]
+        sorted_weights = flat_weights[sort_order]
+        sorted_x = x_flat[sorted_token_ids]
+
+        counts = torch.bincount(sorted_expert_ids, minlength=self.num_routed).tolist()  # single sync
+        x_chunks = torch.split(sorted_x, counts)
+        weight_chunks = torch.split(sorted_weights, counts)
+        token_id_chunks = torch.split(sorted_token_ids, counts)
+
+        for expert_id, (x_chunk, w_chunk, tok_chunk) in enumerate(zip(x_chunks, weight_chunks, token_id_chunks)):
+            if x_chunk.shape[0] == 0:  # cheap Python-int check, no GPU sync
+                continue
+            expert_out = self.routed_experts[expert_id](x_chunk) * w_chunk.unsqueeze(-1)
+            out.index_add_(0, tok_chunk, expert_out)
+
+        # load-balancing aux loss: encourage uniform expert usage (Switch-style)
+        tokens_per_expert = F.one_hot(top_idx, self.num_routed).sum(dim=1).float().mean(dim=0)  # (num_routed,)
+        prob_per_expert = router_probs.mean(dim=0)  # (num_routed,)
+        self.last_aux_loss = self.aux_loss_coef * self.num_routed * (tokens_per_expert * prob_per_expert).sum()
+
+        return out.view(b, t, h)
+
+
+class Block(nn.Module):
+    def __init__(self, cfg: MoEConfig, layer_idx: int):
+        super().__init__()
+        self.attn_norm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
+        self.attn = GQAAttention(cfg)
+        self.ffn_norm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
+
+        self.is_moe = layer_idx >= cfg.first_k_dense_layers
+        if self.is_moe:
+            self.ffn = DeepSeekMoE(cfg)
+        else:
+            self.ffn = SwiGLU(cfg.hidden_size, cfg.dense_intermediate_size)
+
+    def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+        x = x + self.attn(self.attn_norm(x), cos, sin)
+        x = x + self.ffn(self.ffn_norm(x))
+        return x
+
+
+class MoELanguageModel(nn.Module):
+    def __init__(self, cfg: MoEConfig):
+        super().__init__()
+        self.cfg = cfg
+        self.embed = nn.Embedding(cfg.vocab_size, cfg.hidden_size)
+        self.layers = nn.ModuleList([Block(cfg, i) for i in range(cfg.num_layers)])
+        self.final_norm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
+        self.lm_head = nn.Linear(cfg.hidden_size, cfg.vocab_size, bias=False)
+
+        if cfg.tie_word_embeddings:
+            self.lm_head.weight = self.embed.weight
+
+        self._rope_cache = None
+        self.gradient_checkpointing = False
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module: nn.Module):
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+    def _rope(self, seq_len: int, device, dtype):
+        if self._rope_cache is None or self._rope_cache[0].shape[0] < seq_len:
+            self._rope_cache = precompute_rope(
+                self.cfg.head_dim, max(seq_len, self.cfg.max_seq_len), self.cfg.rope_theta, device, dtype
+            )
+        cos, sin = self._rope_cache
+        return cos[:seq_len], sin[:seq_len]
+
+    def forward(self, input_ids: torch.Tensor, labels: torch.Tensor | None = None):
+        x = self.embed(input_ids)
+        cos, sin = self._rope(input_ids.shape[1], x.device, x.dtype)
+
+        aux_losses = []
+        for layer in self.layers:
+            if self.gradient_checkpointing and self.training:
+                x = ckpt.checkpoint(layer, x, cos, sin, use_reentrant=False)
+            else:
+                x = layer(x, cos, sin)
+            if layer.is_moe and layer.ffn.last_aux_loss is not None:
+                aux_losses.append(layer.ffn.last_aux_loss)
+
+        x = self.final_norm(x)
+        logits = self.lm_head(x)
+
+        if labels is None:
+            return logits, None
+
+        ce = F.cross_entropy(
+            logits[:, :-1, :].reshape(-1, logits.size(-1)),
+            labels[:, 1:].reshape(-1),
+            ignore_index=-100,
+        )
+        aux = torch.stack(aux_losses).sum() if aux_losses else torch.zeros((), device=x.device)
+        loss = ce + aux
+        return logits, loss
+
+    def num_params(self, active_only: bool = False) -> int:
+        if not active_only:
+            return sum(p.numel() for p in self.parameters())
+
+        total = 0
+        for name, p in self.named_parameters():
+            if ".routed_experts." in name:
+                continue  # only a top_k/num_routed fraction is active per token
+            total += p.numel()
+
+        cfg = self.cfg
+        per_expert_params = 3 * cfg.hidden_size * cfg.moe_intermediate_size
+        num_moe_layers = cfg.num_layers - cfg.first_k_dense_layers
+        active_routed_params = num_moe_layers * cfg.num_experts_per_tok * per_expert_params
+        return total + active_routed_params
+
+
+if __name__ == "__main__":
+    cfg = MoEConfig()
+    model = MoELanguageModel(cfg)
+    total = model.num_params()
+    active = model.num_params(active_only=True)
+    print(f"total params:  {total / 1e6:.1f}M")
+    print(f"active params: {active / 1e6:.1f}M")
+
+    x = torch.randint(0, cfg.vocab_size, (2, 128))
+    logits, loss = model(x, labels=x)
+    print("logits:", logits.shape, "loss:", loss.item())