Upload v6_model.py with huggingface_hub

v6_model.py

@@ -0,0 +1,256 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-5):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        norm = x.norm(dim=-1, keepdim=True) * (1.0 / (x.size(-1) ** 0.5))
+        return self.weight * x / (norm + self.eps)
+
+
+def _rotate_half(x: torch.Tensor) -> torch.Tensor:
+    x1 = x[..., ::2]
+    x2 = x[..., 1::2]
+    return torch.stack((-x2, x1), dim=-1).flatten(-2)
+
+
+def _build_rope_cache(
+    seq_len: int,
+    dim: int,
+    device: torch.device,
+    dtype: torch.dtype,
+    theta: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    positions = torch.arange(seq_len, device=device, dtype=torch.float32)
+    freq_seq = torch.arange(0, dim, 2, device=device, dtype=torch.float32)
+    inv_freq = theta ** (-freq_seq / dim)
+    angles = torch.outer(positions, inv_freq)
+    cos = torch.repeat_interleave(torch.cos(angles), 2, dim=-1).to(dtype=dtype)
+    sin = torch.repeat_interleave(torch.sin(angles), 2, dim=-1).to(dtype=dtype)
+    return cos[None, None, :, :], sin[None, None, :, :]
+
+
+def _apply_rope(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    q_out = (q * cos) + (_rotate_half(q) * sin)
+    k_out = (k * cos) + (_rotate_half(k) * sin)
+    return q_out, k_out
+
+
+class ExpertEncoderMultiHot(nn.Module):
+    def __init__(
+        self,
+        num_experts: int,
+        num_layers: int,
+        d_model: int,
+        layer_hidden: int,
+        layer_proj: int,
+        dropout: float,
+        layer_gating: bool,
+    ):
+        super().__init__()
+        self.num_experts = num_experts
+        self.num_layers = num_layers
+        self.layer_gating = layer_gating
+        if layer_gating:
+            self.layer_gate = nn.Parameter(torch.zeros(num_layers))
+        self.layer_norm = nn.LayerNorm(num_experts)
+        self.layer_mlp = nn.Sequential(
+            nn.Linear(num_experts, layer_hidden),
+            nn.ReLU(),
+            nn.Linear(layer_hidden, layer_proj),
+        )
+        self.proj = nn.Linear(num_layers * layer_proj, d_model)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, expert_idx: torch.Tensor) -> torch.Tensor:
+        bsz, seq_len, num_layers, _topk = expert_idx.shape
+        multihot = torch.zeros(
+            (bsz, seq_len, num_layers, self.num_experts),
+            device=expert_idx.device,
+            dtype=torch.float32,
+        )
+        multihot.scatter_(-1, expert_idx, 1.0)
+        if self.layer_gating:
+            gate = torch.sigmoid(self.layer_gate).view(1, 1, num_layers, 1)
+            multihot = multihot * gate
+        multihot = self.layer_norm(multihot)
+        layer_repr = self.layer_mlp(multihot)
+        flat = layer_repr.reshape(bsz, seq_len, num_layers * layer_repr.size(-1))
+        return self.dropout(self.proj(flat))
+
+
+class EncoderBlock(nn.Module):
+    def __init__(
+        self,
+        d_model: int,
+        n_head: int,
+        d_ff: int,
+        dropout: float,
+        position_type: str,
+        rope_theta: float,
+        qk_norm: bool,
+        qk_norm_eps: float,
+    ):
+        super().__init__()
+        if d_model % n_head != 0:
+            raise ValueError(f"d_model={d_model} must be divisible by n_head={n_head}")
+        self.n_head = n_head
+        self.d_model = d_model
+        self.d_head = d_model // n_head
+        if self.d_head % 2 != 0:
+            raise ValueError(f"Rotary head dimension must be even, got {self.d_head}")
+        self.position_type = position_type
+        self.rope_theta = rope_theta
+        self.qk_norm = qk_norm
+        self.attn_norm = RMSNorm(d_model)
+        self.mlp_norm = RMSNorm(d_model)
+        self.attn = nn.Linear(d_model, 3 * d_model)
+        self.proj = nn.Linear(d_model, d_model)
+        self.dropout = nn.Dropout(dropout)
+        self.fc = nn.Linear(d_model, d_ff)
+        self.fc_out = nn.Linear(d_ff, d_model)
+        if self.qk_norm:
+            self.q_norm = RMSNorm(self.d_head, eps=qk_norm_eps)
+            self.k_norm = RMSNorm(self.d_head, eps=qk_norm_eps)
+        else:
+            self.q_norm = None
+            self.k_norm = None
+
+    def forward(self, x: torch.Tensor, attention_mask: torch.Tensor | None) -> torch.Tensor:
+        bsz, seq_len, d_model = x.shape
+        mask = None
+        if attention_mask is not None:
+            attention_mask = attention_mask.to(torch.bool)
+            mask = attention_mask.unsqueeze(-1).to(x.dtype)
+            x = x * mask
+
+        qkv = self.attn(self.attn_norm(x))
+        q, k, v = qkv.split(d_model, dim=-1)
+        q = q.view(bsz, seq_len, self.n_head, self.d_head).transpose(1, 2)
+        k = k.view(bsz, seq_len, self.n_head, self.d_head).transpose(1, 2)
+        v = v.view(bsz, seq_len, self.n_head, self.d_head).transpose(1, 2)
+
+        if self.qk_norm:
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+
+        if self.position_type == "rope":
+            cos, sin = _build_rope_cache(
+                seq_len=seq_len,
+                dim=self.d_head,
+                device=q.device,
+                dtype=q.dtype,
+                theta=self.rope_theta,
+            )
+            q, k = _apply_rope(q, k, cos, sin)
+
+        attn = F.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+            attn_mask=attention_mask[:, None, None, :] if attention_mask is not None else None,
+            dropout_p=0.0,
+            is_causal=False,
+        )
+        attn = attn.transpose(1, 2).contiguous().view(bsz, seq_len, d_model)
+        x = x + self.dropout(self.proj(attn))
+        mlp = self.fc(self.mlp_norm(x))
+        mlp = torch.relu(mlp).pow(2)
+        x = x + self.dropout(self.fc_out(mlp))
+        if mask is not None:
+            x = x * mask
+        return x
+
+
+class EncoderOnlyModel(nn.Module):
+    def __init__(
+        self,
+        vocab_size: int,
+        num_experts: int,
+        num_layers: int,
+        topk: int,
+        d_model: int,
+        n_head: int,
+        d_ff: int,
+        n_layer: int,
+        dropout: float,
+        max_len: int,
+        layer_gating: bool,
+        logit_softcap: float,
+        layer_hidden: int,
+        layer_proj: int,
+        position_type: str = "learned",
+        rope_theta: float = 10000.0,
+        qk_norm: bool = True,
+        qk_norm_eps: float = 1e-5,
+    ):
+        super().__init__()
+        if position_type not in {"learned", "rope"}:
+            raise ValueError(f"Unsupported position_type: {position_type}")
+        self.encoder_in = ExpertEncoderMultiHot(
+            num_experts=num_experts,
+            num_layers=num_layers,
+            d_model=d_model,
+            layer_hidden=layer_hidden,
+            layer_proj=layer_proj,
+            dropout=dropout,
+            layer_gating=layer_gating,
+        )
+        self.position_type = position_type
+        self.rope_theta = rope_theta
+        self.qk_norm = qk_norm
+        self.qk_norm_eps = qk_norm_eps
+        self.max_len = max_len
+        if self.position_type == "learned":
+            self.pos_emb = nn.Embedding(max_len, d_model)
+        else:
+            self.pos_emb = None
+        self.blocks = nn.ModuleList(
+            [
+                EncoderBlock(
+                    d_model=d_model,
+                    n_head=n_head,
+                    d_ff=d_ff,
+                    dropout=dropout,
+                    position_type=position_type,
+                    rope_theta=rope_theta,
+                    qk_norm=qk_norm,
+                    qk_norm_eps=qk_norm_eps,
+                )
+                for _ in range(n_layer)
+            ]
+        )
+        self.norm = RMSNorm(d_model)
+        self.head = nn.Linear(d_model, vocab_size, bias=False)
+        self.logit_softcap = logit_softcap
+        self.topk = topk
+
+    def forward(self, expert_idx: torch.Tensor, attention_mask: torch.Tensor | None) -> torch.Tensor:
+        bsz, seq_len = expert_idx.shape[:2]
+        if seq_len > self.max_len:
+            raise ValueError(f"seq_len={seq_len} exceeds configured max_len={self.max_len}")
+        x = self.encoder_in(expert_idx)
+        if self.position_type == "learned":
+            pos_ids = torch.arange(seq_len, device=expert_idx.device)
+            pos_ids = pos_ids.unsqueeze(0).expand(bsz, -1)
+            x = x + self.pos_emb(pos_ids)
+        if attention_mask is not None:
+            x = x * attention_mask.unsqueeze(-1).to(x.dtype)
+        for block in self.blocks:
+            x = block(x, attention_mask)
+        x = self.norm(x)
+        logits = self.head(x)
+        if self.logit_softcap and self.logit_softcap > 0:
+            logits = self.logit_softcap * torch.tanh(logits / self.logit_softcap)
+        return logits