marketing-shibata50/OpenMythos
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

tiny tests

Browse source
This commit is contained in:
Kye Gomez 2026-04-22 12:48:33 -04:00
parent 7d78ebec79
commit 963e11277d
6 changed files with 1046 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

536
tests/small_benchmark.py Normal file
View file

@ -0,0 +1,536 @@
#!/usr/bin/env python3
"""
Side-by-side training of OpenMythos vs. a vanilla GQA transformer on a small
HuggingFace dataset (wikitext-2 by default).
Both models share the same tiny config and see the exact same batches in the
same order, so per-step loss + throughput are directly comparable. The baseline
is a dense GQA + SwiGLU stack whose unique-layer depth matches the recurrent
block's unique-parameter depth (prelude + 1 + coda), so parameter counts land
in the same ballpark.
python training/small_benchmark.py
python training/small_benchmark.py --steps 500 --device cuda
"""
from __future__ import annotations
import argparse
import time
from collections import deque
from dataclasses import dataclass, field
from typing import Deque
import torch
import torch.nn as nn
import torch.nn.functional as F
from datasets import load_dataset
from torch.utils.data import DataLoader, Dataset
from transformers import AutoTokenizer
from open_mythos import MythosConfig, OpenMythos
from open_mythos.main import (
RMSNorm,
TransformerBlock,
precompute_rope_freqs,
)
# ---------------------------------------------------------------------------
# Baseline: dense GQA + SwiGLU transformer
# ---------------------------------------------------------------------------
class BaselineTransformer(nn.Module):
"""Vanilla decoder-only transformer with dense SwiGLU FFNs.
Reuses OpenMythos's TransformerBlock (attention + FFN kernels are identical)
so any measured delta reflects the looped recurrent-depth architecture, not
kernel differences. Supports both attn_type="gqa" and "mla".
"""
def __init__(self, cfg: MythosConfig, n_layers: int):
super().__init__()
self.cfg = cfg
self.embed = nn.Embedding(cfg.vocab_size, cfg.dim)
self.layers = nn.ModuleList(
[TransformerBlock(cfg, use_moe=False) for _ in range(n_layers)]
)
self.norm = RMSNorm(cfg.dim)
self.head = nn.Linear(cfg.dim, cfg.vocab_size, bias=False)
self.head.weight = self.embed.weight # weight tying
# MLA applies RoPE to qk_rope_head_dim only; GQA rotates the full head_dim.
rope_dim = (
cfg.qk_rope_head_dim if cfg.attn_type == "mla" else cfg.dim // cfg.n_heads
)
self.register_buffer(
"freqs_cis",
precompute_rope_freqs(rope_dim, cfg.max_seq_len, cfg.rope_theta),
persistent=False,
)
self._init_weights()
def _init_weights(self) -> None:
for m in self.modules():
if isinstance(m, nn.Linear):
nn.init.normal_(m.weight, std=0.02)
elif isinstance(m, nn.Embedding):
nn.init.normal_(m.weight, std=0.02)
@staticmethod
def _causal_mask(T: int, device: torch.device) -> torch.Tensor:
mask = torch.full((1, 1, T, T), float("-inf"), device=device)
return torch.triu(mask, diagonal=1)
def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
T = input_ids.shape[1]
x = self.embed(input_ids)
freqs_cis = self.freqs_cis[:T]
mask = self._causal_mask(T, x.device) if T > 1 else None
for i, layer in enumerate(self.layers):
x = layer(x, freqs_cis, mask, cache_key=f"layer_{i}")
return self.head(self.norm(x))
# ---------------------------------------------------------------------------
# Dataset: tokenize once, pack into fixed-length next-token pairs
# ---------------------------------------------------------------------------
class PackedLMDataset(Dataset):
"""Flatten an HF text dataset into one token buffer, slice fixed-length pairs.
Accepts either map-style or streaming (`IterableDataset`) HF datasets —
iteration stops once `max_tokens` are collected, so large corpora like
TinyStories can be streamed without downloading the whole thing.
"""
def __init__(
self,
hf_ds,
tokenizer,
seq_len: int,
max_tokens: int,
text_field: str = "text",
):
buf: list[int] = []
for sample in hf_ds:
text = sample[text_field]
if not text or not text.strip():
continue
buf.extend(tokenizer.encode(text, add_special_tokens=False))
if len(buf) >= max_tokens:
break
self.seq_len = seq_len
n_pairs = max(1, (len(buf) - 1) // seq_len)
buf = buf[: n_pairs * seq_len + 1]
self.data = torch.tensor(buf, dtype=torch.long)
def __len__(self) -> int:
return (len(self.data) - 1) // self.seq_len
def __getitem__(self, idx: int):
s = idx * self.seq_len
chunk = self.data[s : s + self.seq_len + 1]
return chunk[:-1], chunk[1:]
# ---------------------------------------------------------------------------
# Metrics
# ---------------------------------------------------------------------------
@dataclass
class Metrics:
total_loss: float = 0.0
total_tokens: int = 0
total_time: float = 0.0
steps: int = 0
first_losses: list[float] = field(default_factory=list)
last_losses: Deque[float] = field(default_factory=lambda: deque(maxlen=10))
def update(self, loss: float, tokens: int, seconds: float) -> None:
self.total_loss += loss
self.total_tokens += tokens
self.total_time += seconds
self.steps += 1
if len(self.first_losses) < 10:
self.first_losses.append(loss)
self.last_losses.append(loss)
@property
def avg_loss(self) -> float:
return self.total_loss / max(1, self.steps)
@property
def tok_per_sec(self) -> float:
return self.total_tokens / max(1e-9, self.total_time)
@property
def initial_loss(self) -> float:
return sum(self.first_losses) / max(1, len(self.first_losses))
@property
def final_loss(self) -> float:
return sum(self.last_losses) / max(1, len(self.last_losses))
# ---------------------------------------------------------------------------
# Training step
# ---------------------------------------------------------------------------
def train_step(
model: nn.Module,
x: torch.Tensor,
y: torch.Tensor,
optimizer: torch.optim.Optimizer,
device: torch.device,
vocab_size: int,
) -> tuple[float, float]:
"""Run one optimizer step; return (loss, wall-clock seconds)."""
t0 = time.perf_counter()
model.train()
optimizer.zero_grad()
logits = model(x)
loss = F.cross_entropy(logits.view(-1, vocab_size), y.view(-1))
loss.backward()
nn.utils.clip_grad_norm_(model.parameters(), 1.0)
optimizer.step()
if device.type == "cuda":
torch.cuda.synchronize()
return loss.item(), time.perf_counter() - t0
@torch.no_grad()
def evaluate(
model: nn.Module,
loader: DataLoader,
device: torch.device,
vocab_size: int,
max_batches: int | None = None,
n_loops: int | None = None,
) -> float:
"""Mean cross-entropy over (up to `max_batches`) of the loader.
`n_loops` is only forwarded to OpenMythos; for any other module the kwarg
is dropped, so the same function benchmarks baseline and mythos uniformly.
"""
model.eval()
total_loss = 0.0
total_tokens = 0
for i, (x, y) in enumerate(loader):
if max_batches is not None and i >= max_batches:
break
x = x.to(device, non_blocking=True)
y = y.to(device, non_blocking=True)
if isinstance(model, OpenMythos):
logits = model(x, n_loops=n_loops)
else:
logits = model(x)
# sum-reduction so we weight by token count, not batch count
loss = F.cross_entropy(logits.view(-1, vocab_size), y.view(-1), reduction="sum")
total_loss += loss.item()
total_tokens += y.numel()
return total_loss / max(1, total_tokens)
# ---------------------------------------------------------------------------
# Config + utilities
# ---------------------------------------------------------------------------
def build_tiny_cfg(vocab_size: int, seq_len: int) -> MythosConfig:
"""Tiny shared config with MLA attention — runs in reasonable time on CPU.
MLA LoRA ranks and head dims scale with `dim=128` instead of the
2048-dim-sized defaults (q_lora_rank=1536, qk_nope_head_dim=128, ...),
which would otherwise dominate the parameter count at this scale.
"""
return MythosConfig(
vocab_size=vocab_size,
dim=128,
n_heads=4,
n_kv_heads=2,
max_seq_len=seq_len,
max_loop_iters=4,
prelude_layers=1,
coda_layers=1,
attn_type="mla",
kv_lora_rank=64,
q_lora_rank=128,
qk_rope_head_dim=16,
qk_nope_head_dim=32,
v_head_dim=32,
n_experts=4,
n_shared_experts=1,
n_experts_per_tok=2,
expert_dim=128,
lora_rank=4,
rope_theta=10000.0,
dropout=0.0,
)
def count_params(model: nn.Module) -> int:
return sum(p.numel() for p in model.parameters())
def fmt_count(n: float) -> str:
for unit in ("", "K", "M", "B"):
if abs(n) < 1000:
return f"{n:.2f}{unit}"
n /= 1000
return f"{n:.2f}T"
# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
def parse_args() -> argparse.Namespace:
p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--steps", type=int, default=1000)
p.add_argument("--batch-size", type=int, default=32)
p.add_argument("--seq-len", type=int, default=256)
p.add_argument("--lr", type=float, default=3e-4)
# Defaults point at TinyStories — simpler vocabulary + shorter documents
# lets a dim=128 model actually reach a meaningful loss in modest time.
p.add_argument("--dataset", default="roneneldan/TinyStories")
p.add_argument(
"--dataset-config",
default="",
help="pass '' for datasets with no config (e.g. TinyStories)",
)
p.add_argument("--train-split", default="train")
p.add_argument("--eval-split", default="validation")
p.add_argument(
"--train-tokens",
type=int,
default=5_000_000,
help="max tokens to materialize for the training buffer",
)
p.add_argument(
"--eval-tokens",
type=int,
default=200_000,
help="max tokens to materialize for the held-out eval buffer",
)
p.add_argument("--text-field", default="text")
p.add_argument("--tokenizer", default="gpt2")
p.add_argument("--log-every", type=int, default=25)
p.add_argument(
"--eval-every",
type=int,
default=200,
help="run held-out eval every N steps (0 disables)",
)
p.add_argument("--eval-batches", type=int, default=20)
p.add_argument(
"--depth-sweep",
default="1,2,4,8,16",
help="comma-separated n_loops values for OpenMythos depth-extrapolation eval",
)
p.add_argument("--seed", type=int, default=0)
p.add_argument(
"--device",
default="cuda" if torch.cuda.is_available() else "cpu",
)
return p.parse_args()
def load_text_ds(name: str, config: str, split: str):
"""Streaming `load_dataset` with optional config (empty string == no config)."""
if config:
return load_dataset(name, config, split=split, streaming=True)
return load_dataset(name, split=split, streaming=True)
def main() -> None:
args = parse_args()
device = torch.device(args.device)
print(
f"[setup] device={device} batch={args.batch_size} "
f"seq_len={args.seq_len} steps={args.steps}"
)
tokenizer = AutoTokenizer.from_pretrained(args.tokenizer)
# AutoTokenizer.vocab_size can be smaller than the head size for BPE
# tokenizers with added tokens; use len(tokenizer) to be safe.
vocab_size = len(tokenizer)
print(f"[setup] tokenizer={args.tokenizer} vocab_size={vocab_size:,}")
# ------------------------------------------------------------------
# Data: streamed train + held-out eval splits
# ------------------------------------------------------------------
print(f"[setup] dataset={args.dataset} config={args.dataset_config or '∅'}")
raw_train = load_text_ds(args.dataset, args.dataset_config, args.train_split)
train_ds = PackedLMDataset(
raw_train, tokenizer, args.seq_len, args.train_tokens, args.text_field
)
raw_eval = load_text_ds(args.dataset, args.dataset_config, args.eval_split)
eval_ds = PackedLMDataset(
raw_eval, tokenizer, args.seq_len, args.eval_tokens, args.text_field
)
print(
f"[setup] train tokens={train_ds.data.numel():,} pairs={len(train_ds)} | "
f"eval tokens={eval_ds.data.numel():,} pairs={len(eval_ds)}"
)
torch.manual_seed(args.seed)
train_loader = DataLoader(
train_ds, batch_size=args.batch_size, shuffle=True, drop_last=True
)
eval_loader = DataLoader(
eval_ds, batch_size=args.batch_size, shuffle=False, drop_last=False
)
# ------------------------------------------------------------------
# Models — same init seed so both start from the same embedding
# ------------------------------------------------------------------
cfg = build_tiny_cfg(vocab_size, args.seq_len)
torch.manual_seed(args.seed)
mythos = OpenMythos(cfg).to(device)
# Parameter-matched depth: prelude + one unique recurrent block + coda.
baseline_layers = cfg.prelude_layers + 1 + cfg.coda_layers
torch.manual_seed(args.seed)
baseline = BaselineTransformer(cfg, n_layers=baseline_layers).to(device)
n_m, n_b = count_params(mythos), count_params(baseline)
print(
f"[setup] OpenMythos params = {fmt_count(n_m)} ({n_m:,})\n"
f"[setup] Baseline params = {fmt_count(n_b)} ({n_b:,}) "
f"[{baseline_layers} layers]"
)
print(
f"[setup] Mythos runtime depth = prelude({cfg.prelude_layers}) + "
f"loops({cfg.max_loop_iters}) + coda({cfg.coda_layers}) = "
f"{cfg.prelude_layers + cfg.max_loop_iters + cfg.coda_layers}"
)
opt_m = torch.optim.AdamW(
mythos.parameters(), lr=args.lr, betas=(0.9, 0.95), weight_decay=0.1
)
opt_b = torch.optim.AdamW(
baseline.parameters(), lr=args.lr, betas=(0.9, 0.95), weight_decay=0.1
)
mm, bm = Metrics(), Metrics()
eval_history: list[tuple[int, float, float]] = [] # (step, mythos_eval, base_eval)
header = (
f"\n{'step':>6} | {'mythos loss':>12} | {'base loss':>10} | "
f"{'mythos tok/s':>13} | {'base tok/s':>11}"
)
print(header)
print("-" * len(header))
# ------------------------------------------------------------------
# Training loop with periodic held-out eval
# ------------------------------------------------------------------
data_iter = iter(train_loader)
t_total = time.perf_counter()
for step in range(1, args.steps + 1):
try:
x, y = next(data_iter)
except StopIteration:
data_iter = iter(train_loader)
x, y = next(data_iter)
x = x.to(device, non_blocking=True)
y = y.to(device, non_blocking=True)
tokens = x.numel()
loss_m, dt_m = train_step(mythos, x, y, opt_m, device, vocab_size)
loss_b, dt_b = train_step(baseline, x, y, opt_b, device, vocab_size)
mm.update(loss_m, tokens, dt_m)
bm.update(loss_b, tokens, dt_b)
if step == 1 or step % args.log_every == 0:
print(
f"{step:>6} | {loss_m:>12.4f} | {loss_b:>10.4f} | "
f"{tokens / dt_m:>13,.0f} | {tokens / dt_b:>11,.0f}"
)
if args.eval_every and step % args.eval_every == 0:
eval_m = evaluate(
mythos, eval_loader, device, vocab_size, args.eval_batches
)
eval_b = evaluate(
baseline, eval_loader, device, vocab_size, args.eval_batches
)
eval_history.append((step, eval_m, eval_b))
print(
f" [eval @ step {step}] mythos {eval_m:.4f} baseline {eval_b:.4f} "
f"(Δ = {eval_m - eval_b:+.4f})"
)
total_wall = time.perf_counter() - t_total
# ------------------------------------------------------------------
# Summary
# ------------------------------------------------------------------
bar = "=" * 70
print(f"\n{bar}\nSummary ({args.steps} steps, wall clock {total_wall:.1f}s)\n{bar}")
print(f" {'':<24} {'OpenMythos':>16} {'Baseline':>16}")
print(f" {'params':<24} {fmt_count(n_m):>16} {fmt_count(n_b):>16}")
print(
f" {'initial train (first 10)':<24} "
f"{mm.initial_loss:>16.4f} {bm.initial_loss:>16.4f}"
)
print(
f" {'final train (last 10)':<24} "
f"{mm.final_loss:>16.4f} {bm.final_loss:>16.4f}"
)
print(
f" {'avg train (all steps)':<24} "
f"{mm.avg_loss:>16.4f} {bm.avg_loss:>16.4f}"
)
print(
f" {'train time (sec)':<24} "
f"{mm.total_time:>16.2f} {bm.total_time:>16.2f}"
)
print(
f" {'avg tok/s':<24} " f"{mm.tok_per_sec:>16,.0f} {bm.tok_per_sec:>16,.0f}"
)
print(
f" {'sec/step':<24} "
f"{mm.total_time / max(1, mm.steps):>16.4f} "
f"{bm.total_time / max(1, bm.steps):>16.4f}"
)
# ------------------------------------------------------------------
# Depth extrapolation: OpenMythos eval loss as a function of n_loops.
# Trained at cfg.max_loop_iters; we run inference with a sweep to see
# whether additional loops keep improving (depth extrapolation) or the
# model collapses outside the trained regime.
# ------------------------------------------------------------------
loops_sweep = sorted({int(s) for s in args.depth_sweep.split(",") if s.strip()})
print(f"\n{bar}\nDepth extrapolation (held-out eval, full eval set)\n{bar}")
baseline_eval = evaluate(baseline, eval_loader, device, vocab_size)
print(f" Baseline (fixed depth) : eval loss = {baseline_eval:.4f}")
# First collect all sweep losses, then print with deltas vs. the trained depth.
sweep: list[tuple[int, float]] = []
for nl in loops_sweep:
sweep.append(
(nl, evaluate(mythos, eval_loader, device, vocab_size, n_loops=nl))
)
trained_loss = next((loss for nl, loss in sweep if nl == cfg.max_loop_iters), None)
print(f" OpenMythos (trained at n_loops={cfg.max_loop_iters}):")
print(f" {'n_loops':>8} {'eval loss':>10} {'Δ vs trained':>14}")
for nl, loss in sweep:
if trained_loss is None or nl == cfg.max_loop_iters:
delta_str = ""
else:
delta_str = f"{loss - trained_loss:+.4f}"
marker = " ←trained" if nl == cfg.max_loop_iters else ""
print(f" {nl:>8} {loss:>10.4f} {delta_str:>14}{marker}")
if __name__ == "__main__":
main()