ASTRAI-labs/pluto-nano-0.5-base
021
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