epoyraz/tinystories-25m
371
1import torch2import torch.nn as nn3import torch.nn.functional as F4import math5from torch.utils.checkpoint import checkpoint6 7 8def soft_cap(logits, cap):9 """Gemma2/modded-nanoGPT logit soft-capping: cap * tanh(logits / cap). No-op if cap falsy."""10 if cap:11 return cap * torch.tanh(logits / cap)12 return logits13 14 15def chunked_cross_entropy(hidden, weight, targets, cap=0, chunk_size=2048):16 """Memory-efficient cross-entropy. Projects hidden -> logits and reduces the loss17 in token-chunks so the full [N, vocab] logits are never materialized at once (each18 chunk's logits are recomputed in backward via checkpointing). Numerically equal to19 F.cross_entropy(soft_cap(hidden @ weight.T), targets, ignore_index=-1)."""20 hidden = hidden.reshape(-1, hidden.size(-1))21 targets = targets.reshape(-1)22 n_valid = (targets != -1).sum().clamp(min=1)23 24 def chunk_loss(h, t, w):25 logits = soft_cap(F.linear(h, w), cap)26 return F.cross_entropy(logits, t, ignore_index=-1, reduction="sum")27 28 use_ckpt = torch.is_grad_enabled() and (hidden.requires_grad or weight.requires_grad)29 total = hidden.new_zeros(())30 for i in range(0, hidden.size(0), chunk_size):31 h, t = hidden[i:i + chunk_size], targets[i:i + chunk_size]32 if use_ckpt:33 total = total + checkpoint(chunk_loss, h, t, weight, use_reentrant=False)34 else:35 total = total + chunk_loss(h, t, weight)36 return total / n_valid37 38 39# --- mHC: Manifold-Constrained Hyper-Connections ---40 41def sinkhorn(log_alpha, n_iters=5):42 for _ in range(n_iters):43 log_alpha = log_alpha - torch.logsumexp(log_alpha, dim=-1, keepdim=True)44 log_alpha = log_alpha - torch.logsumexp(log_alpha, dim=-2, keepdim=True)45 return log_alpha.exp()46 47 48class MHCResidual(nn.Module):49 def __init__(self, n_streams):50 super().__init__()51 self.n_streams = n_streams52 self.log_alpha = nn.Parameter(torch.zeros(n_streams, n_streams))53 54 def forward(self, streams, update):55 W = sinkhorn(self.log_alpha)56 mixed = torch.einsum("ij,bjte->bite", W, streams)57 mixed[:, 0] = mixed[:, 0] + update58 return mixed59 60 61class MHCExpand(nn.Module):62 def __init__(self, n_streams, n_embd):63 super().__init__()64 self.n_streams = n_streams65 self.proj = nn.Linear(n_embd, n_streams * n_embd) if n_streams > 1 else None66 67 def forward(self, x):68 if self.n_streams == 1:69 return x.unsqueeze(1)70 B, T, C = x.shape71 return self.proj(x).view(B, self.n_streams, T, C)72 73 74class MHCCollapse(nn.Module):75 def __init__(self, n_streams, n_embd):76 super().__init__()77 self.n_streams = n_streams78 self.proj = nn.Linear(n_streams * n_embd, n_embd) if n_streams > 1 else None79 80 def forward(self, streams):81 if self.n_streams == 1:82 return streams.squeeze(1)83 B, S, T, C = streams.shape84 return self.proj(streams.permute(0, 2, 1, 3).reshape(B, T, S * C))85 86 87# --- BitNet: Ternary weight linear layer ---88 89class BitLinear(nn.Module):90 def __init__(self, in_features, out_features, bias=True):91 super().__init__()92 self.in_features = in_features93 self.out_features = out_features94 self.weight = nn.Parameter(torch.empty(out_features, in_features))95 self.bias = nn.Parameter(torch.zeros(out_features)) if bias else None96 self.rms_norm = nn.RMSNorm(in_features)97 nn.init.normal_(self.weight, std=0.02)98 99 def ternary_quantize(self, w):100 alpha = w.abs().mean()101 threshold = alpha * 0.5102 w_ternary = torch.zeros_like(w)103 w_ternary[w > threshold] = alpha104 w_ternary[w < -threshold] = -alpha105 return w_ternary.detach() + (w - w.detach())106 107 def activation_quantize(self, x):108 scale = 127.0 / x.abs().max(dim=-1, keepdim=True).values.clamp(min=1e-5)109 x_scaled = x * scale110 x_q = x_scaled.round().clamp(-128, 127).detach() + (x_scaled - x_scaled.detach())111 return x_q / scale112 113 def forward(self, x):114 x = self.rms_norm(x)115 w_q = self.ternary_quantize(self.weight)116 x_q = self.activation_quantize(x)117 out = F.linear(x_q, w_q, self.bias)118 return out119 120 121class FastBitLinear(nn.Module):122 def __init__(self, in_features, out_features, bias=True):123 super().__init__()124 self.in_features = in_features125 self.out_features = out_features126 self.weight = nn.Parameter(torch.empty(out_features, in_features))127 self.bias = nn.Parameter(torch.zeros(out_features)) if bias else None128 self.rms_norm = nn.RMSNorm(in_features)129 nn.init.normal_(self.weight, std=0.02)130 131 def _int8_forward(self, x):132 w = self.weight.detach()133 alpha = w.abs().mean()134 threshold = alpha * 0.5135 # Pack the ternary weight into a single signed int8 tensor {-1,0,+1} so the136 # whole layer is ONE int8 matmul (dp4a / int8 tensor cores), not two. This is137 # exactly equivalent to (x @ w_pos.T) - (x @ w_neg.T) but ~2x cheaper, and it138 # beats fp16 at prefill/training scale.139 w_ternary = torch.zeros_like(w, dtype=torch.int8)140 w_ternary[w > threshold] = 1141 w_ternary[w < -threshold] = -1142 143 x_max = x.detach().abs().max(dim=-1, keepdim=True).values.clamp(min=1e-5)144 x_scale = 127.0 / x_max145 x_q = (x.detach() * x_scale).round().clamp(-128, 127).to(torch.int8)146 147 shape = x_q.shape148 x_2d = x_q.reshape(-1, shape[-1])149 150 rows = x_2d.shape[0]151 if rows <= 16: # torch._int_mm requires more than 16 rows152 x_2d = torch.nn.functional.pad(x_2d, (0, 0, 0, 17 - rows))153 y = torch._int_mm(x_2d, w_ternary.T)[:rows]154 else:155 y = torch._int_mm(x_2d, w_ternary.T)156 157 y = y.float().reshape(*shape[:-1], self.out_features)158 return y * (alpha / x_scale)159 160 def _ste_forward(self, x):161 alpha = self.weight.abs().mean()162 threshold = alpha * 0.5163 w_ternary = torch.zeros_like(self.weight)164 w_ternary[self.weight > threshold] = alpha165 w_ternary[self.weight < -threshold] = -alpha166 w_q = self.weight + (w_ternary - self.weight).detach()167 168 x_scale = 127.0 / x.abs().max(dim=-1, keepdim=True).values.clamp(min=1e-5)169 x_scaled = x * x_scale170 x_q = x_scaled + (x_scaled.round().clamp(-128, 127) - x_scaled).detach()171 x_q = x_q / x_scale172 173 return F.linear(x_q, w_q, None)174 175 def forward(self, x):176 x = self.rms_norm(x)177 if self.training:178 out = self._ste_forward(x)179 else:180 out = self._int8_forward(x)181 if self.bias is not None:182 out = out + self.bias183 return out184 185 186def make_linear(in_f, out_f, bias=True, use_bitnet=False, use_fast_bitnet=False):187 if use_fast_bitnet:188 return FastBitLinear(in_f, out_f, bias=bias)189 if use_bitnet:190 return BitLinear(in_f, out_f, bias=bias)191 return nn.Linear(in_f, out_f, bias=bias)192 193 194# --- TurboQuant: KV-cache compression for inference ---195 196class PolarQuantizer:197 def __init__(self, bits=4):198 self.bits = bits199 self.levels = 2 ** bits200 201 def quantize(self, tensor):202 norms = tensor.norm(dim=-1, keepdim=True).clamp(min=1e-8)203 unit = tensor / norms204 norm_min = norms.min()205 norm_max = norms.max()206 norm_scale = (norm_max - norm_min) / (self.levels - 1)207 q_norms = ((norms - norm_min) / norm_scale.clamp(min=1e-8)).round().clamp(0, self.levels - 1)208 val_min = unit.min()209 val_max = unit.max()210 val_scale = (val_max - val_min) / (self.levels - 1)211 q_unit = ((unit - val_min) / val_scale.clamp(min=1e-8)).round().clamp(0, self.levels - 1)212 return q_norms, q_unit, (norm_min, norm_scale, val_min, val_scale)213 214 def dequantize(self, q_norms, q_unit, params):215 norm_min, norm_scale, val_min, val_scale = params216 norms = q_norms * norm_scale + norm_min217 unit = q_unit * val_scale + val_min218 unit = unit / unit.norm(dim=-1, keepdim=True).clamp(min=1e-8)219 return unit * norms220 221 222class TurboQuantKVCache:223 def __init__(self, bits=4):224 self.quantizer = PolarQuantizer(bits=bits)225 self.k_cache = []226 self.v_cache = []227 228 def update(self, k_new, v_new):229 qk_norms, qk_unit, k_params = self.quantizer.quantize(k_new)230 qv_norms, qv_unit, v_params = self.quantizer.quantize(v_new)231 self.k_cache.append((qk_norms, qk_unit, k_params))232 self.v_cache.append((qv_norms, qv_unit, v_params))233 234 def get(self):235 ks = [self.quantizer.dequantize(*entry) for entry in self.k_cache]236 vs = [self.quantizer.dequantize(*entry) for entry in self.v_cache]237 return torch.cat(ks, dim=2), torch.cat(vs, dim=2)238 239 def clear(self):240 self.k_cache.clear()241 self.v_cache.clear()242 243 244class KVCache:245 def __init__(self, max_seq_len):246 self.max_seq_len = max_seq_len247 self.k_cache = None248 self.v_cache = None249 self.pos = 0250 251 def _ensure_allocated(self, k_new, v_new):252 B, H, _, D = k_new.shape253 needs_alloc = (254 self.k_cache is None255 or self.k_cache.shape[0] != B256 or self.k_cache.shape[1] != H257 or self.k_cache.shape[3] != D258 or self.k_cache.device != k_new.device259 or self.k_cache.dtype != k_new.dtype260 )261 if needs_alloc:262 self.k_cache = torch.empty(263 B, H, self.max_seq_len, D,264 device=k_new.device,265 dtype=k_new.dtype,266 )267 self.v_cache = torch.empty(268 B, H, self.max_seq_len, D,269 device=v_new.device,270 dtype=v_new.dtype,271 )272 self.pos = 0273 274 def update(self, k_new, v_new):275 self._ensure_allocated(k_new, v_new)276 T = k_new.size(2)277 if self.pos + T > self.max_seq_len:278 raise ValueError(f"KV cache length {self.pos + T} exceeds max_seq_len {self.max_seq_len}")279 self.k_cache[:, :, self.pos:self.pos + T, :].copy_(k_new)280 self.v_cache[:, :, self.pos:self.pos + T, :].copy_(v_new)281 self.pos += T282 283 def get(self):284 if self.k_cache is None:285 return None, None286 return self.k_cache[:, :, :self.pos, :], self.v_cache[:, :, :self.pos, :]287 288 def clear(self):289 self.pos = 0290 291 292# --- MTP: Multi-Token Prediction ---293 294class MTPHead(nn.Module):295 def __init__(self, config, future_idx):296 super().__init__()297 self.future_idx = future_idx298 n_embd = config["n_embd"]299 vocab_size = config["vocab_size"]300 self.logit_cap = config.get("logit_cap", 0)301 self.use_chunked_loss = config.get("use_chunked_loss", False)302 self.loss_chunk_size = config.get("loss_chunk_size", 2048)303 self.proj = nn.Linear(n_embd, n_embd)304 self.ln = nn.LayerNorm(n_embd)305 self.lm_head = nn.Linear(n_embd, vocab_size, bias=False)306 307 def forward(self, hidden, targets=None):308 if targets is not None and self.use_chunked_loss:309 shift = self.future_idx310 if targets.size(1) <= shift:311 return None, None312 # Project only the positions with a future target, then reduce in chunks.313 h = self.ln(self.proj(hidden[:, :-shift]))314 loss = chunked_cross_entropy(315 h, self.lm_head.weight, targets[:, shift:], self.logit_cap, self.loss_chunk_size316 )317 return None, loss318 319 h = self.ln(self.proj(hidden))320 logits = soft_cap(self.lm_head(h), self.logit_cap)321 loss = None322 if targets is not None:323 shift = self.future_idx324 if targets.size(1) > shift:325 logits_shifted = logits[:, :-shift].contiguous()326 targets_shifted = targets[:, shift:].contiguous()327 loss = F.cross_entropy(328 logits_shifted.view(-1, logits_shifted.size(-1)),329 targets_shifted.view(-1),330 ignore_index=-1,331 )332 return logits, loss333 334 335# --- RoPE: Rotary Position Embeddings ---336 337class RotaryEmbedding(nn.Module):338 def __init__(self, dim, max_seq_len=4096, base=10000.0):339 super().__init__()340 inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))341 self.register_buffer("inv_freq", inv_freq)342 self._build_cache(max_seq_len)343 344 def _build_cache(self, seq_len):345 t = torch.arange(seq_len, dtype=self.inv_freq.dtype)346 freqs = torch.outer(t, self.inv_freq)347 emb = torch.cat([freqs, freqs], dim=-1)348 self.register_buffer("cos_cached", emb.cos(), persistent=False)349 self.register_buffer("sin_cached", emb.sin(), persistent=False)350 351 def forward(self, seq_len):352 return self.cos_cached[:seq_len], self.sin_cached[:seq_len]353 354 355def rotate_half(x):356 x1, x2 = x.chunk(2, dim=-1)357 return torch.cat([-x2, x1], dim=-1)358 359 360def apply_rope(q, k, cos, sin):361 cos = cos.unsqueeze(0).unsqueeze(0)362 sin = sin.unsqueeze(0).unsqueeze(0)363 q = q * cos + rotate_half(q) * sin364 k = k * cos + rotate_half(k) * sin365 return q, k366 367 368# --- SwiGLU MLP ---369 370class SwiGLU(nn.Module):371 def __init__(self, config):372 super().__init__()373 n_embd = config["n_embd"]374 hidden = int(4 * n_embd * 2 / 3)375 hidden = ((hidden + 63) // 64) * 64376 use_bitnet = config.get("use_bitnet", False)377 use_fast_bitnet = config.get("use_fast_bitnet", False)378 self.gate = make_linear(n_embd, hidden, bias=False, use_bitnet=use_bitnet, use_fast_bitnet=use_fast_bitnet)379 self.up = make_linear(n_embd, hidden, bias=False, use_bitnet=use_bitnet, use_fast_bitnet=use_fast_bitnet)380 self.down = make_linear(hidden, n_embd, bias=False, use_bitnet=use_bitnet, use_fast_bitnet=use_fast_bitnet)381 382 def forward(self, x):383 return self.down(F.silu(self.gate(x)) * self.up(x))384 385 386class ReLU2MLP(nn.Module):387 """Ungated MLP with squared-ReLU activation (modded-nanoGPT). Simpler and a bit388 faster than SwiGLU; competitive quality at small scale."""389 390 def __init__(self, config):391 super().__init__()392 n_embd = config["n_embd"]393 hidden = 4 * n_embd394 use_bitnet = config.get("use_bitnet", False)395 use_fast_bitnet = config.get("use_fast_bitnet", False)396 self.fc = make_linear(n_embd, hidden, bias=False, use_bitnet=use_bitnet, use_fast_bitnet=use_fast_bitnet)397 self.proj = make_linear(hidden, n_embd, bias=False, use_bitnet=use_bitnet, use_fast_bitnet=use_fast_bitnet)398 399 def forward(self, x):400 return self.proj(F.relu(self.fc(x)).square())401 402 403# --- Core model ---404 405def make_norm(n_embd, use_rmsnorm=False):406 if use_rmsnorm:407 return nn.RMSNorm(n_embd)408 return nn.LayerNorm(n_embd)409 410 411class CausalSelfAttention(nn.Module):412 def __init__(self, config):413 super().__init__()414 self.n_head = config["n_head"]415 self.n_embd = config["n_embd"]416 self.n_kv_head = config.get("n_kv_head", self.n_head)417 if self.n_embd % self.n_head != 0:418 raise ValueError(f"n_embd ({self.n_embd}) must be divisible by n_head ({self.n_head})")419 if self.n_head % self.n_kv_head != 0:420 raise ValueError(f"n_head ({self.n_head}) must be divisible by n_kv_head ({self.n_kv_head})")421 self.head_dim = self.n_embd // self.n_head422 self.use_rope = config.get("use_rope", False)423 self.use_qk_norm = config.get("use_qk_norm", False)424 use_bitnet = config.get("use_bitnet", False)425 use_fast_bitnet = config.get("use_fast_bitnet", False)426 427 self.q_proj = make_linear(self.n_embd, self.n_head * self.head_dim, use_bitnet=use_bitnet, use_fast_bitnet=use_fast_bitnet)428 self.k_proj = make_linear(self.n_embd, self.n_kv_head * self.head_dim, use_bitnet=use_bitnet, use_fast_bitnet=use_fast_bitnet)429 self.v_proj = make_linear(self.n_embd, self.n_kv_head * self.head_dim, use_bitnet=use_bitnet, use_fast_bitnet=use_fast_bitnet)430 self.proj = make_linear(self.n_embd, self.n_embd, use_bitnet=use_bitnet, use_fast_bitnet=use_fast_bitnet)431 432 # QK-Norm (modded-nanoGPT): RMSNorm Q and K over the head dim before attention.433 if self.use_qk_norm:434 self.q_norm = nn.RMSNorm(self.head_dim)435 self.k_norm = nn.RMSNorm(self.head_dim)436 437 if self.use_rope:438 self.rope = RotaryEmbedding(self.head_dim, max_seq_len=config.get("block_size", 512))439 440 def forward(self, x, kv_cache=None, pos_offset=0):441 B, T, C = x.shape442 q = self.q_proj(x).view(B, T, self.n_head, self.head_dim).transpose(1, 2)443 k = self.k_proj(x).view(B, T, self.n_kv_head, self.head_dim).transpose(1, 2)444 v = self.v_proj(x).view(B, T, self.n_kv_head, self.head_dim).transpose(1, 2)445 446 if self.use_qk_norm:447 q = self.q_norm(q)448 k = self.k_norm(k)449 450 if self.use_rope:451 cos, sin = self.rope(pos_offset + T)452 cos, sin = cos[pos_offset:pos_offset + T], sin[pos_offset:pos_offset + T]453 q, k = apply_rope(q, k, cos, sin)454 455 if self.n_kv_head < self.n_head:456 repeats = self.n_head // self.n_kv_head457 k = k.repeat_interleave(repeats, dim=1)458 v = v.repeat_interleave(repeats, dim=1)459 460 if kv_cache is not None:461 kv_cache.update(k, v)462 k, v = kv_cache.get()463 464 use_causal = (T > 1)465 out = F.scaled_dot_product_attention(q, k, v, is_causal=use_causal)466 out = out.transpose(1, 2).reshape(B, T, C)467 return self.proj(out)468 469 470class MLP(nn.Module):471 def __init__(self, config):472 super().__init__()473 use_bitnet = config.get("use_bitnet", False)474 use_fast_bitnet = config.get("use_fast_bitnet", False)475 self.fc = make_linear(config["n_embd"], 4 * config["n_embd"], use_bitnet=use_bitnet, use_fast_bitnet=use_fast_bitnet)476 self.proj = make_linear(4 * config["n_embd"], config["n_embd"], use_bitnet=use_bitnet, use_fast_bitnet=use_fast_bitnet)477 478 def forward(self, x):479 return self.proj(F.gelu(self.fc(x)))480 481 482class Block(nn.Module):483 def __init__(self, config, layer_idx=0):484 super().__init__()485 self.use_mhc = config.get("use_mhc", False)486 use_rmsnorm = config.get("use_rmsnorm", False)487 self.ln1 = make_norm(config["n_embd"], use_rmsnorm)488 self.attn = CausalSelfAttention(config)489 self.ln2 = make_norm(config["n_embd"], use_rmsnorm)490 if config.get("use_relu2", False):491 self.mlp = ReLU2MLP(config)492 elif config.get("use_swiglu", False):493 self.mlp = SwiGLU(config)494 else:495 self.mlp = MLP(config)496 if self.use_mhc:497 n_streams = config.get("mhc_streams", 4)498 self.mhc_attn = MHCResidual(n_streams)499 self.mhc_mlp = MHCResidual(n_streams)500 501 def forward(self, x, streams=None, kv_cache=None, pos_offset=0):502 if self.use_mhc and streams is not None:503 inp = streams[:, 0]504 attn_out = self.attn(self.ln1(inp), kv_cache=kv_cache, pos_offset=pos_offset)505 streams = self.mhc_attn(streams, attn_out)506 mlp_inp = streams[:, 0]507 mlp_out = self.mlp(self.ln2(mlp_inp))508 streams = self.mhc_mlp(streams, mlp_out)509 return streams510 else:511 x = x + self.attn(self.ln1(x), kv_cache=kv_cache, pos_offset=pos_offset)512 x = x + self.mlp(self.ln2(x))513 return x514 515 516class GPT(nn.Module):517 def __init__(self, config):518 super().__init__()519 self.config = config520 self.use_mhc = config.get("use_mhc", False)521 self.use_mtp = config.get("use_mtp", False)522 self.use_rope = config.get("use_rope", False)523 self.mtp_heads_n = config.get("mtp_heads", 4)524 self.mtp_weight = config.get("mtp_weight", 0.1)525 self.use_turboquant = config.get("use_turboquant", False)526 self.turboquant_bits = config.get("turboquant_bits", 4)527 self.use_activation_checkpointing = config.get("use_activation_checkpointing", False)528 self.logit_cap = config.get("logit_cap", 0)529 self.use_chunked_loss = config.get("use_chunked_loss", False)530 self.loss_chunk_size = config.get("loss_chunk_size", 2048)531 use_rmsnorm = config.get("use_rmsnorm", False)532 533 self.tok_emb = nn.Embedding(config["vocab_size"], config["n_embd"])534 if not self.use_rope:535 self.pos_emb = nn.Embedding(config["block_size"], config["n_embd"])536 self.blocks = nn.ModuleList([Block(config, i) for i in range(config["n_layer"])])537 self.ln_f = make_norm(config["n_embd"], use_rmsnorm)538 self.lm_head = nn.Linear(config["n_embd"], config["vocab_size"], bias=False)539 self.tok_emb.weight = self.lm_head.weight540 541 if self.use_mhc:542 n_streams = config.get("mhc_streams", 4)543 self.mhc_expand = MHCExpand(n_streams, config["n_embd"])544 self.mhc_collapse = MHCCollapse(n_streams, config["n_embd"])545 546 if self.use_mtp:547 self.mtp_heads = nn.ModuleList([548 MTPHead(config, future_idx=i + 1) for i in range(self.mtp_heads_n)549 ])550 if config.get("tie_mtp_lm_head", True):551 for head in self.mtp_heads:552 head.lm_head.weight = self.lm_head.weight553 554 self.apply(self._init_weights)555 556 # Zero-init the output projection of each block (attention out-proj + MLP557 # down-proj), muP-style (modded-nanoGPT / nanochat). Each block starts as a558 # near-identity residual and learns to contribute, which helps convergence.559 if config.get("use_zero_init", False):560 for block in self.blocks:561 torch.nn.init.zeros_(block.attn.proj.weight)562 mlp_out = getattr(block.mlp, "down", None)563 if mlp_out is None:564 mlp_out = block.mlp.proj # MLP / ReLU2MLP name the out-proj "proj"565 torch.nn.init.zeros_(mlp_out.weight)566 567 def _init_weights(self, module):568 if isinstance(module, (nn.Linear, BitLinear)):569 torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)570 if module.bias is not None:571 torch.nn.init.zeros_(module.bias)572 elif isinstance(module, nn.Embedding):573 torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)574 575 def _compute_hidden(self, idx):576 B, T = idx.shape577 if T > self.config["block_size"]:578 raise ValueError(f"Input length {T} exceeds block_size {self.config['block_size']}")579 x = self.tok_emb(idx)580 if not self.use_rope:581 pos = torch.arange(T, device=idx.device)582 x = x + self.pos_emb(pos)583 584 if self.use_mhc:585 streams = self.mhc_expand(x)586 for block in self.blocks:587 if self.training and self.use_activation_checkpointing:588 streams = checkpoint(lambda s, b=block: b(x, streams=s), streams, use_reentrant=False)589 else:590 streams = block(x, streams=streams)591 x = self.mhc_collapse(streams)592 else:593 for block in self.blocks:594 if self.training and self.use_activation_checkpointing:595 x = checkpoint(block, x, use_reentrant=False)596 else:597 x = block(x)598 599 return self.ln_f(x)600 601 def forward(self, idx, targets=None, return_hidden=False):602 hidden = self._compute_hidden(idx)603 loss = None604 # Chunked loss avoids materializing the full [N, vocab] logits during training.605 # It can't return logits, so fall back to the dense path when logits are needed.606 if targets is not None and self.use_chunked_loss and not return_hidden:607 logits = None608 loss = chunked_cross_entropy(609 hidden, self.lm_head.weight, targets, self.logit_cap, self.loss_chunk_size610 )611 else:612 logits = soft_cap(self.lm_head(hidden), self.logit_cap)613 if targets is not None:614 loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)615 if targets is not None and self.use_mtp:616 for head in self.mtp_heads:617 _, mtp_loss = head(hidden, targets)618 if mtp_loss is not None:619 loss = loss + self.mtp_weight * mtp_loss620 if return_hidden:621 return logits, loss, hidden622 return logits, loss623 624 def _forward_inference(self, x, kv_caches, pos_offset=0, return_hidden=False):625 if self.use_mhc:626 streams = self.mhc_expand(x)627 for block, cache in zip(self.blocks, kv_caches or [None] * len(self.blocks)):628 streams = block(x, streams=streams, kv_cache=cache, pos_offset=pos_offset)629 x = self.mhc_collapse(streams)630 else:631 for block, cache in zip(self.blocks, kv_caches or [None] * len(self.blocks)):632 x = block(x, kv_cache=cache, pos_offset=pos_offset)633 hidden = self.ln_f(x)634 logits = soft_cap(self.lm_head(hidden), self.logit_cap)635 if return_hidden:636 return logits, hidden637 return logits638 639 def _embed(self, tokens, pos_offset=0):640 x = self.tok_emb(tokens)641 if not self.use_rope:642 T = tokens.shape[1]643 pos = torch.arange(pos_offset, pos_offset + T, device=tokens.device)644 x = x + self.pos_emb(pos)645 return x646 647 def _filter_logits(self, logits, top_k=None, top_p=None, min_p=None):648 if top_k is not None and top_k > 0:649 k = min(top_k, logits.size(-1))650 values, _ = torch.topk(logits, k)651 logits = logits.masked_fill(logits < values[:, [-1]], -float("inf"))652 653 if min_p is not None and min_p > 0:654 probs = F.softmax(logits, dim=-1)655 max_probs = probs.max(dim=-1, keepdim=True).values656 remove = probs < (min_p * max_probs)657 top_token = logits.argmax(dim=-1, keepdim=True)658 remove.scatter_(dim=-1, index=top_token, value=False)659 logits = logits.masked_fill(remove, -float("inf"))660 661 if top_p is not None and 0 < top_p < 1.0:662 sorted_logits, sorted_idx = torch.sort(logits, descending=True, dim=-1)663 sorted_probs = F.softmax(sorted_logits, dim=-1)664 cumulative_probs = sorted_probs.cumsum(dim=-1)665 sorted_remove = cumulative_probs > top_p666 sorted_remove[..., 1:] = sorted_remove[..., :-1].clone()667 sorted_remove[..., 0] = False668 remove = torch.zeros_like(logits, dtype=torch.bool)669 remove.scatter_(dim=-1, index=sorted_idx, src=sorted_remove)670 logits = logits.masked_fill(remove, -float("inf"))671 672 return logits673 674 def _distribution(self, logits, temperature=0.8, top_k=40, top_p=None, min_p=None):675 if temperature <= 0:676 token = logits.argmax(dim=-1, keepdim=True)677 probs = torch.zeros_like(logits)678 probs.scatter_(1, token, 1.0)679 return token, probs680 logits = self._filter_logits(logits / temperature, top_k=top_k, top_p=top_p, min_p=min_p)681 probs = F.softmax(logits, dim=-1)682 token = torch.multinomial(probs, num_samples=1)683 return token, probs684 685 def _make_kv_caches(self, use_turboquant, use_kv_cache=True):686 if not use_kv_cache:687 return None688 if use_turboquant:689 return [TurboQuantKVCache(bits=self.turboquant_bits) for _ in self.blocks]690 return [KVCache(self.config["block_size"]) for _ in self.blocks]691 692 def _trim_or_seed_prompt(self, idx):693 block_size = self.config["block_size"]694 if idx.shape[1] == 0:695 eos_id = 1696 idx = torch.tensor([[eos_id]], dtype=idx.dtype, device=idx.device)697 return idx[:, -block_size:]698 699 def _prefill_generation(self, idx, use_turboquant=False, use_kv_cache=True):700 kv_caches = self._make_kv_caches(use_turboquant, use_kv_cache=use_kv_cache)701 seq_len = idx.shape[1]702 x = self._embed(idx)703 logits, hidden = self._forward_inference(x, kv_caches, pos_offset=0, return_hidden=True)704 return logits, hidden[:, -1:, :], kv_caches, seq_len705 706 def _advance_generation_state(self, idx, idx_next, kv_caches, seq_len, use_turboquant):707 block_size = self.config["block_size"]708 if kv_caches is not None and seq_len < block_size:709 x = self._embed(idx_next, pos_offset=seq_len)710 logits, hidden = self._forward_inference(x, kv_caches, pos_offset=seq_len, return_hidden=True)711 return logits, hidden[:, -1:, :], kv_caches, seq_len + 1712 713 use_kv_cache = kv_caches is not None714 if kv_caches:715 for cache in kv_caches:716 cache.clear()717 idx_cond = idx[:, -block_size:]718 logits, hidden, kv_caches, seq_len = self._prefill_generation(719 idx_cond,720 use_turboquant=use_turboquant,721 use_kv_cache=use_kv_cache,722 )723 return logits, hidden, kv_caches, seq_len724 725 def _generate_autoregressive(726 self,727 idx,728 max_new_tokens,729 temperature=0.8,730 top_k=40,731 top_p=None,732 min_p=None,733 use_turboquant=None,734 use_kv_cache=True,735 ):736 idx = self._trim_or_seed_prompt(idx)737 use_turboquant = self.use_turboquant if use_turboquant is None else use_turboquant738 logits, last_hidden, kv_caches, seq_len = self._prefill_generation(739 idx,740 use_turboquant=use_turboquant,741 use_kv_cache=use_kv_cache,742 )743 744 for i in range(max_new_tokens):745 idx_next, _ = self._distribution(746 logits[:, -1, :],747 temperature=temperature,748 top_k=top_k,749 top_p=top_p,750 min_p=min_p,751 )752 idx = torch.cat([idx, idx_next], dim=1)753 754 if i < max_new_tokens - 1:755 logits, last_hidden, kv_caches, seq_len = self._advance_generation_state(756 idx, idx_next, kv_caches, seq_len, use_turboquant757 )758 759 return idx760 761 def _mtp_draft(self, last_hidden, n_tokens, temperature=0.8, top_k=40, top_p=None, min_p=None):762 draft_tokens = []763 draft_probs = []764 for head in self.mtp_heads[:n_tokens]:765 draft_logits, _ = head(last_hidden)766 token, probs = self._distribution(767 draft_logits[:, -1, :],768 temperature=temperature,769 top_k=top_k,770 top_p=top_p,771 min_p=min_p,772 )773 draft_tokens.append(token)774 draft_probs.append(probs)775 return draft_tokens, draft_probs776 777 def _resample_on_reject(self, target_token, p_probs, q_probs, temperature):778 if temperature <= 0:779 return target_token780 residual = (p_probs - q_probs).clamp(min=0)781 denom = residual.sum(dim=-1, keepdim=True)782 if denom.item() <= 1e-12:783 return target_token784 return torch.multinomial(residual / denom, num_samples=1)785 786 def _mtp_speculative_generate(787 self,788 idx,789 max_new_tokens,790 temperature=0.8,791 top_k=40,792 top_p=None,793 min_p=None,794 speculate_tokens=None,795 use_turboquant=None,796 use_kv_cache=True,797 ):798 use_turboquant = self.use_turboquant if use_turboquant is None else use_turboquant799 # Batched verification needs a single sequence, MTP draft heads, and the800 # plain (rollback-able) KV cache. TurboQuant's cache cannot be rolled back801 # token-by-token, so fall back to autoregressive there.802 if not self.use_mtp or idx.size(0) != 1 or not use_kv_cache or use_turboquant:803 return self._generate_autoregressive(804 idx,805 max_new_tokens,806 temperature=temperature,807 top_k=top_k,808 top_p=top_p,809 min_p=min_p,810 use_turboquant=use_turboquant,811 use_kv_cache=use_kv_cache,812 )813 814 idx = self._trim_or_seed_prompt(idx)815 block_size = self.config["block_size"]816 draft_width = speculate_tokens or self.mtp_heads_n817 draft_width = max(1, min(draft_width, self.mtp_heads_n))818 819 logits, last_hidden, kv_caches, seq_len = self._prefill_generation(820 idx, use_turboquant=False, use_kv_cache=True821 )822 # p0 = main-model logits for the next token (verifies the first draft).823 p0_logits = logits[:, -1, :]824 generated = 0825 826 while generated < max_new_tokens:827 remaining = max_new_tokens - generated828 n_draft = min(draft_width, remaining)829 830 # No room left in the cache window: take one plain step (this slides the831 # window via re-prefill inside _advance_generation_state) and continue.832 if seq_len + n_draft > block_size:833 idx_next, _ = self._distribution(p0_logits, temperature, top_k, top_p, min_p)834 idx = torch.cat([idx, idx_next], dim=1)835 generated += 1836 if generated < max_new_tokens:837 logits, last_hidden, kv_caches, seq_len = self._advance_generation_state(838 idx, idx_next, kv_caches, seq_len, False839 )840 p0_logits = logits[:, -1, :]841 continue842 843 # 1. Draft n tokens cheaply from the MTP heads (no main-model forward).844 draft_tokens, draft_probs = self._mtp_draft(845 last_hidden, n_draft, temperature=temperature, top_k=top_k, top_p=top_p, min_p=min_p846 )847 draft_seq = torch.cat(draft_tokens, dim=1)848 849 # 2. Verify ALL drafts in a SINGLE main-model forward pass.850 x = self._embed(draft_seq, pos_offset=seq_len)851 v_logits, v_hidden = self._forward_inference(852 x, kv_caches, pos_offset=seq_len, return_hidden=True853 )854 855 # 3. Walk the drafts left-to-right; draft j is checked against the main856 # distribution at the previous position (p0 for j=0, else v_logits[j-1]).857 accepted = 0858 reject_token = None859 for j in range(n_draft):860 target_logits = p0_logits if j == 0 else v_logits[:, j - 1, :]861 target_token, p_probs = self._distribution(862 target_logits, temperature, top_k, top_p, min_p863 )864 if temperature <= 0:865 accept = torch.equal(draft_tokens[j], target_token)866 else:867 proposed = draft_tokens[j].item()868 p = p_probs[0, proposed]869 q = draft_probs[j][0, proposed].clamp(min=1e-12)870 accept = torch.rand((), device=idx.device) <= torch.minimum(torch.ones_like(p), p / q)871 if accept:872 accepted += 1873 else:874 reject_token = self._resample_on_reject(875 target_token, p_probs, draft_probs[j], temperature876 )877 break878 879 if accepted == n_draft:880 # Every draft matched the main model: commit them all. The cache881 # already holds them and v_hidden/v_logits give the next draft state882 # for free (no extra forward, no separate bonus token needed).883 idx = torch.cat([idx, draft_seq], dim=1)884 generated += n_draft885 seq_len += n_draft886 last_hidden = v_hidden[:, -1:, :]887 p0_logits = v_logits[:, -1, :]888 else:889 # Commit the accepted prefix plus the corrected token, then roll the890 # cache back to drop the rejected drafts' (now stale) KV entries.891 commit = torch.cat(draft_tokens[:accepted] + [reject_token], dim=1)892 idx = torch.cat([idx, commit], dim=1)893 generated += accepted + 1894 for cache in kv_caches:895 cache.pos = seq_len + accepted896 seq_len += accepted897 if generated < max_new_tokens:898 # reject_token's KV/hidden are not cached yet; one short forward rebases.899 logits, last_hidden, kv_caches, seq_len = self._advance_generation_state(900 idx, reject_token, kv_caches, seq_len, False901 )902 p0_logits = logits[:, -1, :]903 904 return idx905 906 def generate(907 self,908 idx,909 max_new_tokens,910 temperature=0.8,911 top_k=40,912 top_p=None,913 min_p=None,914 speculative=False,915 speculate_tokens=None,916 use_turboquant=None,917 use_kv_cache=True,918 ):919 if speculative:920 return self._mtp_speculative_generate(921 idx,922 max_new_tokens,923 temperature=temperature,924 top_k=top_k,925 top_p=top_p,926 min_p=min_p,927 speculate_tokens=speculate_tokens,928 use_turboquant=use_turboquant,929 use_kv_cache=use_kv_cache,930 )931 return self._generate_autoregressive(932 idx,933 max_new_tokens,934 temperature=temperature,935 top_k=top_k,936 top_p=top_p,937 min_p=min_p,938 use_turboquant=use_turboquant,939 use_kv_cache=use_kv_cache,940 )941 942 943# --- Configs ---944 945BASE_CONFIG = {946 "vocab_size": 16384,947 "block_size": 512,948 "n_embd": 512,949 "n_head": 8,950 "n_layer": 12,951}952 953# Individual techniques954MHC_CONFIG = {**BASE_CONFIG, "use_mhc": True, "mhc_streams": 4}955BITNET_CONFIG = {**BASE_CONFIG, "use_bitnet": True}956FAST_BITNET_CONFIG = {**BASE_CONFIG, "use_fast_bitnet": True}957MTP_CONFIG = {**BASE_CONFIG, "use_mtp": True, "mtp_heads": 4, "mtp_weight": 0.1}958ROPE_CONFIG = {**BASE_CONFIG, "use_rope": True}959GQA_CONFIG = {**BASE_CONFIG, "n_kv_head": 2}960SWIGLU_CONFIG = {**BASE_CONFIG, "use_swiglu": True}961RMSNORM_CONFIG = {**BASE_CONFIG, "use_rmsnorm": True}962TURBOQUANT_CONFIG = {**BASE_CONFIG, "use_turboquant": True, "turboquant_bits": 4}963 964# Combinations965MHC_BITNET_CONFIG = {**BASE_CONFIG, "use_mhc": True, "mhc_streams": 4, "use_bitnet": True}966MHC_MTP_CONFIG = {**BASE_CONFIG, "use_mhc": True, "mhc_streams": 4, "use_mtp": True, "mtp_heads": 4, "mtp_weight": 0.1}967 968# Modern LLaMA-style (RoPE + GQA + SwiGLU + RMSNorm)969MODERN_CONFIG = {**BASE_CONFIG, "use_rope": True, "n_kv_head": 2, "use_swiglu": True, "use_rmsnorm": True}970 971# Everything972ALL_CONFIG = {973 **BASE_CONFIG,974 "use_mhc": True, "mhc_streams": 4,975 "use_bitnet": True,976 "use_mtp": True, "mtp_heads": 4, "mtp_weight": 0.1,977 "use_rope": True, "n_kv_head": 2,978 "use_swiglu": True, "use_rmsnorm": True,979 "use_turboquant": True, "turboquant_bits": 4,980}981 982RECOMMENDED_CONFIG = {983 **BASE_CONFIG,984 "use_rope": True, "n_kv_head": 2,985 "use_swiglu": True, "use_rmsnorm": True,986 "use_mtp": True, "mtp_heads": 4, "mtp_weight": 0.1,987}988 989FAST_2060_CONFIG = {990 **BASE_CONFIG,991 "block_size": 256,992 "n_embd": 384,993 "n_head": 6,994 "n_layer": 8,995 "use_rope": True,996 "n_kv_head": 2,997 "use_swiglu": True,998 "use_rmsnorm": True,999}1000 1001FAST_2060_MTP_CONFIG = {1002 **FAST_2060_CONFIG,1003 "use_mtp": True,1004 "mtp_heads": 2,1005 "mtp_weight": 0.1,1006 "tie_mtp_lm_head": True,1007}1008 1009FAST_2060_MTP_FBITNET_CONFIG = {1010 **FAST_2060_MTP_CONFIG,1011 "use_fast_bitnet": True,1012}1013 1014# modded-nanoGPT-style recipe. QK-Norm helps under any optimizer; ReLU2 and1015# logit_cap only pay off paired with Muon's higher LR. Train with --optimizer muon.1016FAST_2060_MODDED_CONFIG = {1017 **FAST_2060_MTP_CONFIG,1018 "use_swiglu": False, # superseded by ReLU2 below1019 "use_relu2": True,1020 "use_qk_norm": True,1021 "logit_cap": 15.0,1022 "use_zero_init": True, # measured: val 2.13 -> 2.04 at equal steps, free1023}1024 1025# Same modded recipe but WITHOUT MTP (built on FAST_2060_CONFIG, not the _mtp one).1026# This is exactly the config that won the convergence A/B (val 2.13). No MTP means a1027# cleaner pure-CE loss number and faster steps, but no speculative-decoding heads.1028FAST_2060_MODDED_NOMTP_CONFIG = {1029 **FAST_2060_CONFIG,1030 "use_swiglu": False,1031 "use_relu2": True,1032 "use_qk_norm": True,1033 "logit_cap": 15.0,1034 "use_zero_init": True,1035}1036 1037FAST_2060_MTP_TURBO_CONFIG = {1038 **FAST_2060_MTP_CONFIG,1039 "use_turboquant": True,1040 "turboquant_bits": 4,1041}1042 1043TINY_FAST_CONFIG = {1044 **BASE_CONFIG,1045 "block_size": 256,1046 "n_embd": 256,1047 "n_head": 4,1048 "n_layer": 6,1049 "use_rope": True,1050 "n_kv_head": 2,1051 "use_swiglu": True,1052 "use_rmsnorm": True,1053}1054 1055LOW_MEMORY_2060_CONFIG = {1056 **FAST_2060_CONFIG,1057 "use_activation_checkpointing": True,1058}1059 1060CONFIGS = {1061 "base": BASE_CONFIG,1062 "mhc": MHC_CONFIG,1063 "bitnet": BITNET_CONFIG,1064 "mtp": MTP_CONFIG,1065 "rope": ROPE_CONFIG,1066 "gqa": GQA_CONFIG,1067 "swiglu": SWIGLU_CONFIG,1068 "rmsnorm": RMSNORM_CONFIG,1069 "turboquant": TURBOQUANT_CONFIG,1070 "mhc_bitnet": MHC_BITNET_CONFIG,1071 "mhc_mtp": MHC_MTP_CONFIG,1072 "modern": MODERN_CONFIG,1073 "all": ALL_CONFIG,1074 "recommended": RECOMMENDED_CONFIG,1075 "fast_2060": FAST_2060_CONFIG,1076 "fast_2060_mtp": FAST_2060_MTP_CONFIG,1077 "fast_2060_mtp_fbitnet": FAST_2060_MTP_FBITNET_CONFIG,1078 "fast_2060_modded": FAST_2060_MODDED_CONFIG,1079 "fast_2060_modded_nomtp": FAST_2060_MODDED_NOMTP_CONFIG,1080 "fast_2060_mtp_turbo": FAST_2060_MTP_TURBO_CONFIG,1081 "tiny_fast": TINY_FAST_CONFIG,1082 "low_memory_2060": LOW_MEMORY_2060_CONFIG,1083}1084 1085 1086def get_model_config(name="fast_2060", **overrides):1087 if name not in CONFIGS:1088 available = ", ".join(sorted(CONFIGS))1089 raise ValueError(f"Unknown config '{name}'. Available configs: {available}")1090 return {**CONFIGS[name], **{k: v for k, v in overrides.items() if v is not None}}1091 1092 1093MODEL_CONFIG = RECOMMENDED_CONFIG1094 1095if __name__ == "__main__":1096 configs = CONFIGS1097 for name, cfg in configs.items():1098 model = GPT(cfg)1099 n_params = sum(p.numel() for p in model.parameters())1100 x = torch.randint(0, cfg["vocab_size"], (2, 64))1101 logits, loss = model(x, x)1102 print(f"{name:<12} | {n_params:>12,} params ({n_params/1e6:.1f}M) | loss: {loss.item():.2f}")1103