""" ZeroShot-1B — single-file LLaMA-style 1.19B trainer + inference. Architecture trimmed from the original spec to fit microbatch=4 on a 32 GB 5090. Original was FFN=8192, ctx=4096 → ~29 GB activations alone, OOM. This: FFN=5632, ctx=3072 → ~21 GB activations, ~29 GB peak. Subcommands: python train.py base --no-compile python train.py mid --checkpoint ckpt_base_final.pt --no-compile python train.py finetune --checkpoint ckpt_mid_final.pt --no-compile python train.py finetune --checkpoint ckpt_mid_final.pt --no-compile --skip_mid python train.py generate --checkpoint ckpt_sft_final.pt --prompt "..." --chat Auto-resume: re-launching a stage picks up the newest matching ckpt in --ckpt_dir (model + optimizer + step + RNG + dataset sample counter). Pass --fresh to ignore. --checkpoint is for cross-stage *weight* init only (fresh optimizer). Exit codes: 0 — clean finish 2 — recoverable failure (data, OOM, SIGTERM) — emergency ckpt saved, run.sh restarts 130 — Ctrl-C """ import argparse import math import os import queue import random import signal import sys import threading import time from dataclasses import asdict, dataclass from pathlib import Path from typing import Iterator import numpy as np import tiktoken import torch import torch.nn as nn import torch.nn.functional as F from datasets import load_dataset # ============================================================================= # CONFIG # ============================================================================= @dataclass class ModelConfig: vocab_size: int = 50304 # GPT-2 (50257) padded for tensor cores n_layers: int = 24 n_heads: int = 16 n_kv_heads: int = 4 # GQA: 4 KV heads, 4 query heads per group d_model: int = 2048 d_ff: int = 5632 # SwiGLU expansion (~2.75x d_model). 8192 OOMs at ctx>=3072 mb=4. max_seq_len: int = 3072 rope_theta: float = 10000.0 rms_eps: float = 1e-5 @dataclass class TrainConfig: stage: str = "base" micro_batch_size: int = 4 grad_accum: int = 20 # tokens/step = 4 * 20 * 3072 = 245,760 (~0.25M) seq_len: int = 3072 max_steps: int = 25_000 warmup_steps: int = 2_000 peak_lr: float = 3e-4 min_lr: float = 3e-5 weight_decay: float = 0.1 beta1: float = 0.9 beta2: float = 0.95 grad_clip: float = 1.0 log_every: int = 10 ckpt_interval_sec: int = 30 * 60 # time-based; vast instances die at unpredictable steps ckpt_keep: int = 1 ckpt_dir: str = "./checkpoints" seed: int = 42 cost_per_hour: float = 0.365 wandb_project: str = "zeroshot-1b" data_max_retries: int = 5 data_retry_base_delay: float = 2.0 STAGE_PRESETS = { "base": dict( max_steps=25_000, warmup_steps=2_000, peak_lr=3e-4, min_lr=3e-5, ckpt_prefix="ckpt_base", micro_batch_size=4, grad_accum=20, ), "mid": dict( max_steps=2_500, warmup_steps=200, peak_lr=1e-4, min_lr=1e-5, ckpt_prefix="ckpt_mid", micro_batch_size=4, grad_accum=20, ), "finetune": dict( max_steps=2_000, warmup_steps=100, peak_lr=5e-5, min_lr=5e-6, ckpt_prefix="ckpt_sft", micro_batch_size=4, grad_accum=8, ), } # ============================================================================= # MODEL — RMSNorm, RoPE, GQA attention, SwiGLU, transformer # ============================================================================= class RMSNorm(nn.Module): def __init__(self, dim, eps=1e-5): super().__init__() self.eps = eps self.weight = nn.Parameter(torch.ones(dim)) def forward(self, x): x32 = x.float() norm = x32 * torch.rsqrt(x32.pow(2).mean(-1, keepdim=True) + self.eps) return (norm * self.weight.float()).type_as(x) def precompute_rope(head_dim, max_seq_len, theta, device="cpu"): inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim)) t = torch.arange(max_seq_len, device=device).float() freqs = torch.outer(t, inv_freq) return freqs.cos(), freqs.sin() def apply_rope(x, cos, sin): # x: (B, H, T, D). cos/sin: (T, D/2). GPT-NeoX-style split-halves rotation. x1, x2 = x.chunk(2, dim=-1) cos = cos[None, None, :, :] sin = sin[None, None, :, :] return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1) class Attention(nn.Module): def __init__(self, cfg: ModelConfig): super().__init__() assert cfg.d_model % cfg.n_heads == 0 assert cfg.n_heads % cfg.n_kv_heads == 0 self.n_heads = cfg.n_heads self.n_kv_heads = cfg.n_kv_heads self.head_dim = cfg.d_model // cfg.n_heads self.n_rep = cfg.n_heads // cfg.n_kv_heads self.q_proj = nn.Linear(cfg.d_model, cfg.n_heads * self.head_dim, bias=False) self.k_proj = nn.Linear(cfg.d_model, cfg.n_kv_heads * self.head_dim, bias=False) self.v_proj = nn.Linear(cfg.d_model, cfg.n_kv_heads * self.head_dim, bias=False) self.o_proj = nn.Linear(cfg.n_heads * self.head_dim, cfg.d_model, bias=False) def forward(self, x, cos, sin, kv_cache=None, pos=0): B, T, _ = x.shape q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2) k = self.k_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2) v = self.v_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2) cos_t = cos[pos:pos + T] sin_t = sin[pos:pos + T] q = apply_rope(q, cos_t, sin_t) k = apply_rope(k, cos_t, sin_t) if kv_cache is not None: k_cache, v_cache = kv_cache k_cache[:, :, pos:pos + T] = k v_cache[:, :, pos:pos + T] = v k = k_cache[:, :, :pos + T] v = v_cache[:, :, :pos + T] # SDPA's enable_gqa is still flagged experimental and has open bug reports; # explicit repeat is correct on every nightly we'd be using. k = k.repeat_interleave(self.n_rep, dim=1) v = v.repeat_interleave(self.n_rep, dim=1) # is_causal=True uses top-left mask alignment in SDPA, which is correct for # training (Q==K len) and prefill at pos=0. Single-token decode wants no # mask. We never do chunked prefill (pos>0 with T>1), so this is sufficient. is_causal = (kv_cache is None) or (T > 1) out = F.scaled_dot_product_attention(q, k, v, is_causal=is_causal) out = out.transpose(1, 2).contiguous().view(B, T, -1) return self.o_proj(out) class SwiGLU(nn.Module): def __init__(self, cfg: ModelConfig): super().__init__() self.gate_proj = nn.Linear(cfg.d_model, cfg.d_ff, bias=False) self.up_proj = nn.Linear(cfg.d_model, cfg.d_ff, bias=False) self.down_proj = nn.Linear(cfg.d_ff, cfg.d_model, bias=False) def forward(self, x): return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) class Block(nn.Module): def __init__(self, cfg: ModelConfig): super().__init__() self.attn_norm = RMSNorm(cfg.d_model, cfg.rms_eps) self.attn = Attention(cfg) self.mlp_norm = RMSNorm(cfg.d_model, cfg.rms_eps) self.mlp = SwiGLU(cfg) def forward(self, x, cos, sin, kv_cache=None, pos=0): x = x + self.attn(self.attn_norm(x), cos, sin, kv_cache, pos) x = x + self.mlp(self.mlp_norm(x)) return x class LLaMA(nn.Module): def __init__(self, cfg: ModelConfig): super().__init__() self.cfg = cfg self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.d_model) self.layers = nn.ModuleList([Block(cfg) for _ in range(cfg.n_layers)]) self.norm = RMSNorm(cfg.d_model, cfg.rms_eps) cos, sin = precompute_rope(cfg.d_model // cfg.n_heads, cfg.max_seq_len, cfg.rope_theta) self.register_buffer("rope_cos", cos, persistent=False) self.register_buffer("rope_sin", sin, persistent=False) self.apply(self._init_weights) for name, p in self.named_parameters(): if name.endswith("o_proj.weight") or name.endswith("down_proj.weight"): nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * cfg.n_layers)) @staticmethod def _init_weights(m): if isinstance(m, nn.Linear): nn.init.normal_(m.weight, mean=0.0, std=0.02) if m.bias is not None: nn.init.zeros_(m.bias) elif isinstance(m, nn.Embedding): nn.init.normal_(m.weight, mean=0.0, std=0.02) def forward(self, idx, targets=None, kv_caches=None, pos=0, ignore_index=-100): x = self.tok_emb(idx) for i, layer in enumerate(self.layers): kv = kv_caches[i] if kv_caches is not None else None x = layer(x, self.rope_cos, self.rope_sin, kv, pos) x = self.norm(x) logits = F.linear(x, self.tok_emb.weight) # tied lm head loss = None if targets is not None: loss = F.cross_entropy( logits.view(-1, logits.size(-1)).float(), targets.view(-1), ignore_index=ignore_index, ) return logits, loss def num_params(self): return sum(p.numel() for p in self.parameters()) def configure_optimizer(self, lr, weight_decay, betas): # 8-bit AdamW: more stable than bf16 Adam states (variance term is fragile), # less memory than fp32 states. Embedding gets fp32 override per bnb docs. import bitsandbytes as bnb decay, no_decay = [], [] for _, p in self.named_parameters(): if not p.requires_grad: continue (decay if p.dim() >= 2 else no_decay).append(p) groups = [ {"params": decay, "weight_decay": weight_decay}, {"params": no_decay, "weight_decay": 0.0}, ] opt = bnb.optim.AdamW8bit(groups, lr=lr, betas=betas) bnb.optim.GlobalOptimManager.get_instance().register_module_override( self.tok_emb, "weight", {"optim_bits": 32} ) return opt # ============================================================================= # DATA — streaming loaders with retry + background prefetch # ============================================================================= ENC = tiktoken.get_encoding("gpt2") EOT = ENC.eot_token # 50256 IGNORE = -100 USER_TAG = "\n<|user|>\n" ASSISTANT_TAG = "\n<|assistant|>\n" class _StreamError(RuntimeError): pass def _retry_iter(open_fn, max_retries, base_delay): """Run open_fn() to get a fresh iterator, yield from it; on exception recreate via open_fn() with exp backoff. open_fn closes over the right skip count (set on stage start; on hard restart, run.sh re-invokes us and we recreate from the latest ckpt's samples_consumed).""" attempt = 0 while True: try: it = open_fn() for x in it: yield x attempt = 0 return except StopIteration: return except Exception as e: attempt += 1 if attempt > max_retries: raise _StreamError(f"max retries ({max_retries}) exceeded: {e}") from e delay = base_delay * (2 ** (attempt - 1)) print(f"[data] {type(e).__name__}: {e}; retry {attempt}/{max_retries} in {delay:.1f}s", flush=True) time.sleep(delay) def _open_fineweb(seed, skip): ds = load_dataset("HuggingFaceFW/fineweb-edu", name="sample-10BT", split="train", streaming=True) ds = ds.shuffle(seed=seed, buffer_size=10_000) return ds.skip(skip) if skip else ds def _open_cosmopedia(seed, skip): ds = load_dataset("HuggingFaceTB/cosmopedia", name="web_samples_v2", split="train", streaming=True) ds = ds.shuffle(seed=seed, buffer_size=10_000) return ds.skip(skip) if skip else ds def _open_ultrachat(seed, skip): ds = load_dataset("HuggingFaceH4/ultrachat_200k", split="train_sft", streaming=True) ds = ds.shuffle(seed=seed, buffer_size=2_000) return ds.skip(skip) if skip else ds def _doc_tokens(ds_iter, text_key): for sample in ds_iter: text = sample.get(text_key) or "" if not text: continue ids = ENC.encode_ordinary(text) ids.append(EOT) for t in ids: yield t def _pack(token_iter: Iterator[int], seq_len): buf = [] for tok in token_iter: buf.append(tok) if len(buf) >= seq_len + 1: chunk = buf[: seq_len + 1] del buf[: seq_len] x = torch.tensor(chunk[:-1], dtype=torch.long) y = torch.tensor(chunk[1:], dtype=torch.long) yield x, y def base_loader(seq_len, seed, skip_samples, max_retries, base_delay): yield from _pack( _doc_tokens(_retry_iter(lambda: _open_fineweb(seed, skip_samples), max_retries, base_delay), "text"), seq_len, ) def mid_loader(seq_len, seed, skip_samples, max_retries, base_delay, mix_ratio=0.1): """90% FineWeb-Edu / 10% Cosmopedia, switched at document boundaries.""" fw_iter = _doc_tokens( _retry_iter(lambda: _open_fineweb(seed, skip_samples), max_retries, base_delay), "text") co_iter = _doc_tokens( _retry_iter(lambda: _open_cosmopedia(seed + 1, skip_samples // 10), max_retries, base_delay), "text") rng = random.Random(seed) def mixed(): while True: src = co_iter if rng.random() < mix_ratio else fw_iter while True: t = next(src) yield t if t == EOT: break yield from _pack(mixed(), seq_len) def _format_chat(messages): ids, mask = [], [] for msg in messages: role = msg.get("role") content = (msg.get("content") or "").strip() if not content: continue if role == "user": tag = ENC.encode_ordinary(USER_TAG + content) ids.extend(tag); mask.extend([0] * len(tag)) elif role == "assistant": head = ENC.encode_ordinary(ASSISTANT_TAG) body = ENC.encode_ordinary(content) ids.extend(head); mask.extend([0] * len(head)) ids.extend(body); mask.extend([1] * len(body)) ids.append(EOT); mask.append(1) return ids, mask def sft_loader(seq_len, seed, skip_samples, max_retries, base_delay): for sample in _retry_iter(lambda: _open_ultrachat(seed, skip_samples), max_retries, base_delay): ids, mask = _format_chat(sample.get("messages") or []) if len(ids) < 8: continue ids = ids[: seq_len + 1] mask = mask[: seq_len + 1] pad_n = (seq_len + 1) - len(ids) if pad_n > 0: ids += [EOT] * pad_n mask += [0] * pad_n x = torch.tensor(ids[:-1], dtype=torch.long) y = torch.tensor( [t if m else IGNORE for t, m in zip(ids[1:], mask[1:])], dtype=torch.long, ) yield x, y class _Prefetcher: """Background thread fills a bounded queue. Errors propagate on next().""" def __init__(self, gen, max_prefetch=8): self.q = queue.Queue(maxsize=max_prefetch) self._sentinel = object() self._error = None self.thread = threading.Thread(target=self._run, args=(gen,), daemon=True) self.thread.start() def _run(self, gen): try: for x in gen: self.q.put(x) except BaseException as e: self._error = e finally: self.q.put(self._sentinel) def __iter__(self): return self def __next__(self): x = self.q.get() if x is self._sentinel: if self._error is not None: raise self._error raise StopIteration return x class StagedLoader: """Wraps a sample generator into a batch iterator and counts samples consumed.""" def __init__(self, gen_fn, train_cfg, stage_seed, skip_samples=0, prefetch=8): self.samples_consumed = skip_samples self.micro_batch = train_cfg.micro_batch_size gen = gen_fn( seq_len=train_cfg.seq_len, seed=stage_seed, skip_samples=skip_samples, max_retries=train_cfg.data_max_retries, base_delay=train_cfg.data_retry_base_delay, ) self._gen = _Prefetcher(gen, max_prefetch=prefetch) def __iter__(self): return self def __next__(self): xs, ys = [], [] for _ in range(self.micro_batch): x, y = next(self._gen) xs.append(x); ys.append(y) self.samples_consumed += 1 return torch.stack(xs), torch.stack(ys) def build_loader(stage, train_cfg, stage_seed, skip_samples=0): gen = {"base": base_loader, "mid": mid_loader, "finetune": sft_loader}[stage] return StagedLoader(gen, train_cfg, stage_seed, skip_samples) # ============================================================================= # CHECKPOINT + TRAIN HELPERS # ============================================================================= def cosine_lr(step, warmup, max_steps, peak, min_lr): if step < warmup: return peak * (step + 1) / max(1, warmup) if step >= max_steps: return min_lr progress = (step - warmup) / max(1, max_steps - warmup) return min_lr + 0.5 * (peak - min_lr) * (1.0 + math.cos(math.pi * progress)) def find_latest_ckpt(ckpt_dir: Path, prefix: str): files = list(ckpt_dir.glob(f"{prefix}_step*.pt")) return max(files, key=lambda p: p.stat().st_mtime) if files else None def prune_ckpts(ckpt_dir: Path, prefix: str, keep: int): files = sorted(ckpt_dir.glob(f"{prefix}_step*.pt"), key=lambda p: p.stat().st_mtime) for f in files[:-keep]: try: f.unlink() except OSError: pass def get_rng_state(): return { "torch": torch.get_rng_state(), "torch_cuda": torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None, "numpy": np.random.get_state(), "python": random.getstate(), } def set_rng_state(state): torch.set_rng_state(state["torch"].cpu().byte()) if state.get("torch_cuda") is not None and torch.cuda.is_available(): torch.cuda.set_rng_state_all([s.cpu().byte() for s in state["torch_cuda"]]) np.random.set_state(state["numpy"]) random.setstate(state["python"]) def save_checkpoint(path, model, optimizer, step, train_cfg, model_cfg, samples_consumed): payload = { "model": model.state_dict(), "optimizer": optimizer.state_dict(), "step": step, "train_cfg": asdict(train_cfg), "model_cfg": asdict(model_cfg), "samples_consumed": samples_consumed, "rng": get_rng_state(), } tmp = str(path) + ".tmp" torch.save(payload, tmp) os.replace(tmp, path) def load_checkpoint(path, model, optimizer=None, map_location="cuda"): ckpt = torch.load(path, map_location=map_location, weights_only=False) sd = ckpt["model"] if any(k.startswith("_orig_mod.") for k in sd): sd = {k.removeprefix("_orig_mod."): v for k, v in sd.items()} model.load_state_dict(sd) if optimizer is not None and "optimizer" in ckpt: optimizer.load_state_dict(ckpt["optimizer"]) return ckpt def estimate_vram_gb(model_cfg: ModelConfig, train_cfg: TrainConfig, n_params: int): weight_gb = n_params * 2 / 1e9 grad_gb = n_params * 2 / 1e9 optim_gb = n_params * 2 / 1e9 + 0.82 # 8-bit AdamW + fp32 emb override fixed = weight_gb + grad_gb + optim_gb B, T = train_cfg.micro_batch_size, train_cfg.seq_len per_layer = B * T * (13 * model_cfg.d_model + 8 * model_cfg.d_ff) * 2 / 1e9 return fixed, per_layer * model_cfg.n_layers def print_startup(model, train_cfg, model_cfg, ckpt_prefix, start_step): n = model.num_params() fixed, acts = estimate_vram_gb(model_cfg, train_cfg, n) tps = train_cfg.micro_batch_size * train_cfg.grad_accum * train_cfg.seq_len bar = "=" * 64 print(bar) print(f" ZeroShot-1B — stage: {train_cfg.stage} (prefix: {ckpt_prefix})") print(bar) print(f" params: {n/1e9:.3f}B ({n:,})") print(f" dtype: bfloat16 optim: AdamW8bit (bnb) ctx: {train_cfg.seq_len}") print(f" arch: L={model_cfg.n_layers} H={model_cfg.n_heads} KV={model_cfg.n_kv_heads} " f"d={model_cfg.d_model} ffn={model_cfg.d_ff}") print(f" micro_batch: {train_cfg.micro_batch_size} grad_accum: {train_cfg.grad_accum} " f"tokens/step: {tps:,}") print(f" schedule: cosine warmup={train_cfg.warmup_steps} peak={train_cfg.peak_lr:.1e} " f"min={train_cfg.min_lr:.1e} steps={train_cfg.max_steps}") print(f" total tokens: {tps*train_cfg.max_steps/1e9:.2f}B (start step {start_step})") print(f" vram (est): fixed {fixed:.1f} + acts {acts:.1f} = {fixed+acts:.1f} GB / 32 GB") print(f" ckpt every: {train_cfg.ckpt_interval_sec//60} min wall → " f"{train_cfg.ckpt_dir}/ (keep last {train_cfg.ckpt_keep})") print(bar, flush=True) def fmt_eta(seconds): if seconds <= 0 or math.isinf(seconds): return "?" h = int(seconds // 3600); m = int((seconds % 3600) // 60) return f"{h:d}h{m:02d}m" # ============================================================================= # TRAIN LOOP # ============================================================================= def train(stage: str, args): preset = STAGE_PRESETS[stage] train_cfg = TrainConfig(stage=stage, **{k: v for k, v in preset.items() if k != "ckpt_prefix"}) ckpt_prefix = preset["ckpt_prefix"] if args.lr is not None: train_cfg.peak_lr = args.lr if args.min_lr is not None: train_cfg.min_lr = args.min_lr if args.batch_size is not None: train_cfg.micro_batch_size = args.batch_size if args.grad_accum is not None: train_cfg.grad_accum = args.grad_accum if args.steps is not None: train_cfg.max_steps = args.steps if args.warmup is not None: train_cfg.warmup_steps = args.warmup if args.ckpt_dir is not None: train_cfg.ckpt_dir = args.ckpt_dir if args.seed is not None: train_cfg.seed = args.seed ckpt_dir = Path(train_cfg.ckpt_dir) ckpt_dir.mkdir(parents=True, exist_ok=True) device = "cuda" torch.manual_seed(train_cfg.seed) np.random.seed(train_cfg.seed) random.seed(train_cfg.seed) torch.set_float32_matmul_precision("high") model_cfg = ModelConfig() model = LLaMA(model_cfg).to(device=device, dtype=torch.bfloat16) optimizer = model.configure_optimizer( lr=train_cfg.peak_lr, weight_decay=train_cfg.weight_decay, betas=(train_cfg.beta1, train_cfg.beta2), ) start_step = 0 samples_consumed = 0 resume_path = None if args.fresh else find_latest_ckpt(ckpt_dir, ckpt_prefix) if resume_path is not None: print(f"[resume] {resume_path}") ckpt = load_checkpoint(resume_path, model, optimizer, map_location=device) start_step = ckpt["step"] samples_consumed = ckpt.get("samples_consumed", 0) if "rng" in ckpt: set_rng_state(ckpt["rng"]) elif args.checkpoint is not None: init_path = Path(args.checkpoint) if not init_path.is_absolute() and not init_path.exists(): init_path = ckpt_dir / args.checkpoint print(f"[init-from] {init_path} (weights only, fresh optimizer)") load_checkpoint(init_path, model, optimizer=None, map_location=device) elif stage != "base": raise SystemExit(f"{stage} requires --checkpoint pointing at the prior stage's ckpt " f"(no auto-resume ckpt found in {ckpt_dir})") stage_seed = train_cfg.seed + {"base": 0, "mid": 1000, "finetune": 2000}[stage] model = torch.compile(model, mode="default") print_startup(model, train_cfg, model_cfg, ckpt_prefix, start_step) use_wandb = bool(os.environ.get("WANDB_API_KEY")) and not args.no_wandb if use_wandb: import wandb wandb.init( project=train_cfg.wandb_project, name=f"{stage}-{int(time.time())}", config={**asdict(train_cfg), **asdict(model_cfg)}, resume="allow", ) loader = build_loader(stage, train_cfg, stage_seed, samples_consumed) data_iter = iter(loader) tokens_per_step = train_cfg.micro_batch_size * train_cfg.grad_accum * train_cfg.seq_len # ---- emergency-ckpt helpers ---- _state = {"step": start_step} def emergency_save(reason): path = ckpt_dir / f"{ckpt_prefix}_emergency.pt" save_checkpoint(path, model, optimizer, _state["step"], train_cfg, model_cfg, loader.samples_consumed) print(f"[emergency] saved {path} reason: {reason}", flush=True) def on_sigterm(signum, frame): print("[signal] SIGTERM received", flush=True) emergency_save("SIGTERM") sys.exit(2) signal.signal(signal.SIGTERM, on_sigterm) # ---- the loop ---- model.train() run_start = time.time() last_ckpt_time = time.time() t_window = time.time() tokens_window = 0 loss_window = 0.0 loss_count = 0 try: for step in range(start_step, train_cfg.max_steps): _state["step"] = step lr = cosine_lr(step, train_cfg.warmup_steps, train_cfg.max_steps, train_cfg.peak_lr, train_cfg.min_lr) for g in optimizer.param_groups: g["lr"] = lr optimizer.zero_grad(set_to_none=True) accum_loss = 0.0 for _ in range(train_cfg.grad_accum): x, y = next(data_iter) x = x.to(device, non_blocking=True) y = y.to(device, non_blocking=True) _, loss = model(x, targets=y) (loss / train_cfg.grad_accum).backward() accum_loss += loss.item() accum_loss /= train_cfg.grad_accum grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), train_cfg.grad_clip) optimizer.step() loss_window += accum_loss loss_count += 1 tokens_window += tokens_per_step if (step + 1) % train_cfg.log_every == 0: dt = time.time() - t_window tps = tokens_window / max(dt, 1e-6) avg_loss = loss_window / loss_count elapsed = time.time() - run_start eta = (train_cfg.max_steps - (step + 1)) * (dt / train_cfg.log_every) cost = (elapsed / 3600.0) * train_cfg.cost_per_hour print( f"step {step+1:>6d}/{train_cfg.max_steps} | " f"loss {avg_loss:.4f} | lr {lr:.2e} | grad {grad_norm.item():.2f} | " f"{tps/1e3:.1f}k tok/s | " f"elapsed {fmt_eta(elapsed)} | eta {fmt_eta(eta)} | ${cost:.2f}" ) if use_wandb: import wandb wandb.log({ "loss": avg_loss, "lr": lr, "grad_norm": grad_norm.item(), "tokens_per_sec": tps, "step": step + 1, "elapsed_sec": elapsed, "cost_usd": cost, }) t_window = time.time() tokens_window = 0 loss_window = 0.0 loss_count = 0 if time.time() - last_ckpt_time >= train_cfg.ckpt_interval_sec: path = ckpt_dir / f"{ckpt_prefix}_step{step+1:07d}.pt" save_checkpoint(path, model, optimizer, step + 1, train_cfg, model_cfg, loader.samples_consumed) prune_ckpts(ckpt_dir, ckpt_prefix, train_cfg.ckpt_keep) last_ckpt_time = time.time() print(f"[ckpt] {path.name} saved at step {step+1}, " f"can resume from step {step+1}", flush=True) except _StreamError as e: print(f"[fatal] dataloader: {e}", flush=True) emergency_save(f"_StreamError: {e}") sys.exit(2) except torch.cuda.OutOfMemoryError as e: print(f"[fatal] OOM: {e}", flush=True) emergency_save("OOM") sys.exit(2) except KeyboardInterrupt: print("[interrupt] user", flush=True) emergency_save("KeyboardInterrupt") sys.exit(130) final = ckpt_dir / f"{ckpt_prefix}_final.pt" save_checkpoint(final, model, optimizer, train_cfg.max_steps, train_cfg, model_cfg, loader.samples_consumed) elapsed = time.time() - run_start cost = (elapsed / 3600.0) * train_cfg.cost_per_hour print(f"[done] {final} total elapsed {fmt_eta(elapsed)} approx cost ${cost:.2f}") # ============================================================================= # INFERENCE — pre-allocated KV cache # ============================================================================= def alloc_kv_cache(cfg: ModelConfig, batch_size, max_len, device, dtype): head_dim = cfg.d_model // cfg.n_heads return [ ( torch.zeros(batch_size, cfg.n_kv_heads, max_len, head_dim, device=device, dtype=dtype), torch.zeros(batch_size, cfg.n_kv_heads, max_len, head_dim, device=device, dtype=dtype), ) for _ in range(cfg.n_layers) ] @torch.no_grad() def sample(model, prompt_ids, max_new=200, temperature=0.8, top_k=50, device="cuda"): model.eval() cfg = model.cfg T0 = len(prompt_ids) max_len = min(T0 + max_new, cfg.max_seq_len) kv = alloc_kv_cache(cfg, 1, max_len, device, torch.bfloat16) idx = torch.tensor([prompt_ids], dtype=torch.long, device=device) logits, _ = model(idx, kv_caches=kv, pos=0) next_logits = logits[:, -1, :] out = list(prompt_ids) for i in range(max_new): if T0 + i >= cfg.max_seq_len: break if temperature <= 0: tok = next_logits.argmax(-1, keepdim=True) else: scaled = next_logits / temperature if top_k: v, _ = torch.topk(scaled, min(top_k, scaled.size(-1))) scaled = scaled.masked_fill(scaled < v[:, [-1]], -float("inf")) probs = torch.softmax(scaled.float(), dim=-1) tok = torch.multinomial(probs, 1) tid = tok.item() out.append(tid) if tid == EOT: break logits, _ = model(tok, kv_caches=kv, pos=T0 + i) next_logits = logits[:, -1, :] return out def generate_cmd(args): device = "cuda" ckpt = torch.load(args.checkpoint, map_location=device, weights_only=False) cfg = ModelConfig(**ckpt["model_cfg"]) if "model_cfg" in ckpt else ModelConfig() model = LLaMA(cfg).to(device=device, dtype=torch.bfloat16) sd = ckpt["model"] if any(k.startswith("_orig_mod.") for k in sd): sd = {k.removeprefix("_orig_mod."): v for k, v in sd.items()} model.load_state_dict(sd) prompt = f"\n<|user|>\n{args.prompt}\n<|assistant|>\n" if args.chat else args.prompt ids = ENC.encode_ordinary(prompt) t0 = time.time() out = sample(model, ids, args.max_new, args.temperature, args.top_k, device) dt = time.time() - t0 new_tokens = len(out) - len(ids) print(ENC.decode(out)) print(f"\n[gen] {new_tokens} new tokens in {dt:.2f}s = {new_tokens/dt:.1f} tok/s") # ============================================================================= # CLI # ============================================================================= def build_parser(): p = argparse.ArgumentParser() sub = p.add_subparsers(dest="cmd", required=True) def add_train_args(sp): sp.add_argument("--no-compile", action="store_true", help="Required on Blackwell; compile is never invoked here regardless.") sp.add_argument("--checkpoint", type=str, default=None, help="Initialize weights from this file (for cross-stage seeding).") sp.add_argument("--fresh", action="store_true", help="Ignore any existing ckpt in --ckpt_dir for this stage; start over.") sp.add_argument("--ckpt_dir", type=str, default=None) sp.add_argument("--lr", type=float, default=None) sp.add_argument("--min_lr", type=float, default=None) sp.add_argument("--batch_size", type=int, default=None) sp.add_argument("--grad_accum", type=int, default=None) sp.add_argument("--steps", type=int, default=None) sp.add_argument("--warmup", type=int, default=None) sp.add_argument("--seed", type=int, default=None) sp.add_argument("--no-wandb", action="store_true") add_train_args(sub.add_parser("base")) add_train_args(sub.add_parser("mid")) ft = sub.add_parser("finetune"); add_train_args(ft) ft.add_argument("--skip_mid", action="store_true", help="UX flag: confirms starting SFT directly from a mid checkpoint.") g = sub.add_parser("generate") g.add_argument("--checkpoint", required=True) g.add_argument("--prompt", default="Once upon a time") g.add_argument("--max_new", type=int, default=200) g.add_argument("--temperature", type=float, default=0.8) g.add_argument("--top_k", type=int, default=50) g.add_argument("--chat", action="store_true", help="Wrap prompt in <|user|>/<|assistant|> tags (use with SFT checkpoints).") return p def main(): args = build_parser().parse_args() if args.cmd == "generate": generate_cmd(args) else: train(args.cmd, args) if __name__ == "__main__": main()