|
|
|
|
|
|
|
|
|
|
|
|
|
|
| """
|
| Referencyjny ZWYKLY transformer (byte-level nanoGPT-style) ~25M — apples-to-apples vs BDH-25M.
|
| Ta sama data (train.bin/val.bin uint8), ten sam scale (~25M), ten sam byte-level (vocab256).
|
| Rozni sie TYLKO architektura (standard causal transformer vs BDH fast-weights) -> czysta referencja.
|
|
|
| CLI mirror train_bdh.py. GPU ROCm/CUDA bf16, cosine+warmup+clip, ckpt/resume, logging.
|
| Autor: Hart (N-02).
|
|
|
| Smoke throughput (bez danych PII, syntetyczny bufor): --synthetic --steps 60
|
| Realny: --data-dir . --run-id gpt25m_run1 --steps 30000
|
| """
|
| import argparse
|
| import json
|
| import math
|
| import os
|
| import time
|
| import queue
|
| import threading
|
| from contextlib import nullcontext
|
|
|
| import numpy as np
|
| import torch
|
| import torch.nn as nn
|
| import torch.nn.functional as F
|
|
|
|
|
| def get_args():
|
| p = argparse.ArgumentParser()
|
| p.add_argument("--data-dir", default=".")
|
| p.add_argument("--out-dir", default=None)
|
| p.add_argument("--run-id", default="gpt25m_run1")
|
| p.add_argument("--steps", type=int, default=30000)
|
| p.add_argument("--batch", type=int, default=32)
|
| p.add_argument("--block", type=int, default=256)
|
| p.add_argument("--n-layer", type=int, default=8)
|
| p.add_argument("--n-embd", type=int, default=512)
|
| p.add_argument("--n-head", type=int, default=8)
|
| p.add_argument("--lr", type=float, default=6e-4)
|
| p.add_argument("--min-lr", type=float, default=6e-5)
|
| p.add_argument("--warmup", type=int, default=200)
|
| p.add_argument("--wd", type=float, default=0.1)
|
| p.add_argument("--grad-clip", type=float, default=1.0)
|
| p.add_argument("--log-every", type=int, default=50)
|
| p.add_argument("--eval-every", type=int, default=500)
|
| p.add_argument("--eval-iters", type=int, default=50)
|
| p.add_argument("--ckpt-every", type=int, default=1000)
|
| p.add_argument("--seed", type=int, default=1337)
|
| p.add_argument("--resume", action="store_true")
|
| p.add_argument("--synthetic", action="store_true", help="smoke throughput na losowym uint8 (bez danych)")
|
| p.add_argument("--vocab", type=int, default=256)
|
| p.add_argument("--dtype", default="uint8", help="bin dtype: uint8 (byte) | uint16 (BPE)")
|
| p.add_argument("--optimizer", choices=["adamw", "muon"], default="adamw",
|
| help="adamw (default, backward-compat) | muon (Newton-Schulz ortho dla 2D-weights + AdamW dla reszty)")
|
| p.add_argument("--muon-lr", type=float, default=0.02,
|
| help="peak LR dla Muon (macierzowe params); AdamW-aux uzywa --lr. Muon skalowany ta sama cosine-schedule co AdamW przez lr_mult=muon_lr/lr")
|
| return p.parse_args()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| def zeropower_via_newtonschulz5(G, steps=5):
|
| """Ortogonalizacja macierzy G przez quintic Newton-Schulz (bf16). Zwraca
|
| macierz ~ U V^T z SVD(G)=U S V^T. Wspolczynniki (a,b,c) z impl. Kellera."""
|
| assert G.ndim == 2
|
| a, b, c = (3.4445, -4.7750, 2.0315)
|
| X = G.bfloat16()
|
| transposed = G.size(0) > G.size(1)
|
| if transposed:
|
| X = X.T
|
| X = X / (X.norm() + 1e-7)
|
| for _ in range(steps):
|
| A = X @ X.T
|
| B = b * A + c * (A @ A)
|
| X = a * X + B @ X
|
| if transposed:
|
| X = X.T
|
| return X
|
|
|
|
|
| class Muon(torch.optim.Optimizer):
|
| """Momentum-SGD z ortogonalizowanym update. weight_decay domyslnie 0 (Muon-params
|
| czysto; WD trzymamy na AdamW-aux). lr_mult pozwala petli lr-schedule skalowac
|
| Muon proporcjonalnie do AdamW."""
|
| def __init__(self, params, lr=0.02, lr_mult=1.0, momentum=0.95, nesterov=True,
|
| ns_steps=5, weight_decay=0.0):
|
| defaults = dict(lr=lr, lr_mult=lr_mult, momentum=momentum, nesterov=nesterov,
|
| ns_steps=ns_steps, weight_decay=weight_decay)
|
| super().__init__(params, defaults)
|
|
|
| @torch.no_grad()
|
| def step(self, closure=None):
|
| loss = None
|
| if closure is not None:
|
| with torch.enable_grad():
|
| loss = closure()
|
| for group in self.param_groups:
|
| lr = group["lr"]; momentum = group["momentum"]; wd = group["weight_decay"]
|
| for p in group["params"]:
|
| g = p.grad
|
| if g is None:
|
| continue
|
| if g.ndim > 2:
|
| g = g.reshape(g.size(0), -1)
|
| state = self.state[p]
|
| if "momentum_buffer" not in state:
|
| state["momentum_buffer"] = torch.zeros_like(g)
|
| buf = state["momentum_buffer"]
|
| buf.mul_(momentum).add_(g)
|
| g = g.add(buf, alpha=momentum) if group["nesterov"] else buf
|
| u = zeropower_via_newtonschulz5(g, steps=group["ns_steps"])
|
| if wd != 0:
|
| p.mul_(1 - lr * wd)
|
|
|
| scale = max(1.0, p.size(0) / p.size(1)) ** 0.5
|
| p.add_(u.reshape(p.shape).to(p.dtype), alpha=-lr * scale)
|
| return loss
|
|
|
|
|
| class MuonWithAuxAdam:
|
| """Kontener: Muon dla macierzowych ukrytych wag + AdamW dla reszty. Wystawia
|
| param_groups/step/zero_grad/state_dict tak, by petla treningowa dzialala bez zmian."""
|
| def __init__(self, muon, adamw):
|
| self.muon = muon
|
| self.adamw = adamw
|
|
|
| @property
|
| def param_groups(self):
|
| return self.muon.param_groups + self.adamw.param_groups
|
|
|
| @property
|
| def state(self):
|
| return {**self.muon.state, **self.adamw.state}
|
|
|
| def step(self, closure=None):
|
| self.muon.step()
|
| self.adamw.step()
|
|
|
| def zero_grad(self, set_to_none=True):
|
| self.muon.zero_grad(set_to_none=set_to_none)
|
| self.adamw.zero_grad(set_to_none=set_to_none)
|
|
|
| def state_dict(self):
|
| return {"muon": self.muon.state_dict(), "adamw": self.adamw.state_dict()}
|
|
|
| def load_state_dict(self, sd):
|
| self.muon.load_state_dict(sd["muon"])
|
| self.adamw.load_state_dict(sd["adamw"])
|
|
|
|
|
| def build_optimizer(a, model):
|
| """--optimizer adamw -> DOKLADNIE poprzedni AdamW (backward-compat).
|
| --optimizer muon -> Muon(2D-hidden) + AdamW(embeddingi/head/norm/bias)."""
|
| if a.optimizer == "adamw":
|
| return torch.optim.AdamW(model.parameters(), lr=a.lr, weight_decay=a.wd, betas=(0.9, 0.95))
|
| muon_params, adamw_params, seen = [], [], set()
|
| for name, p in model.named_parameters():
|
| if not p.requires_grad or id(p) in seen:
|
| continue
|
| seen.add(id(p))
|
| is_embed_or_head = name.startswith(("tok.", "pos.", "head."))
|
| if p.ndim >= 2 and not is_embed_or_head:
|
| muon_params.append(p)
|
| else:
|
| adamw_params.append(p)
|
| lr_mult = a.muon_lr / a.lr if a.lr > 0 else 1.0
|
| muon = Muon(muon_params, lr=a.muon_lr, lr_mult=lr_mult, weight_decay=0.0)
|
| adamw = torch.optim.AdamW(adamw_params, lr=a.lr, weight_decay=a.wd, betas=(0.9, 0.95))
|
| return MuonWithAuxAdam(muon, adamw)
|
|
|
|
|
| class Block(nn.Module):
|
| def __init__(self, d, nh, block):
|
| super().__init__()
|
| self.ln1 = nn.LayerNorm(d)
|
| self.ln2 = nn.LayerNorm(d)
|
| self.qkv = nn.Linear(d, 3 * d)
|
| self.proj = nn.Linear(d, d)
|
| self.mlp = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))
|
| self.nh = nh
|
| self.d = d
|
|
|
| def forward(self, x):
|
| B, T, D = x.size()
|
| h = self.ln1(x)
|
| q, k, v = self.qkv(h).split(self.d, dim=2)
|
| q = q.view(B, T, self.nh, D // self.nh).transpose(1, 2)
|
| k = k.view(B, T, self.nh, D // self.nh).transpose(1, 2)
|
| v = v.view(B, T, self.nh, D // self.nh).transpose(1, 2)
|
| y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
|
| y = y.transpose(1, 2).contiguous().view(B, T, D)
|
| x = x + self.proj(y)
|
| x = x + self.mlp(self.ln2(x))
|
| return x
|
|
|
|
|
| class GPT(nn.Module):
|
| def __init__(self, vocab, n_layer, n_embd, n_head, block):
|
| super().__init__()
|
| self.tok = nn.Embedding(vocab, n_embd)
|
| self.pos = nn.Embedding(block, n_embd)
|
| self.blocks = nn.ModuleList([Block(n_embd, n_head, block) for _ in range(n_layer)])
|
| self.lnf = nn.LayerNorm(n_embd)
|
| self.head = nn.Linear(n_embd, vocab, bias=False)
|
| self.head.weight = self.tok.weight
|
| self.block = block
|
| self.apply(self._init)
|
|
|
| def _init(self, m):
|
| if isinstance(m, nn.Linear):
|
| nn.init.normal_(m.weight, 0.0, 0.02)
|
| if m.bias is not None:
|
| nn.init.zeros_(m.bias)
|
| elif isinstance(m, nn.Embedding):
|
| nn.init.normal_(m.weight, 0.0, 0.02)
|
|
|
| def forward(self, idx, targets=None):
|
| B, T = idx.size()
|
| pos = torch.arange(T, device=idx.device)
|
| x = self.tok(idx) + self.pos(pos)[None]
|
| for b in self.blocks:
|
| x = b(x)
|
| logits = self.head(self.lnf(x))
|
| loss = None
|
| if targets is not None:
|
| loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
|
| return logits, loss
|
|
|
|
|
| def main():
|
| a = get_args()
|
| out_dir = a.out_dir or os.path.join(a.data_dir, "runs", a.run_id)
|
| os.makedirs(out_dir, exist_ok=True)
|
| log_path = os.path.join(out_dir, "train.log")
|
| metrics_path = os.path.join(out_dir, "metrics.jsonl")
|
| ckpt_path = os.path.join(out_dir, "ckpt.pt")
|
|
|
| def log(msg):
|
| line = f"[{time.strftime('%H:%M:%S')}] {msg}"
|
| print(line, flush=True)
|
| with open(log_path, "a", encoding="utf-8") as f:
|
| f.write(line + "\n")
|
|
|
| torch.manual_seed(a.seed)
|
| torch.backends.cuda.matmul.allow_tf32 = True
|
| torch.backends.cudnn.allow_tf32 = True
|
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| use_bf16 = device.type == "cuda" and torch.cuda.is_bf16_supported()
|
| ptdtype = torch.bfloat16 if use_bf16 else torch.float32
|
| ctx = torch.amp.autocast(device_type=device.type, dtype=ptdtype) if device.type == "cuda" else nullcontext()
|
| log(f"device={device} bf16={use_bf16} dev={torch.cuda.get_device_name(0) if device.type=='cuda' else 'cpu'}")
|
|
|
| if a.synthetic:
|
| rng = np.random.default_rng(a.seed)
|
| train_data = rng.integers(0, 256, size=8_000_000, dtype=np.uint8)
|
| val_data = train_data[:200_000]
|
| log("SYNTHETIC uint8 (smoke throughput, zero danych PII)")
|
| else:
|
| train_data = np.memmap(os.path.join(a.data_dir, "train.bin"), dtype=np.dtype(a.dtype), mode="r")
|
| val_data = np.memmap(os.path.join(a.data_dir, "val.bin"), dtype=np.dtype(a.dtype), mode="r")
|
| log(f"dane: train={len(train_data):,}B block={a.block} batch={a.batch} tok/step={a.block*a.batch:,}")
|
|
|
| def _make_batch_cpu(split, generator=None):
|
| """Wektoryzowane budowanie batcha: JEDEN numpy fancy-index zamiast
|
| python-loop per-item. sliding_window_view daje strided-view (N-block, block+1)
|
| BEZ kopiowania; windows[ix] materializuje tylko wybrane wiersze naraz.
|
| Zwraca (x,y) long CPU (pinned jesli cuda). Rozklad batchy IDENTYCZNY jak
|
| stary torch.stack-loop: x=data[i:i+block], y=data[i+1:i+1+block]."""
|
| data = train_data if split == "train" else val_data
|
| ix = torch.randint(len(data) - a.block - 1, (a.batch,), generator=generator)
|
|
|
| windows = np.lib.stride_tricks.sliding_window_view(data, a.block + 1)
|
| sel = windows[ix.numpy()]
|
| x = torch.from_numpy(sel[:, :-1].astype(np.int64))
|
| y = torch.from_numpy(sel[:, 1:].astype(np.int64))
|
| if device.type == "cuda":
|
| x = x.pin_memory(); y = y.pin_memory()
|
| return x, y
|
|
|
| def _to_device(x, y):
|
| if device.type == "cuda":
|
| return x.to(device, non_blocking=True), y.to(device, non_blocking=True)
|
| return x.to(device), y.to(device)
|
|
|
| def get_batch(split, generator=None):
|
| return _to_device(*_make_batch_cpu(split, generator))
|
|
|
| class Prefetcher:
|
| """Async double-buffer: 1 background-thread buduje NASTEPNY batch na CPU
|
| (pinned) podczas gdy GPU liczy biezacy. queue depth=2. Konsument robi
|
| .next() -> H2D-copy (non_blocking) w watku glownym. Watek uzywa wlasnego
|
| torch.Generator (seeded), wiec ciag train-batchy jest deterministyczny i
|
| NIEZALEZNY od timingu watku oraz od RNG val-loopa (dystrybucja bez zmian)."""
|
| def __init__(self, split, generator, depth=2):
|
| self.split = split
|
| self.gen = generator
|
| self.q = queue.Queue(maxsize=depth)
|
| self._stop = threading.Event()
|
| self.t = threading.Thread(target=self._worker, daemon=True)
|
| self.t.start()
|
|
|
| def _worker(self):
|
| while not self._stop.is_set():
|
| try:
|
| item = _make_batch_cpu(self.split, self.gen)
|
| except Exception as e:
|
| self.q.put(e)
|
| return
|
| while not self._stop.is_set():
|
| try:
|
| self.q.put(item, timeout=0.5)
|
| break
|
| except queue.Full:
|
| continue
|
|
|
| def next(self):
|
| item = self.q.get()
|
| if isinstance(item, Exception):
|
| raise item
|
| return _to_device(*item)
|
|
|
| def close(self):
|
| self._stop.set()
|
|
|
| try:
|
| self.q.get_nowait()
|
| except queue.Empty:
|
| pass
|
|
|
| model = GPT(a.vocab, a.n_layer, a.n_embd, a.n_head, a.block).to(device)
|
| nparam = sum(p.numel() for p in model.parameters())
|
| log(f"model GPT-ref: {nparam/1e6:.1f}M param (L{a.n_layer} d{a.n_embd} h{a.n_head})")
|
| opt = build_optimizer(a, model)
|
| log(f"optimizer={a.optimizer}" + (f" muon_lr={a.muon_lr} (mult={a.muon_lr/a.lr:.1f}x)" if a.optimizer == "muon" else ""))
|
|
|
| start_step = 0
|
| if a.resume and os.path.exists(ckpt_path):
|
| ck = torch.load(ckpt_path, map_location=device)
|
| model.load_state_dict(ck["model"]); opt.load_state_dict(ck["opt"]); start_step = ck["step"]
|
| log(f"RESUME @ {start_step}")
|
|
|
| def lr_at(s):
|
| if s < a.warmup:
|
| return a.lr * (s + 1) / a.warmup
|
| if s >= a.steps:
|
| return a.min_lr
|
| r = (s - a.warmup) / max(1, a.steps - a.warmup)
|
| return a.min_lr + 0.5 * (a.lr - a.min_lr) * (1 + math.cos(math.pi * r))
|
|
|
| @torch.no_grad()
|
| def eval_val():
|
| model.eval()
|
| ls = []
|
| for _ in range(a.eval_iters):
|
| xb, yb = get_batch("val")
|
| with ctx:
|
| _, loss = model(xb, yb)
|
| ls.append(loss.item())
|
| model.train()
|
| return sum(ls) / len(ls)
|
|
|
| model.train()
|
| log(f"START gpt-ref: steps={a.steps} (od {start_step}) lr={a.lr}->{a.min_lr}")
|
|
|
|
|
| train_gen = torch.Generator()
|
| train_gen.manual_seed(a.seed)
|
| prefetcher = Prefetcher("train", train_gen)
|
| t0 = time.time(); running = 0.0
|
| for step in range(start_step, a.steps):
|
| lr = lr_at(step)
|
| for g in opt.param_groups:
|
| g["lr"] = lr * g.get("lr_mult", 1.0)
|
| xb, yb = prefetcher.next()
|
| with ctx:
|
| _, loss = model(xb, yb)
|
| loss.backward()
|
| gn = torch.nn.utils.clip_grad_norm_(model.parameters(), a.grad_clip) if a.grad_clip > 0 else 0.0
|
| opt.step(); opt.zero_grad(set_to_none=True)
|
| running += loss.item()
|
| if (step + 1) % a.log_every == 0:
|
| dt = time.time() - t0
|
| tok_s = a.log_every * a.block * a.batch / dt
|
| mem = torch.cuda.max_memory_allocated()/1e9 if device.type == "cuda" else 0.0
|
| log(f"step {step+1}/{a.steps} loss {running/a.log_every:.4f} lr {lr:.2e} gnorm {float(gn):.2f} {tok_s:,.0f} tok/s peakVRAM {mem:.1f}GB")
|
| with open(metrics_path, "a", encoding="utf-8") as f:
|
| f.write(json.dumps({"step": step+1, "loss": running/a.log_every, "lr": lr, "tok_s": tok_s}) + "\n")
|
| running = 0.0; t0 = time.time()
|
| if (step + 1) % a.eval_every == 0:
|
| log(f" >> VAL loss {eval_val():.4f} @ {step+1}")
|
| if (step + 1) % a.ckpt_every == 0 and not a.synthetic:
|
| torch.save({"model": model.state_dict(), "opt": opt.state_dict(), "step": step+1}, ckpt_path)
|
| log(f"ckpt @ {step+1}")
|
| prefetcher.close()
|
| log("DONE")
|
|
|
|
|
| if __name__ == "__main__":
|
| main()
|
|
|
