Spaces:
Paused
Paused
File size: 6,875 Bytes
6f75f79 06d2e6b 6f3cb01 06d2e6b 6f75f79 06d2e6b 6f75f79 06d2e6b 6f75f79 06d2e6b 6f3cb01 06d2e6b 6f75f79 06d2e6b 6f75f79 06d2e6b 6f75f79 06d2e6b 6f75f79 06d2e6b 6f75f79 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 | """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"]
@dataclass(frozen=True)
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]
|