Spaces:
Paused
Paused
| """A small GPT harness whose linear layers are swappable (dense / qat / counter / counter_rms), | |
| so the same model is the parity comparison. Pure PyTorch, runs on CPU/CUDA. | |
| """ | |
| from __future__ import annotations | |
| import math | |
| from dataclasses import dataclass | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from .baselines import make_linear | |
| from .counter import CompactCounterLinear | |
| __all__ = ["GPTConfig", "CONFIGS", "GPT"] | |
| class GPTConfig: | |
| vocab_size: int | |
| block_size: int | |
| n_layer: int | |
| n_head: int | |
| n_embd: int | |
| # --- architecture levers (default = the original dense-FFN transformer, unchanged) --- | |
| # ffn: the FFN sub-block. "dense" = fc->gelu->fc2 (original). "moe" = Counter-MoE (M4, the | |
| # validated win: capacity in experts at ~dense active compute). "memory" = product-key | |
| # CounterMemoryFFN (M1, a compute-saver; lower quality than dense, opt-in). | |
| ffn: str = "dense" | |
| ffn_experts: int = 8 # M4: E total experts (capacity, not active compute) | |
| ffn_top_k: int = 2 # M4: experts visited per token (active compute) | |
| ffn_grouped: bool = False # M4: grouped GEMM experts (torch._grouped_mm, no python expert loop) | |
| ffn_cells: int = 16384 # M1: number of memory cells (perfect square) | |
| ffn_k: int = 8 # M1: cells retrieved per token | |
| CONFIGS = { | |
| "micro": GPTConfig(32, 16, 2, 2, 32), | |
| "tiny": GPTConfig(2048, 128, 4, 4, 256), | |
| "s512": GPTConfig(8192, 256, 8, 8, 512), | |
| "small": GPTConfig(8192, 256, 12, 12, 768), | |
| } | |
| def _counter_numeric(kw: dict) -> dict: | |
| """The counter knobs the FFN-level modules (M1/M4) accept -- they take only C/lr/lr_scale, | |
| not the linear-level kw (tile_rows, rms_mode, cache_mode, ...). Filter so an arm built with a | |
| full counter_kw can still select ffn='moe'/'memory' without passing unexpected kwargs.""" | |
| return {k: kw[k] for k in ("C", "lr", "lr_scale") if k in kw} | |
| class Block(nn.Module): | |
| def __init__(self, kind: str, cfg: GPTConfig, **kw) -> None: | |
| super().__init__() | |
| d, nl = cfg.n_embd, cfg.n_layer | |
| rg = 1.0 / math.sqrt(2.0 * nl) | |
| self.n_embd = d # expose for MoD / dim inference | |
| self.ln1 = nn.LayerNorm(d) | |
| self.q = make_linear(kind, d, d, 1.0, **kw) | |
| self.k = make_linear(kind, d, d, 1.0, **kw) | |
| self.v = make_linear(kind, d, d, 1.0, **kw) | |
| self.proj = make_linear(kind, d, d, rg, **kw) | |
| self.ln2 = nn.LayerNorm(d) | |
| self.n_head = cfg.n_head | |
| # FFN sub-block: dense (original) or a counter architecture lever (M4 moe / M1 memory). | |
| self.ffn_kind = cfg.ffn | |
| if cfg.ffn == "dense": | |
| self.fc = make_linear(kind, d, 4 * d, 1.0, **kw) | |
| self.fc2 = make_linear(kind, 4 * d, d, rg, **kw) | |
| self.ffn = None | |
| elif cfg.ffn == "moe": | |
| from .moe_ffn import CounterMoEFFN | |
| self.ffn = CounterMoEFFN(d, n_experts=cfg.ffn_experts, top_k=cfg.ffn_top_k, | |
| grouped=cfg.ffn_grouped, **_counter_numeric(kw)) | |
| elif cfg.ffn == "memory": | |
| from .memory_ffn import CounterMemoryFFN | |
| self.ffn = CounterMemoryFFN(d, n_cells=cfg.ffn_cells, k=cfg.ffn_k, | |
| **_counter_numeric(kw)) | |
| else: | |
| raise ValueError(f"unknown ffn kind: {cfg.ffn!r} (dense|moe|memory)") | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| b, t, c = x.shape | |
| h = self.ln1(x) | |
| q = self.q(h).view(b, t, self.n_head, -1).transpose(1, 2) | |
| k = self.k(h).view(b, t, self.n_head, -1).transpose(1, 2) | |
| v = self.v(h).view(b, t, self.n_head, -1).transpose(1, 2) | |
| a = F.scaled_dot_product_attention(q, k, v, is_causal=True) | |
| a = a.transpose(1, 2).contiguous().view(b, t, c) | |
| x = x + self.proj(a) | |
| h2 = self.ln2(x) | |
| x = x + (self.fc2(F.gelu(self.fc(h2))) if self.ffn is None else self.ffn(h2)) | |
| return x | |
| class GPT(nn.Module): | |
| """kind in {dense, qat, counter, counter_rms, counter_packed, counter_triton, slowfast (M3), | |
| group (M2)} selects the LINEAR type; cfg.ffn in {dense, moe (M4), memory (M1)} selects the FFN | |
| sub-block. counter_kw are forwarded to the counter layers (lr, lr_scale, C, tile_rows, | |
| rms_mode, ...); the FFN-level levers take only the C/lr/lr_scale subset.""" | |
| def __init__(self, kind: str, cfg: GPTConfig, **counter_kw) -> None: | |
| super().__init__() | |
| self.cfg = cfg | |
| self.kind = kind | |
| self.tok = nn.Embedding(cfg.vocab_size, cfg.n_embd) | |
| self.pos = nn.Embedding(cfg.block_size, cfg.n_embd) | |
| nn.init.normal_(self.tok.weight, std=0.02) | |
| nn.init.normal_(self.pos.weight, std=0.02) | |
| self.blocks = nn.ModuleList([Block(kind, cfg, **counter_kw) for _ in range(cfg.n_layer)]) | |
| self.lnf = nn.LayerNorm(cfg.n_embd) | |
| self.head = nn.Linear(cfg.n_embd, cfg.vocab_size, bias=False) | |
| self.head.weight = self.tok.weight # tie input/output embedding | |
| def forward(self, idx: torch.Tensor, targets: torch.Tensor | None = None): | |
| _, t = idx.shape | |
| if t > self.cfg.block_size: | |
| raise ValueError(f"sequence length {t} exceeds block size {self.cfg.block_size}") | |
| pos = torch.arange(t, device=idx.device) | |
| x = self.tok(idx) + self.pos(pos)[None] | |
| for blk in self.blocks: | |
| x = blk(x) | |
| logits = self.head(self.lnf(x)) | |
| loss = None | |
| if targets is not None: | |
| loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.reshape(-1)) | |
| aux = self.aux_loss() # M1/M4 load-balance term (None if no such FFN) | |
| if aux is not None: | |
| loss = loss + aux | |
| return logits, loss | |
| def aux_loss(self) -> torch.Tensor | None: | |
| """Sum the FFN load-balance aux losses (M4 Counter-MoE / M1 memory) weighted by each | |
| module's own aux_loss_weight, from the last forward. None when no FFN exposes one (e.g. | |
| the dense FFN) -- so the training loss is unchanged for the original transformer.""" | |
| total = None | |
| for m in self.modules(): | |
| w = float(getattr(m, "aux_loss_weight", 0.0) or 0.0) | |
| if w > 0.0 and hasattr(m, "last_aux_loss"): | |
| term = w * m.last_aux_loss | |
| total = term if total is None else total + term | |
| return total | |
| def counter_layers(self) -> list[CompactCounterLinear]: | |
| return [m for m in self.modules() if isinstance(m, CompactCounterLinear)] | |
| def trainable_parameters(self): | |
| """Conventional Parameters (embeddings, norms, head) — counter layers self-update | |
| and expose no parameters(), so this is exactly what an AdamW should own.""" | |
| return [p for p in self.parameters() if p.requires_grad] | |