ASTRAI-labs/pluto-nano-0.5

1"""2ASTRAI Pluto — native architecture for the Pluto family.3 4A standalone decoder-only Transformer with:5  * RMSNorm + RoPE (no learned positional embeddings)6  * Causal SDPA attention (multi-head, optional GQA)7  * Top-K Mixture-of-Experts (SwiGLU experts), no required shared expert8  * Multi-Token Prediction heads (training-only)9  * Tied input/output embedding10  * Router auxiliary loss (load balance) + z-loss11 12Not derived from any HuggingFace base model — fresh implementation in plain13PyTorch. Save/load uses a `pluto_config.json` + a safetensors weights file.14 15Naming: `PlutoModel` / `PlutoForCausalLM`. The `_meta` dict on the config holds16size hyper-params; routing / aux-loss config is on its own dataclass.17"""18from __future__ import annotations19 20import json21import math22import os23from dataclasses import asdict, dataclass, field24from pathlib import Path25from typing import Optional26 27import torch28import torch.nn as nn29import torch.nn.functional as F30 31 32# ─── Config ─────────────────────────────────────────────────────────────33 34@dataclass35class PlutoConfig:36    # Architecture (multilingual Nano — d=384, layers=16, GQA, 32k vocab)37    vocab_size: int = 3276838    hidden_size: int = 38439    intermediate_size_expert: int = 153640    intermediate_size_shared: int = 0   # 0 = no shared expert41    n_layers: int = 1642    n_heads: int = 643    n_kv_heads: int = 2                 # GQA: 6→2 → ~50 % attn-param saving44    n_experts: int = 35                 # 5 langs × 7 experts each45    top_k: int = 1                      # max sparsity → ~50 M active inference46    n_languages: int = 5                # en, pt, es, zh, hi47    max_position_embeddings: int = 409648    rope_theta: float = 1_000_000.049    rms_norm_eps: float = 1e-650    tie_word_embeddings: bool = True51 52    # MTP — training-only aux heads53    mtp_depth: int = 254    mtp_loss_weight: float = 0.1555 56    # Routing aux losses57    router_aux_loss_coef: float = 0.0158    router_z_loss_coef: float = 0.00159 60    # Bookkeeping61    model_type: str = "astrai_pluto"62    pad_token_id: int | None = None63    bos_token_id: int | None = None64    eos_token_id: int | None = None65 66    # Tokenizer config (saved for convenience)67    tokenizer_name: str | None = None68 69    def to_dict(self) -> dict:70        return asdict(self)71 72    @classmethod73    def from_dict(cls, d: dict) -> "PlutoConfig":74        # ignore extra keys silently for forward-compat75        known = {f.name for f in cls.__dataclass_fields__.values()}76        return cls(**{k: v for k, v in d.items() if k in known})77 78    def save(self, output_dir: str | Path) -> None:79        os.makedirs(output_dir, exist_ok=True)80        with open(Path(output_dir) / "pluto_config.json", "w") as f:81            json.dump(self.to_dict(), f, indent=2)82 83    @classmethod84    def load(cls, model_dir: str | Path) -> "PlutoConfig":85        with open(Path(model_dir) / "pluto_config.json") as f:86            return cls.from_dict(json.load(f))87 88 89# ─── Layers ─────────────────────────────────────────────────────────────90 91class RMSNorm(nn.Module):92    def __init__(self, dim: int, eps: float = 1e-6):93        super().__init__()94        self.weight = nn.Parameter(torch.ones(dim))95        self.eps = eps96 97    def forward(self, x: torch.Tensor) -> torch.Tensor:98        # Compute in fp32 for numerical stability, return in input dtype99        out = x.float()100        norm = out.pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()101        return (out * norm).to(x.dtype) * self.weight102 103 104def _rope_freqs(dim: int, base: float, device, dtype=torch.float32) -> torch.Tensor:105    inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, device=device, dtype=dtype) / dim))106    return inv_freq107 108 109def _rope_cache(seq_len: int, dim: int, base: float, device) -> tuple[torch.Tensor, torch.Tensor]:110    inv_freq = _rope_freqs(dim, base, device)111    t = torch.arange(seq_len, device=device, dtype=torch.float32)112    freqs = torch.outer(t, inv_freq)113    cos = freqs.cos()114    sin = freqs.sin()115    return cos, sin116 117 118def _apply_rope(q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):119    # q, k: [B, H, T, Dh]; cos, sin: [T, Dh/2]120    def rotate(x: torch.Tensor) -> torch.Tensor:121        x1, x2 = x[..., ::2], x[..., 1::2]122        rot = torch.stack((-x2 * sin + x1 * cos, x1 * sin + x2 * cos), dim=-1)123        return rot.flatten(-2)124    return rotate(q), rotate(k)125 126 127class PlutoAttention(nn.Module):128    """Causal SDPA attention with optional GQA + RoPE."""129    def __init__(self, cfg: PlutoConfig):130        super().__init__()131        assert cfg.hidden_size % cfg.n_heads == 0132        self.cfg = cfg133        self.head_dim = cfg.hidden_size // cfg.n_heads134        self.q_proj = nn.Linear(cfg.hidden_size, cfg.n_heads * self.head_dim, bias=False)135        self.k_proj = nn.Linear(cfg.hidden_size, cfg.n_kv_heads * self.head_dim, bias=False)136        self.v_proj = nn.Linear(cfg.hidden_size, cfg.n_kv_heads * self.head_dim, bias=False)137        self.o_proj = nn.Linear(cfg.hidden_size, cfg.hidden_size, bias=False)138 139    def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:140        B, T, D = x.shape141        H = self.cfg.n_heads142        Hk = self.cfg.n_kv_heads143        Dh = self.head_dim144 145        q = self.q_proj(x).view(B, T, H, Dh).transpose(1, 2)    # [B, H, T, Dh]146        k = self.k_proj(x).view(B, T, Hk, Dh).transpose(1, 2)   # [B, Hk, T, Dh]147        v = self.v_proj(x).view(B, T, Hk, Dh).transpose(1, 2)148        q, k = _apply_rope(q, k, cos[:T].to(q.dtype), sin[:T].to(q.dtype))149        # GQA: expand kv if Hk < H150        if Hk != H:151            repeats = H // Hk152            k = k.repeat_interleave(repeats, dim=1)153            v = v.repeat_interleave(repeats, dim=1)154        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)155        y = y.transpose(1, 2).contiguous().view(B, T, D)156        return self.o_proj(y)157 158 159class SwiGLU(nn.Module):160    def __init__(self, dim: int, hidden: int):161        super().__init__()162        self.w_gate = nn.Linear(dim, hidden, bias=False)163        self.w_up = nn.Linear(dim, hidden, bias=False)164        self.w_down = nn.Linear(hidden, dim, bias=False)165 166    def forward(self, x: torch.Tensor) -> torch.Tensor:167        return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x))168 169 170class PlutoMoE(nn.Module):171    """Top-K MoE using grouped matmul (torch._grouped_mm).172 173    Expert weights are kept as 3 stacked tensors of shape [E, D, H] (gate, up)174    and [E, H, D] (down) so the whole layer is 3 grouped GEMMs per forward.175 176    Currently specialised for top_k == 1 (sort once, no aggregation). Top-K>1177    falls back to the per-expert loop.178 179    Optional shared expert (always active) if intermediate_size_shared > 0.180    """181    def __init__(self, cfg: PlutoConfig):182        super().__init__()183        self.cfg = cfg184        E, D, H = cfg.n_experts, cfg.hidden_size, cfg.intermediate_size_expert185        self.router = nn.Linear(D, E, bias=False)186        # SwiGLU expert weights stacked along the expert dim.187        # `_grouped_mm(A, B, offs)` expects B in [E, K, N] for A in [M, K]188        # → output [M, N]. So we store:189        #   W_gate: [E, D, H]  →  x @ W_gate → [M, H]190        #   W_up:   [E, D, H]191        #   W_down: [E, H, D]192        self.W_gate = nn.Parameter(torch.empty(E, D, H))193        self.W_up   = nn.Parameter(torch.empty(E, D, H))194        self.W_down = nn.Parameter(torch.empty(E, H, D))195        # Init: Kaiming-like, scaled down so initial residual is well-behaved.196        std_in = 1.0 / math.sqrt(D)197        std_h  = 1.0 / math.sqrt(H)198        nn.init.normal_(self.W_gate, std=std_in)199        nn.init.normal_(self.W_up,   std=std_in)200        nn.init.normal_(self.W_down, std=std_h)201        self.shared = (SwiGLU(D, cfg.intermediate_size_shared)202                       if cfg.intermediate_size_shared > 0 else None)203 204    @staticmethod205    def _offsets_from_counts(counts: torch.Tensor) -> torch.Tensor:206        # Convert [E] counts → end-offset tensor [E] of int32.207        # `torch._grouped_mm` consumes end-offsets (exclusive cumsum).208        return counts.cumsum(0).to(torch.int32)209 210    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, dict]:211        B, T, D = x.shape212        E = self.cfg.n_experts213        x_flat = x.reshape(B * T, D)214        logits = self.router(x_flat)                                  # [B*T, E]215 216        if self.cfg.top_k == 1:217            # Sort tokens by expert id → contiguous expert ranges → grouped GEMM218            top_idx = logits.argmax(dim=-1)                           # [B*T]219            sort_idx = top_idx.argsort(stable=True)220            x_sorted = x_flat[sort_idx]                                # [B*T, D]221 222            counts = torch.bincount(top_idx, minlength=E)             # [E]223            offsets = self._offsets_from_counts(counts)               # [E] end-offsets224 225            # Grouped SwiGLU: each token uses ONE expert.226            gate = torch._grouped_mm(x_sorted, self.W_gate, offsets)  # [B*T, H]227            up   = torch._grouped_mm(x_sorted, self.W_up,   offsets)  # [B*T, H]228            hidden = F.silu(gate) * up229            out_sorted = torch._grouped_mm(hidden, self.W_down, offsets)  # [B*T, D]230 231            # Un-sort232            inverse = torch.empty_like(sort_idx)233            inverse[sort_idx] = torch.arange(sort_idx.size(0), device=x.device)234            out = out_sorted[inverse]235        else:236            # Top-K>1 fallback: slower loop. Kept for completeness.237            topk_vals, topk_idx = logits.topk(self.cfg.top_k, dim=-1)238            topk_w = F.softmax(topk_vals, dim=-1)239            out = torch.zeros_like(x_flat)240            for k in range(self.cfg.top_k):241                ids = topk_idx[..., k]242                w = topk_w[..., k].unsqueeze(-1)243                # Per-K grouped GEMM244                sort_idx = ids.argsort(stable=True)245                x_sorted = x_flat[sort_idx]246                counts = torch.bincount(ids, minlength=E)247                offsets = self._offsets_from_counts(counts)248                gate = torch._grouped_mm(x_sorted, self.W_gate, offsets)249                up   = torch._grouped_mm(x_sorted, self.W_up,   offsets)250                hidden = F.silu(gate) * up251                out_sorted = torch._grouped_mm(hidden, self.W_down, offsets)252                inverse = torch.empty_like(sort_idx)253                inverse[sort_idx] = torch.arange(sort_idx.size(0), device=x.device)254                out = out + out_sorted[inverse] * w255            top_idx = topk_idx[..., 0]   # for aux-loss bookkeeping below256 257        if self.shared is not None:258            out = out + self.shared(x_flat)259        out = out.reshape(B, T, D)260 261        # Auxiliary losses (Switch Transformer load-balance + ST-MoE z-loss)262        aux: dict = {}263        if self.training:264            probs = F.softmax(logits.float(), dim=-1)265            expert_freq = probs.mean(dim=0)                            # [E]266            counts_norm = (counts.float() / counts.float().sum().clamp_min(1.0))267            aux["aux_load"] = (expert_freq * counts_norm).sum() * self.cfg.n_experts268            aux["aux_z"]    = (logits.float().logsumexp(-1) ** 2).mean()269        return out, aux270 271 272class PlutoBlock(nn.Module):273    def __init__(self, cfg: PlutoConfig):274        super().__init__()275        self.ln1 = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)276        self.attn = PlutoAttention(cfg)277        self.ln2 = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)278        self.moe = PlutoMoE(cfg)279 280    def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> tuple[torch.Tensor, dict]:281        x = x + self.attn(self.ln1(x), cos, sin)282        y, aux = self.moe(self.ln2(x))283        x = x + y284        return x, aux285 286 287# ─── Models ─────────────────────────────────────────────────────────────288 289class PlutoModel(nn.Module):290    """Decoder backbone: token embed → N blocks → final RMSNorm."""291    def __init__(self, cfg: PlutoConfig):292        super().__init__()293        self.cfg = cfg294        self.embed_tokens = nn.Embedding(cfg.vocab_size, cfg.hidden_size)295        self.blocks = nn.ModuleList([PlutoBlock(cfg) for _ in range(cfg.n_layers)])296        self.final_norm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)297        self.register_buffer("_rope_initialised", torch.tensor(False), persistent=False)298        self._rope_cos = None299        self._rope_sin = None300 301    def _ensure_rope(self, seq_len: int, device, dtype):302        head_dim = self.cfg.hidden_size // self.cfg.n_heads303        if (self._rope_cos is None or self._rope_cos.size(0) < seq_len304                or self._rope_cos.device != device):305            cos, sin = _rope_cache(self.cfg.max_position_embeddings, head_dim,306                                    self.cfg.rope_theta, device)307            self._rope_cos = cos.to(dtype)308            self._rope_sin = sin.to(dtype)309 310    def forward(self, input_ids: torch.Tensor) -> tuple[torch.Tensor, list[dict]]:311        B, T = input_ids.shape312        h = self.embed_tokens(input_ids)313        self._ensure_rope(T, h.device, h.dtype)314        aux_list = []315        for blk in self.blocks:316            h, aux = blk(h, self._rope_cos, self._rope_sin)317            aux_list.append(aux)318        h = self.final_norm(h)319        return h, aux_list320 321 322class PlutoForCausalLM(nn.Module):323    """LM head + optional MTP heads. Returns full loss in `forward`."""324    def __init__(self, cfg: PlutoConfig):325        super().__init__()326        self.cfg = cfg327        self.model = PlutoModel(cfg)328        self.lm_head = nn.Linear(cfg.hidden_size, cfg.vocab_size, bias=False)329        if cfg.tie_word_embeddings:330            self.lm_head.weight = self.model.embed_tokens.weight331        # MTP — training-only auxiliary heads that predict tokens further ahead.332        self.mtp_heads = nn.ModuleList([333            nn.Linear(cfg.hidden_size, cfg.vocab_size, bias=False)334            for _ in range(cfg.mtp_depth)335        ])336 337    def forward(self, input_ids: torch.Tensor, labels: torch.Tensor | None = None,338                attention_mask: torch.Tensor | None = None,339                ) -> dict:340        # We only honour `labels` from the training harness (HF API).341        if labels is None:342            labels = input_ids343        h, aux_list = self.model(input_ids)344        logits = self.lm_head(h)345        out = {"logits": logits}346 347        # Main next-token loss. Trainer is expected to pass `input_ids = ids[:-1]`348        # and `labels = ids[1:]` so they already align (no internal shift).349        if labels is not None and labels.size(1) == logits.size(1):350            ce = F.cross_entropy(351                logits.float().view(-1, logits.size(-1)),352                labels.view(-1),353                ignore_index=-100,354            )355            loss = ce356            # MTP auxiliary losses: head d predicts the token d positions ahead.357            # Skip entirely when mtp_loss_weight == 0 to save the per-head matmul358            # against the full vocab — that head alone is ~15-20 % of step time.359            if self.cfg.mtp_depth > 0 and self.cfg.mtp_loss_weight > 0:360                mtp_total = 0.0361                for d, head in enumerate(self.mtp_heads, start=1):362                    if labels.size(1) <= d: continue363                    logits_d = head(h)[:, :-d, :].contiguous()364                    labels_d = labels[:,  d:].contiguous()365                    mtp_total = mtp_total + F.cross_entropy(366                        logits_d.float().view(-1, logits_d.size(-1)),367                        labels_d.view(-1),368                        ignore_index=-100,369                    )370                loss = loss + self.cfg.mtp_loss_weight * (mtp_total / max(self.cfg.mtp_depth, 1))371            # Router aux losses (averaged over layers)372            if aux_list and "aux_load" in aux_list[0]:373                aux_load = torch.stack([a["aux_load"] for a in aux_list]).mean()374                aux_z = torch.stack([a["aux_z"] for a in aux_list]).mean()375                loss = (loss + self.cfg.router_aux_loss_coef * aux_load376                        + self.cfg.router_z_loss_coef * aux_z)377            out["loss"] = loss378        return out379 380 381# ─── Save / load ────────────────────────────────────────────────────────382 383def save_pluto(model: PlutoForCausalLM, output_dir: str | Path) -> None:384    model.cfg.save(output_dir)385    from safetensors.torch import save_model386    # `save_model` handles tied weights (embed↔lm_head) by deduplicating them.387    # We must NOT permanently move the model to CPU — restore device after save.388    devices = {p.device for p in model.parameters()}389    device = next(iter(devices)) if len(devices) == 1 else None390    model_cpu = model.cpu()391    save_model(model_cpu, str(Path(output_dir) / "model.safetensors"))392    if device is not None and device.type != "cpu":393        model.to(device)394 395 396def load_pluto(model_dir: str | Path, dtype=torch.bfloat16, map_location="cpu") -> PlutoForCausalLM:397    cfg = PlutoConfig.load(model_dir)398    model = PlutoForCausalLM(cfg).to(dtype)399    from safetensors.torch import load_file400    state = load_file(str(Path(model_dir) / "model.safetensors"), device=str(map_location))401    model.load_state_dict(state, strict=False)402    return model403 404 405# ─── Param accounting ──────────────────────────────────────────────────406 407def count_params(model: nn.Module) -> int:408    return sum(p.numel() for p in model.parameters())409 410 411def estimate_active_params(cfg: PlutoConfig) -> dict:412    """At-inference active params (MTP heads NOT counted, since they are training-only)."""413    head_dim = cfg.hidden_size // cfg.n_heads414    attn_per_layer = (415        cfg.hidden_size * cfg.n_heads * head_dim       # q_proj416        + cfg.hidden_size * cfg.n_kv_heads * head_dim  # k_proj417        + cfg.hidden_size * cfg.n_kv_heads * head_dim  # v_proj418        + cfg.hidden_size * cfg.hidden_size            # o_proj419    )420    expert_size = 3 * cfg.hidden_size * cfg.intermediate_size_expert  # SwiGLU421    shared_size = (3 * cfg.hidden_size * cfg.intermediate_size_shared422                   if cfg.intermediate_size_shared > 0 else 0)423    active_per_layer = attn_per_layer + cfg.top_k * expert_size + shared_size424    active_total = active_per_layer * cfg.n_layers425    # lm_head is also "active" (full matmul against vocab)426    active_total += cfg.vocab_size * cfg.hidden_size427 428    total_experts = expert_size * cfg.n_experts * cfg.n_layers429    total_shared = shared_size * cfg.n_layers430    total_attn = attn_per_layer * cfg.n_layers431    emb_params = cfg.vocab_size * cfg.hidden_size432    lm_head_params = 0 if cfg.tie_word_embeddings else cfg.vocab_size * cfg.hidden_size433    mtp_params = cfg.mtp_depth * cfg.vocab_size * cfg.hidden_size434    total_params = (total_experts + total_shared + total_attn + emb_params435                    + lm_head_params + mtp_params436                    + 2 * cfg.n_layers * cfg.hidden_size   # RMSNorm weights437                    + cfg.hidden_size)438    return {439        "total_params": total_params,440        "active_inference_params": active_total,441        "expert_total_params": total_experts,442        "attn_total_params": total_attn,443        "embedding_params": emb_params,444        "lm_head_params": lm_head_params,445        "mtp_head_params": mtp_params,446    }447 448 449if __name__ == "__main__":450    cfg = PlutoConfig()451    stats = estimate_active_params(cfg)452    for k, v in stats.items():453        print(f"  {k:<28} {v/1e6:>8.2f} M")454    print(f"  active/total ratio          {stats['active_inference_params']/stats['total_params']*100:>5.2f} %")455 456    m = PlutoForCausalLM(cfg)457    n_real = count_params(m)458    print(f"\n  real (actual) total         {n_real/1e6:>8.2f} M")459    x = torch.randint(0, cfg.vocab_size, (2, 32))460    out = m(x, labels=x)461    print(f"  fwd OK   logits {tuple(out['logits'].shape)}  loss={out['loss'].item():.4f}")462