""" NoCap-Test Wave-1 trainer — extended fork of train_gpt2.py. Design contract: with all new flags at their defaults, this script reproduces the baseline trainer's semantics (same model, same loss, same schedule, same timing protocol: the clock pauses during validation). Every Wave-1 experiment is opt-in: --precision {bf16,fp16,fp32} fp16 (+GradScaler) enables T4 dev; bf16 = baseline --zloss C PaLM z-loss C*mean(logsumexp(logits)^2) [idea #21] --logit_prior PATH fixed unigram log-prior added to logits [idea #3] --ema_beta B, --ema_start_frac EMA of weights, dual-evaluated [idea #2] --val_stride S also report sliding-window val loss [idea #1] (official disjoint number always reported first) --compile {0,1}, --seed, --device {cuda,cpu}, --run_name, --val_tokens (DEBUG only) Model 'd2' is a debug-size config for CPU smoke tests only. """ import os import sys import csv import json import uuid import math import glob import time import hashlib import re import argparse import contextlib from dataclasses import dataclass import numpy as np import torch from torch import nn import torch.distributed as dist import torch.nn.functional as F from torch.nn.parallel import DistributedDataParallel as DDP from torch.distributed import init_process_group, destroy_process_group with open(sys.argv[0]) as f: code = f.read() TRAINER_SHA256 = hashlib.sha256(code.encode("utf-8")).hexdigest() def _muon_sha256(): path = os.path.join( os.path.dirname(os.path.abspath(__file__)), "muon_opt_split_qkv_cheap_v1.py" ) if not os.path.isfile(path): return "" with open(path, "rb") as f: return hashlib.sha256(f.read()).hexdigest() def _sha256_file(path): h = hashlib.sha256() with open(path, "rb") as f: for chunk in iter(lambda: f.read(8 << 20), b""): h.update(chunk) return h.hexdigest() def _validate_checkpoint_ancestors(value): """Return a defensive copy of an exact checkpoint hash chain.""" if not isinstance(value, list): raise RuntimeError("checkpoint_ancestor_sha256 must be a list") out = [] for digest in value: if not isinstance(digest, str) or re.fullmatch(r"[0-9a-f]{64}", digest) is None: raise RuntimeError("checkpoint ancestry contains a non-SHA256 value") if digest in out: raise RuntimeError("checkpoint ancestry contains a duplicate/cycle") out.append(digest) return out def _resume_recipe(arg_dict): """Canonical mathematical/data recipe; storage location may differ.""" ignored = {"resume", "run_name", "output_dir", "allow_legacy_resume"} return {k: v for k, v in arg_dict.items() if k not in ignored} # ----------------------------------------------------------------------------- # Model (identical to baseline, plus optional logit prior buffer + zloss term) class Rotary(torch.nn.Module): def __init__(self, dim, base=10000): super().__init__() inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim)) self.register_buffer("inv_freq", inv_freq) self.seq_len_cached = None self.cos_cached = None self.sin_cached = None def forward(self, x): seq_len = x.shape[1] if seq_len != self.seq_len_cached: self.seq_len_cached = seq_len t = torch.arange(seq_len, device=x.device).type_as(self.inv_freq) freqs = torch.outer(t, self.inv_freq).to(x.device) self.cos_cached = freqs.cos() self.sin_cached = freqs.sin() return self.cos_cached[None, :, None, :], self.sin_cached[None, :, None, :] def apply_rotary_emb(x, cos, sin): assert x.ndim == 4 d = x.shape[3] // 2 x1 = x[..., :d] x2 = x[..., d:] y1 = x1 * cos + x2 * sin y2 = x1 * (-sin) + x2 * cos return torch.cat([y1, y2], 3) def rmsnorm(x0, eps=1e-6): x = x0.float() x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps) return x.type_as(x0) class CausalSelfAttention(nn.Module): def __init__(self, config): super().__init__() self.n_head = config.n_head self.n_embd = config.n_embd self.head_dim = self.n_embd // self.n_head assert self.n_embd % self.n_head == 0 self.c_attn = nn.Linear(self.n_embd, 3 * self.n_embd, bias=False) self.c_proj = nn.Linear(self.n_embd, self.n_embd, bias=False) self.rotary = Rotary(self.head_dim) def forward(self, x): B, T, C = x.size() qkv = self.c_attn(x) q, k, v = qkv.split(self.n_embd, dim=2) k = k.view(B, T, self.n_head, self.head_dim) q = q.view(B, T, self.n_head, self.head_dim) v = v.view(B, T, self.n_head, self.head_dim) cos, sin = self.rotary(q) q = apply_rotary_emb(q, cos, sin) k = apply_rotary_emb(k, cos, sin) y = F.scaled_dot_product_attention( q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), is_causal=True ) y = y.transpose(1, 2).contiguous().view(B, T, C) y = self.c_proj(y) return y class MLP(nn.Module): """Baseline 4x GELU MLP, or iso-FLOP SwiGLU (idea from Wave-3 research: d_ff = 8d/3 gives exactly the same params/FLOPs as the 4x GELU MLP).""" def __init__(self, config, mlp_type="gelu"): super().__init__() self.mlp_type = mlp_type if mlp_type == "swiglu": d_ff = int(8 * config.n_embd / 3) self.c_gate = nn.Linear(config.n_embd, d_ff, bias=False) self.c_up = nn.Linear(config.n_embd, d_ff, bias=False) self.c_proj = nn.Linear(d_ff, config.n_embd, bias=False) else: self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=False) self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=False) def forward(self, x): if self.mlp_type == "swiglu": return self.c_proj(F.silu(self.c_gate(x)) * self.c_up(x)) return self.c_proj(F.gelu(self.c_fc(x))) class CanonConv(nn.Module): """Canon-C: zero-initialized causal depthwise conv residual (kernel 4). Zero init => training starts EXACTLY at the baseline function.""" def __init__(self, dim, kernel=4): super().__init__() self.kernel = kernel # fork the RNG: constructing Conv1d consumes random draws, which would # silently shift every subsequent layer's init away from the baseline # stream (caught by the step-0 equality smoke test). Weights are zeroed # anyway, so the draws are meaningless — keep the global stream intact. with torch.random.fork_rng(devices=[]): torch.manual_seed(0) self.conv = nn.Conv1d(dim, dim, kernel, groups=dim, bias=False) nn.init.zeros_(self.conv.weight) def forward(self, x): # x: (B, T, C) y = F.pad(x.transpose(1, 2), (self.kernel - 1, 0)) # causal left-pad return self.conv(y).transpose(1, 2) class Block(nn.Module): def __init__(self, config, mlp_type="gelu", canon=False, parallel=False): super().__init__() self.attn = CausalSelfAttention(config) self.mlp = MLP(config, mlp_type) self.attn_scale = 1 / math.sqrt(2 * config.n_layer) self.parallel = parallel self.canon_a = CanonConv(config.n_embd) if canon else None # pre-MLP canon site only exists in the serial block layout self.canon_m = CanonConv(config.n_embd) if (canon and not parallel) else None def forward(self, x): if self.canon_a is not None: x = x + self.canon_a(x) if self.parallel: xn = rmsnorm(x) x = x + self.attn_scale * self.attn(xn) + self.mlp(xn) else: x = x + self.attn_scale * self.attn(rmsnorm(x)) if self.canon_m is not None: x = x + self.canon_m(x) x = x + self.mlp(rmsnorm(x)) return x @dataclass class GPTConfig: vocab_size: int = 50257 n_layer: int = 12 n_head: int = 12 n_embd: int = 768 class GPT(nn.Module): def __init__(self, config, mtp=False, mtp_chunk=256, ls_hi_q=0.98, ls_hi_w=0.3, ls_lo_q=0.10, ls_lo_w=0.5, mlp_type="gelu", canon=False, parallel_block=False, trainable_logit_prior=False): super().__init__() self.config = config self.transformer = nn.ModuleDict( dict( wte=nn.Embedding(config.vocab_size, config.n_embd), h=nn.ModuleList([ Block(config, mlp_type=mlp_type, canon=canon, parallel=parallel_block) for _ in range(config.n_layer)]), ) ) self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False) self.transformer.wte.weight = self.lm_head.weight # weight tying # A vocabulary intercept initialized from the FineWeb unigram prior. # Default remains a fixed buffer. The opt-in trainable form adds one # small convex multinomial-logistic block (50,257 scalar parameters). prior = torch.zeros(config.vocab_size) if trainable_logit_prior: self.logit_prior = nn.Parameter(prior) else: self.register_buffer("logit_prior", prior) # --- idea #11: sequential MTP module (DeepSeek-style, contiguous chunk). # Train-only; never used at eval, so the base model stays a valid # probability model. lambda is carried as a buffer so its per-step decay # never triggers recompilation. self.mtp_chunk = mtp_chunk if mtp: self.mtp_block = Block(config) self.mtp_merge = nn.Linear(2 * config.n_embd, config.n_embd, bias=False) else: self.mtp_block = None self.mtp_merge = None self.register_buffer("mtp_lambda_t", torch.zeros(())) # --- idea #13: online token-loss shaping constants self.ls_hi_q, self.ls_hi_w = ls_hi_q, ls_hi_w self.ls_lo_q, self.ls_lo_w = ls_lo_q, ls_lo_w self.register_buffer("shape_gate", torch.zeros(())) # 0=off, 1=on def forward(self, idx, targets=None, return_logits=True, zloss=0.0, use_shape=False, mtp_offset=-1): b, t = idx.size() x = self.transformer.wte(idx) for block in self.transformer.h: x = block(x) x = rmsnorm(x) if targets is not None: logits = self.lm_head(x) logits = logits + self.logit_prior # zeros unless a prior was loaded tflat = targets.view(-1) ce_tok = F.cross_entropy( logits.view(-1, logits.size(-1)), tflat, ignore_index=-1, reduction="none", ) valid = (tflat >= 0).float() if use_shape: cef = ce_tok.detach().float() qhi = torch.quantile(cef, self.ls_hi_q) qlo = torch.quantile(cef, self.ls_lo_q) w = torch.where(cef > qhi, torch.full_like(cef, self.ls_hi_w), torch.ones_like(cef)) w = w * torch.where(cef < qlo, torch.full_like(cef, self.ls_lo_w), torch.ones_like(cef)) w = 1.0 + self.shape_gate * (w - 1.0) # gate=0 -> plain CE w = w * valid loss = (w * ce_tok).sum() / w.sum().clamp(min=1.0) else: loss = (ce_tok * valid).sum() / valid.sum().clamp(min=1.0) if zloss > 0.0: lse = torch.logsumexp(logits.float(), dim=-1) # (b, t) mask = (targets >= 0).float() zterm = (lse.pow(2) * mask).sum() / mask.sum().clamp(min=1.0) loss = loss + zloss * zterm if self.mtp_block is not None and mtp_offset >= 0: # predict token t+2: combine hidden at p with embedding of token # p+1 (= idx[:, p+1]), one extra block, shared (tied) head. o, C = mtp_offset, self.mtp_chunk h = x[:, o:o + C, :] e = self.transformer.wte(idx[:, o + 1:o + 1 + C]) h2 = self.mtp_block(self.mtp_merge(torch.cat([h, e], dim=-1))) h2 = rmsnorm(h2) lg2 = self.lm_head(h2) + self.logit_prior tgt2 = targets[:, o + 1:o + 1 + C] ce2 = F.cross_entropy( lg2.reshape(-1, lg2.size(-1)), tgt2.reshape(-1), ignore_index=-1, ) loss = loss + self.mtp_lambda_t * ce2 else: logits = self.lm_head(x[:, [-1], :]) logits = logits + self.logit_prior loss = None if not return_logits: logits = None return logits, loss def configure_optimizers(self, weight_decay, learning_rate, betas, device_type, optimizer_name="adamw", muon_lr=None, muon_weight_decay=None, muon_ns_steps=5, muon_momentum=0.95, prior_lr=None, muon_split_qkv=False, match_joint_frobenius=True): if optimizer_name == "adamw": prior_params = [] model_params = [] for name, param in self.named_parameters(): (prior_params if name == "logit_prior" else model_params).append(param) groups = [{"params": model_params, "weight_decay": weight_decay}] if prior_params: groups.append({ "params": prior_params, "lr": learning_rate if prior_lr is None else prior_lr, "weight_decay": 0.0, "is_logit_prior": True, }) return torch.optim.AdamW(groups, lr=learning_rate, betas=betas) if optimizer_name == "muon": # Attribution: Keller Jordan Muon. Embed/head/1D → AdamW. from muon_opt_split_qkv_cheap_v1 import build_muon_adam_optimizer mlr = learning_rate if muon_lr is None else muon_lr return build_muon_adam_optimizer( self, lr_adam=learning_rate, lr_muon=mlr, weight_decay=weight_decay, betas=betas, muon_weight_decay=muon_weight_decay, muon_ns_steps=muon_ns_steps, muon_momentum=muon_momentum, prior_lr=prior_lr, muon_split_qkv=muon_split_qkv, match_joint_frobenius=match_joint_frobenius, ) raise ValueError("unknown optimizer %s" % optimizer_name) # ----------------------------------------------------------------------------- # Training data loader (identical to baseline) def _peek_data_shard(filename): with open(filename, "rb") as f: header = np.frombuffer(f.read(256 * 4), dtype=np.int32) if header[0] != 20240520: print("ERROR: magic number mismatch in the data .bin file!") exit(1) assert header[1] == 1, "unsupported version" return header[2] def _load_data_shard(filename): with open(filename, "rb") as f: header = np.frombuffer(f.read(256 * 4), dtype=np.int32) assert header[0] == 20240520, "magic number mismatch in the data .bin file" assert header[1] == 1, "unsupported version" ntok = header[2] tokens = np.frombuffer(f.read(), dtype=np.uint16) assert len(tokens) == ntok, "number of tokens read does not match header?" return tokens class DistributedDataLoader: def __init__(self, filename_pattern, B, T, process_rank, num_processes, device): self.process_rank = process_rank self.num_processes = num_processes self.B = B self.T = T self.device = device self.files = sorted(glob.glob(filename_pattern)) assert len(self.files) > 0, "no files match pattern %s" % filename_pattern ntok_total = np.int64(0) for fname in self.files: shard_ntok = _peek_data_shard(fname) assert shard_ntok >= num_processes * B * T + 1 ntok_total += shard_ntok self.ntok_total = ntok_total print0( "DataLoader: total number of tokens: {:,} across {} files".format( ntok_total, len(self.files) ) ) self.reset() def reset(self): self.current_shard = 0 self.current_position = self.process_rank * self.B * self.T self.tokens = _load_data_shard(self.files[self.current_shard]) def advance(self): self.current_shard = (self.current_shard + 1) % len(self.files) self.current_position = self.process_rank * self.B * self.T self.tokens = _load_data_shard(self.files[self.current_shard]) def next_batch(self): B, T = self.B, self.T buf = self.tokens[self.current_position : self.current_position + B * T + 1] buf = torch.tensor(buf.astype(np.int32), dtype=torch.long) x = (buf[:-1]).view(B, T) y = (buf[1:]).view(B, T) self.current_position += B * T * self.num_processes if self.current_position + (B * T * self.num_processes + 1) > len(self.tokens): self.advance() return x.to(self.device), y.to(self.device) def state_dict(self): """Exact stream cursor for preemption-safe single-process resumes.""" return { "current_shard": int(self.current_shard), "current_position": int(self.current_position), "files": list(self.files), "B": int(self.B), "T": int(self.T), "process_rank": int(self.process_rank), "num_processes": int(self.num_processes), } def load_state_dict(self, state): expected = { "files": list(self.files), "B": int(self.B), "T": int(self.T), "process_rank": int(self.process_rank), "num_processes": int(self.num_processes), } for key, value in expected.items(): if state.get(key) != value: raise RuntimeError( "train-loader resume mismatch for %s: checkpoint=%r current=%r" % (key, state.get(key), value) ) shard = int(state["current_shard"]) position = int(state["current_position"]) if not (0 <= shard < len(self.files)): raise RuntimeError("invalid checkpoint train shard %d" % shard) self.current_shard = shard self.current_position = position self.tokens = _load_data_shard(self.files[self.current_shard]) # ----------------------------------------------------------------------------- # Validation: exact replication of the baseline's disjoint protocol, plus the # optional boundary-aware sliding-window evaluation (idea #1). def load_val_tokens(filename_pattern, n_tokens): files = sorted(glob.glob(filename_pattern)) assert len(files) > 0, "no val files match %s" % filename_pattern toks = _load_data_shard(files[0]) assert len(toks) >= n_tokens + 1, "val shard smaller than requested val tokens" return torch.tensor(toks[: n_tokens + 1].astype(np.int64), dtype=torch.long) @torch.no_grad() def evaluate(model, val_tokens, T, batch_size, device, mode="disjoint", stride=None): """Mean CE over exactly (len(val_tokens)-1) target tokens. disjoint: contiguous non-overlapping windows of length T (baseline protocol). sliding : overlapping windows (stride s); first window scores all T targets, each later window scores only its last s targets, so every target is scored exactly once with >= T-s tokens of real context. Still a valid probability model over the entire val set. """ model.eval() n_targets = val_tokens.numel() - 1 windows = [] # (start, n_scored_from_end) if mode == "disjoint": assert n_targets % T == 0 for start in range(0, n_targets, T): windows.append((start, T)) else: s = stride assert s and 0 < s <= T and (n_targets - T) % s == 0 windows.append((0, T)) for start in range(s, n_targets - T + s, s): windows.append((start, s)) total_loss = 0.0 total_count = 0 for i in range(0, len(windows), batch_size): chunk = windows[i : i + batch_size] xs, ys = [], [] for start, n_scored in chunk: x = val_tokens[start : start + T] y = val_tokens[start + 1 : start + T + 1].clone() if n_scored < T: y[: T - n_scored] = -1 xs.append(x) ys.append(y) xb = torch.stack(xs).to(device) yb = torch.stack(ys).to(device) logits, _ = model(xb, yb, return_logits=True) loss_sum = F.cross_entropy( logits.view(-1, logits.size(-1)), yb.view(-1), ignore_index=-1, reduction="sum", ) total_loss += loss_sum.item() total_count += int((yb >= 0).sum().item()) assert total_count == n_targets, "scored %d != expected %d" % ( total_count, n_targets, ) model.train() return total_loss / total_count # ----------------------------------------------------------------------------- VAL_TOKENS = 1_048_576 # do not change for real runs — this IS the metric def print0(*args, **kwargs): if int(os.environ.get("RANK", 0)) == 0: print(*args, **kwargs) if __name__ == "__main__": parser = argparse.ArgumentParser() # baseline args parser.add_argument("--input_bin", type=str, default="data/fineweb10B/fineweb_train_*.bin") parser.add_argument("--input_val_bin", type=str, default="data/fineweb10B/fineweb_val_*.bin") parser.add_argument("--output_dir", type=str, default="") parser.add_argument("--model", type=str, default="d12") parser.add_argument("--batch_size", type=int, default=4) parser.add_argument("--grad_accumulation_steps", type=int, default=1) parser.add_argument("--sequence_length", type=int, default=64) parser.add_argument("--num_iterations", type=int, default=10) parser.add_argument("--learning_rate", type=float, default=1e-4) parser.add_argument("--warmup_iters", type=int, default=0) parser.add_argument("--warmdown_iters", type=int, default=0) parser.add_argument("--weight_decay", type=float, default=0.0) parser.add_argument("--val_loss_every", type=int, default=0) parser.add_argument("--val_batch_size", type=int, default=16) parser.add_argument( "--no_validation_data", type=int, default=0, help=( "Fail-closed train-only mode for curation anchors (0/1). When 1, " "requires --val_loss_every 0, requires an empty --input_val_bin, " "and never opens validation data." ), ) parser.add_argument( "--anchor_prepare_manifest_sha256", type=str, default="", help=( "Cross-fit anchors only: SHA-256 of the immutable preparation " "manifest. Must be paired with --anchor_corpus_sha256." ), ) parser.add_argument( "--anchor_corpus_sha256", type=str, default="", help=( "Cross-fit anchors only: canonical SHA-256 identity of the exact " "fold corpus. Must be paired with --anchor_prepare_manifest_sha256." ), ) parser.add_argument("--save_every", type=int, default=5000) parser.add_argument("--log_wandb", action="store_true") # wave-1 args parser.add_argument("--precision", type=str, default="bf16", choices=["bf16", "fp16", "fp32"]) parser.add_argument("--zloss", type=float, default=0.0) parser.add_argument("--logit_prior", type=str, default="") parser.add_argument( "--trainable_logit_prior", type=int, default=0, help="Optimize the loaded unigram vocabulary intercept (0=fixed, 1=trainable)", ) parser.add_argument( "--prior_lr", type=float, default=None, help="Adam LR for a trainable logit prior (default: --learning_rate)", ) parser.add_argument("--ema_beta", type=float, default=0.0) parser.add_argument("--ema_start_frac", type=float, default=0.3) parser.add_argument("--val_stride", type=int, default=0) parser.add_argument("--compile", type=int, default=1) parser.add_argument("--seed", type=int, default=1337) parser.add_argument("--device", type=str, default="cuda", choices=["cuda", "cpu"]) parser.add_argument("--run_name", type=str, default="") parser.add_argument("--save_final", type=int, default=1) parser.add_argument("--val_tokens", type=int, default=VAL_TOKENS, help="DEBUG ONLY — changing this invalidates the metric") # wave-2 args parser.add_argument("--mtp_lambda", type=float, default=0.0, help="idea #11: MTP aux loss weight (0 = module off)") parser.add_argument("--mtp_chunk", type=int, default=256) parser.add_argument("--mtp_decay_frac", type=float, default=0.75, help="lambda decays linearly to 0 by this fraction of training") parser.add_argument("--loss_shape", type=int, default=0, help="idea #13: online token-loss shaping (0/1)") parser.add_argument("--ls_hi_q", type=float, default=0.98) parser.add_argument("--ls_hi_w", type=float, default=0.3) parser.add_argument("--ls_lo_q", type=float, default=0.10) parser.add_argument("--ls_lo_w", type=float, default=0.5) parser.add_argument("--ls_start_frac", type=float, default=0.15) parser.add_argument("--optimizer", type=str, default="adamw", choices=["adamw", "muon"], help="adamw (default) or muon hybrid (2D Muon + AdamW embed/head)") parser.add_argument("--muon_lr", type=float, default=None, help="Muon spectral LR for 2D weights (default: same as --learning_rate; " "try 0.02 for pure-Muon-scale proxies)") parser.add_argument( "--muon_weight_decay", type=float, default=None, help="Weight decay for Muon matrix groups only (default: inherit --weight_decay)", ) parser.add_argument( "--muon_ns_steps", type=int, default=5, help="Newton-Schulz iterations per Muon update (historical default: 5)", ) parser.add_argument( "--muon_momentum", type=float, default=0.95, help="Muon heavy-ball momentum coefficient (historical default: 0.95)", ) parser.add_argument( "--muon_split_qkv", type=int, default=0, help=( "Apply frozen NS5 independently to fused Q/K/V row blocks " "(0=historical joint map, 1=split)" ), ) parser.add_argument( "--muon_split_qkv_match_joint", type=int, default=1, help=( "When split-QKV is on: 1=dual-polar online joint norm match (taxed); " "0=fixed sqrt(3) scale only (timing-cheap)" ), ) parser.add_argument( "--muon_qkv_diagnostic_steps", type=str, default="", help=( "Comma-separated update indices for symmetric off-clock joint/split " "QKV norm diagnostics (empty disables diagnostics)" ), ) parser.add_argument( "--muon_to_adam_step", type=int, default=-1, help=( "Continuation experiment: first update that hands hidden matrix " "groups from Muon to fresh AdamW state (-1 disables)" ), ) parser.add_argument( "--muon_to_adam_lr", type=float, default=None, help="Base LR for handed-off matrix AdamW groups (default: --learning_rate)", ) parser.add_argument( "--single_process_fastpath", type=int, default=0, help="On WORLD_SIZE=1, skip DDP wrapping and logging all-reduce (0/1)", ) parser.add_argument( "--log_every", type=int, default=1, help="Write/print train loss every N steps; validation timing remains exact", ) parser.add_argument("--resume", type=str, default="", help="Path to ckpt_step*.pt to resume (restores model, optim, step, train_time_ms). " "Does not change the training recipe — for preemption recovery only.") parser.add_argument( "--allow_legacy_resume", type=int, default=0, help="UNSAFE exploratory escape hatch for old checkpoints without exact resume metadata", ) # wave-3 architecture flags (defaults = exact baseline architecture) parser.add_argument("--mlp", type=str, default="gelu", choices=["gelu", "swiglu"], help="idea #9: iso-FLOP SwiGLU MLP") parser.add_argument("--canon", type=int, default=0, help="idea #18: zero-init causal depthwise-conv residuals") parser.add_argument("--parallel_block", type=int, default=0, help="parallel attention+FFN block (GPT-J style)") parser.add_argument( "--sdp_backend", type=str, default="auto", choices=["auto", "flash", "mem", "math"], help="Force CUDA SDPA backend for F.scaled_dot_product_attention " "(auto=PyTorch default; flash|mem|math force one path). " "Board-v3 S1 systems lever for t_step.") args = parser.parse_args() B, T = args.batch_size, args.sequence_length assert args.model in {"d2", "d12", "d24", "d36", "d48"} assert args.muon_ns_steps >= 1 if args.muon_split_qkv not in (0, 1): raise RuntimeError("--muon_split_qkv must be 0 or 1") if args.muon_split_qkv_match_joint not in (0, 1): raise RuntimeError("--muon_split_qkv_match_joint must be 0 or 1") if args.muon_split_qkv and args.optimizer != "muon": raise RuntimeError("--muon_split_qkv=1 requires --optimizer muon") if args.muon_split_qkv_match_joint == 0 and not args.muon_split_qkv: raise RuntimeError( "--muon_split_qkv_match_joint=0 requires --muon_split_qkv=1" ) try: muon_qkv_diagnostic_steps = tuple( sorted( { int(value) for value in args.muon_qkv_diagnostic_steps.split(",") if value != "" } ) ) except ValueError as exc: raise RuntimeError( "--muon_qkv_diagnostic_steps must be comma-separated integers" ) from exc if any( step < 0 or step >= args.num_iterations for step in muon_qkv_diagnostic_steps ): raise RuntimeError( "Muon QKV diagnostic steps must be in [0, num_iterations)" ) if muon_qkv_diagnostic_steps and args.optimizer != "muon": raise RuntimeError("Muon QKV diagnostics require --optimizer muon") assert args.log_every >= 1 if args.muon_to_adam_step >= 0: if args.optimizer != "muon": raise RuntimeError("--muon_to_adam_step requires --optimizer muon") if args.muon_to_adam_step >= args.num_iterations: raise RuntimeError( "--muon_to_adam_step must be smaller than --num_iterations" ) if args.muon_to_adam_lr is not None and args.muon_to_adam_lr <= 0: raise RuntimeError("--muon_to_adam_lr must be positive") if args.no_validation_data not in (0, 1): raise RuntimeError("--no_validation_data must be 0 or 1") if args.no_validation_data and args.val_loss_every != 0: raise RuntimeError( "--no_validation_data=1 requires --val_loss_every 0" ) if args.no_validation_data and args.input_val_bin: raise RuntimeError( "--no_validation_data=1 requires an empty --input_val_bin" ) anchor_hashes = ( args.anchor_prepare_manifest_sha256, args.anchor_corpus_sha256, ) if bool(anchor_hashes[0]) != bool(anchor_hashes[1]): raise RuntimeError( "anchor preparation and corpus SHA-256 values must be supplied together" ) for digest in anchor_hashes: if digest and re.fullmatch(r"[0-9a-f]{64}", digest) is None: raise RuntimeError("anchor provenance values must be lowercase SHA-256 hex") if anchor_hashes[0] and not args.no_validation_data: raise RuntimeError( "anchor provenance is accepted only with --no_validation_data=1" ) if args.trainable_logit_prior and not args.logit_prior: raise RuntimeError("--trainable_logit_prior requires --logit_prior initialization") # allow running without torchrun (single process) if "RANK" not in os.environ: os.environ["RANK"] = "0" os.environ["LOCAL_RANK"] = "0" os.environ["WORLD_SIZE"] = "1" os.environ.setdefault("MASTER_ADDR", "127.0.0.1") os.environ.setdefault("MASTER_PORT", str(29500 + os.getpid() % 1000)) use_cuda = args.device == "cuda" if use_cuda: assert torch.cuda.is_available(), "--device cuda requires a GPU" backend = "nccl" if use_cuda else "gloo" init_process_group(backend=backend) ddp_rank = int(os.environ["RANK"]) ddp_local_rank = int(os.environ["LOCAL_RANK"]) ddp_world_size = int(os.environ["WORLD_SIZE"]) assert args.grad_accumulation_steps % ddp_world_size == 0 args.grad_accumulation_steps //= ddp_world_size single_process_fastpath = bool(args.single_process_fastpath) if single_process_fastpath and ddp_world_size != 1: raise RuntimeError("--single_process_fastpath requires WORLD_SIZE=1") device = ("cuda:%d" % ddp_local_rank) if use_cuda else "cpu" if use_cuda: torch.cuda.set_device(device) master_process = ddp_rank == 0 print0("Running pytorch %s on %s" % (torch.version.__version__, device)) # Optional SDPA backend pin (board_v3 S1). Must run after CUDA is ready. sdp = getattr(args, "sdp_backend", "auto") or "auto" if use_cuda and sdp != "auto": # Force a single backend: disable the others so SDPA cannot silently fall back. torch.backends.cuda.enable_flash_sdp(sdp == "flash") torch.backends.cuda.enable_mem_efficient_sdp(sdp == "mem") torch.backends.cuda.enable_math_sdp(sdp == "math") print0( "sdp_backend=%s (flash=%s mem=%s math=%s)" % ( sdp, torch.backends.cuda.flash_sdp_enabled(), torch.backends.cuda.mem_efficient_sdp_enabled(), torch.backends.cuda.math_sdp_enabled(), ) ) else: print0("sdp_backend=auto (PyTorch default selection)") torch.manual_seed(args.seed) if use_cuda: torch.cuda.manual_seed_all(args.seed) if args.precision == "bf16": if use_cuda and not torch.cuda.is_bf16_supported(): raise SystemExit( "This GPU has no bf16 support (e.g. T4/Turing). Use --precision fp16." ) ctx = torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16) if use_cuda else contextlib.nullcontext() elif args.precision == "fp16": assert use_cuda, "fp16 autocast requires cuda" ctx = torch.amp.autocast(device_type="cuda", dtype=torch.float16) else: ctx = contextlib.nullcontext() scaler = torch.amp.GradScaler("cuda", enabled=(args.precision == "fp16" and use_cuda)) if args.log_wandb and master_process: import wandb import datetime start_time = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") wandb.init(project="benchmark_gpt2", name="gpt2-%s %s" % (args.model, start_time)) wandb.config.update(args) tokens_per_iter = B * T * ddp_world_size * args.grad_accumulation_steps print0("tokens per iteration: {:,}".format(tokens_per_iter)) train_loader = DistributedDataLoader(args.input_bin, B, T, ddp_rank, ddp_world_size, device) if args.no_validation_data: val_tokens_buf = None print0("validation data disabled (train-only anchor mode)") else: val_tokens_buf = load_val_tokens(args.input_val_bin, args.val_tokens) num_vocab = 50257 model_config = { "d2": GPTConfig(vocab_size=num_vocab, n_layer=2, n_head=2, n_embd=64), "d12": GPTConfig(vocab_size=num_vocab, n_layer=12, n_head=12, n_embd=768), "d24": GPTConfig(vocab_size=num_vocab, n_layer=24, n_head=16, n_embd=1024), "d36": GPTConfig(vocab_size=num_vocab, n_layer=36, n_head=20, n_embd=1280), "d48": GPTConfig(vocab_size=num_vocab, n_layer=48, n_head=25, n_embd=1600), }[args.model] gpt = GPT(model_config, mtp=(args.mtp_lambda > 0), mtp_chunk=args.mtp_chunk, ls_hi_q=args.ls_hi_q, ls_hi_w=args.ls_hi_w, ls_lo_q=args.ls_lo_q, ls_lo_w=args.ls_lo_w, mlp_type=args.mlp, canon=bool(args.canon), parallel_block=bool(args.parallel_block), trainable_logit_prior=bool(args.trainable_logit_prior)) n_params = sum(p.numel() for p in gpt.parameters()) print0("model params: {:,} (mlp={}, canon={}, parallel={})".format( n_params, args.mlp, bool(args.canon), bool(args.parallel_block))) # MTP chunk must fit: offsets cycle over a bounded set (compile-friendly) if args.mtp_lambda > 0: assert args.mtp_chunk + 2 <= T, "mtp_chunk too large for sequence length" mtp_offsets = sorted(set( min(o, T - args.mtp_chunk - 1) for o in range(0, T, max(args.mtp_chunk, 1)))) print0("MTP on: chunk %d, offsets %s, lambda %.3g decaying to 0 by step %d" % (args.mtp_chunk, mtp_offsets, args.mtp_lambda, int(args.mtp_decay_frac * args.num_iterations))) else: mtp_offsets = None if args.logit_prior: prior = torch.load(args.logit_prior, map_location="cpu") assert prior.shape == (num_vocab,), "prior must be (%d,)" % num_vocab with torch.no_grad(): gpt.logit_prior.copy_(prior.float()) print0("loaded %s logit prior from %s (min %.2f max %.2f)" % ( "trainable" if args.trainable_logit_prior else "fixed", args.logit_prior, prior.min().item(), prior.max().item())) gpt = gpt.train().to(device) if args.compile and use_cuda: import torch._inductor.config as iconfig if hasattr(iconfig, "coordinate_descent_tuning"): iconfig.coordinate_descent_tuning = True print0("compiling the model...") model = torch.compile(gpt) else: model = gpt if single_process_fastpath: print0("single_process_fastpath=1 (DDP wrapper/all-reduce disabled)") else: model = DDP(model, device_ids=[ddp_local_rank] if use_cuda else None) raw_model = gpt # clean, uncompiled reference sharing the same parameters optimizer = raw_model.configure_optimizers( weight_decay=args.weight_decay, learning_rate=args.learning_rate, betas=(0.9, 0.95), device_type=device, optimizer_name=getattr(args, "optimizer", "adamw"), muon_lr=getattr(args, "muon_lr", None), muon_weight_decay=getattr(args, "muon_weight_decay", None), muon_ns_steps=getattr(args, "muon_ns_steps", 5), muon_momentum=getattr(args, "muon_momentum", 0.95), prior_lr=getattr(args, "prior_lr", None), muon_split_qkv=bool(getattr(args, "muon_split_qkv", 0)), match_joint_frobenius=bool(getattr(args, "muon_split_qkv_match_joint", 1)), ) # Per-group base LR so schedule scales Muon and Adam groups independently for g in optimizer.param_groups: g["base_lr"] = g["lr"] print0("optimizer=%s groups=%d" % ( getattr(args, "optimizer", "adamw"), len(optimizer.param_groups))) qkv_diagnostic_parameter = None qkv_diagnostic_group = None if muon_qkv_diagnostic_steps: qkv_name = "transformer.h.0.attn.c_attn.weight" qkv_diagnostic_parameter = dict(raw_model.named_parameters()).get(qkv_name) if qkv_diagnostic_parameter is None: raise RuntimeError("Muon QKV diagnostic parameter is missing: %s" % qkv_name) for group in optimizer.param_groups: if group.get("use_muon", False) and any( parameter is qkv_diagnostic_parameter for parameter in group["params"] ): qkv_diagnostic_group = group break if qkv_diagnostic_group is None: raise RuntimeError("Muon QKV diagnostic parameter is not in a Muon group") print0( "Muon QKV diagnostics: parameter=%s steps=%s" % (qkv_name, muon_qkv_diagnostic_steps) ) # ---- EMA of weights (idea #2). fp32 buffers, updated post-optimizer-step. ema_active = args.ema_beta > 0.0 ema_start_step = int(args.num_iterations * args.ema_start_frac) ema_params = None named_params = [p for p in raw_model.parameters()] resume_step = 0 resume_train_time_ms = 0.0 pending_x = pending_y = None checkpoint_ancestor_sha256 = [] if getattr(args, "resume", ""): print0("resuming from %s" % args.resume) ck = torch.load(args.resume, map_location="cpu") legacy = not all( key in ck for key in ( "optimizer", "args", "train_loader", "pending_x", "pending_y", "trainer_sha256", "muon_sha256", "checkpoint_ancestor_sha256", ) ) if legacy and not args.allow_legacy_resume: raise RuntimeError( "legacy/incomplete checkpoint cannot be used for an exact resume; " "start clean (or pass --allow_legacy_resume 1 for exploratory use only)" ) if not legacy: expected_recipe = _resume_recipe(ck["args"]) current_recipe = _resume_recipe(vars(args)) if expected_recipe != current_recipe: changed = sorted( key for key in set(expected_recipe) | set(current_recipe) if expected_recipe.get(key) != current_recipe.get(key) ) raise RuntimeError( "resume recipe mismatch for keys %s" % ", ".join(changed) ) if ck["trainer_sha256"] != TRAINER_SHA256: raise RuntimeError( "trainer hash mismatch on resume: checkpoint=%s current=%s" % (ck["trainer_sha256"], TRAINER_SHA256) ) current_muon_sha = _muon_sha256() if ck["muon_sha256"] != current_muon_sha: raise RuntimeError( "Muon hash mismatch on resume: checkpoint=%s current=%s" % (ck["muon_sha256"], current_muon_sha) ) checkpoint_ancestor_sha256 = _validate_checkpoint_ancestors( ck["checkpoint_ancestor_sha256"] ) resumed_checkpoint_sha256 = _sha256_file(args.resume) if resumed_checkpoint_sha256 in checkpoint_ancestor_sha256: raise RuntimeError( "resumed checkpoint already appears in its own ancestry" ) checkpoint_ancestor_sha256.append(resumed_checkpoint_sha256) raw_model.load_state_dict(ck["model"]) if "optimizer" in ck: try: optimizer.load_state_dict(ck["optimizer"]) except Exception as e: raise RuntimeError("optimizer state load failed on exact resume") from e elif not args.allow_legacy_resume: raise RuntimeError("checkpoint has no optimizer state") if ck.get("ema") is not None and ema_active: if isinstance(ck["ema"], dict): name_list = [n for n, _ in raw_model.named_parameters()] ema_params = [ck["ema"][n].detach().float().to(device) for n in name_list if n in ck["ema"]] if len(ema_params) != len(name_list): print0("WARN: ema resume size mismatch — dropping ema") ema_params = None else: ema_params = [t.float() for t in ck["ema"]] resume_step = int(ck.get("step", 0)) resume_train_time_ms = float(ck.get("training_time_ms", 0.0)) loader_state = ck.get("train_loader") if loader_state is not None: if ck.get("pending_x") is None or ck.get("pending_y") is None: raise RuntimeError( "checkpoint has a train-loader cursor but no pending batch" ) train_loader.load_state_dict(loader_state) pending_x = ck["pending_x"].to(device) pending_y = ck["pending_y"].to(device) print0( "resume: restored train shard=%d position=%d plus pending batch" % (loader_state["current_shard"], loader_state["current_position"]) ) else: if not args.allow_legacy_resume: raise RuntimeError("checkpoint has no exact train-loader state") print0("UNSAFE legacy resume: FineWeb stream restarts at shard 0") print0("resume: start_step=%d train_time_ms=%.1f" % (resume_step, resume_train_time_ms)) if pending_x is not None: x, y = pending_x, pending_y else: x, y = train_loader.next_batch() def ema_update(): torch._foreach_mul_(ema_params, args.ema_beta) torch._foreach_add_(ema_params, named_params, alpha=1.0 - args.ema_beta) @contextlib.contextmanager def ema_weights(): backup = [p.detach().clone() for p in named_params] with torch.no_grad(): for p, e in zip(named_params, ema_params): p.copy_(e) try: yield finally: with torch.no_grad(): for p, b in zip(named_params, backup): p.copy_(b) def get_lr(it): assert it <= args.num_iterations if it < args.warmup_iters: return args.learning_rate * (it + 1) / args.warmup_iters elif args.warmdown_iters == 0 or it < args.num_iterations - args.warmdown_iters: return args.learning_rate # (guard: warmdown 0 would divide by zero) else: decay_ratio = (args.num_iterations - it) / args.warmdown_iters return args.learning_rate * decay_ratio run_id = args.run_name or str(uuid.uuid4()) run_dir = os.path.join(args.output_dir or "runs", run_id) logfile = None csvfile = None qkv_diagnostic_file = None resuming = bool(getattr(args, "resume", "")) and resume_step > 0 if master_process: os.makedirs(run_dir, exist_ok=True) cfg_dump = dict(vars(args)) cfg_dump["torch_version"] = str(torch.version.__version__) cfg_dump["resume_step"] = resume_step cfg_dump["trainer_sha256"] = TRAINER_SHA256 cfg_dump["muon_sha256"] = _muon_sha256() cfg_dump["modal_task_id"] = os.environ.get("MODAL_TASK_ID", "") cfg_dump["gpu_name"] = os.environ.get("NOCAP_GPU_ATTESTATION", "") config_name = ( "config.resume_step%06d.json" % resume_step if resuming else "config.json" ) with open(os.path.join(run_dir, config_name), "w") as f: json.dump(cfg_dump, f, indent=2) logfile = os.path.join(run_dir, "main.log") csvfile = os.path.join(run_dir, "log.csv") if muon_qkv_diagnostic_steps: qkv_diagnostic_file = os.path.join( run_dir, "muon_qkv_diagnostics.jsonl" ) if resuming: with open(logfile, "a") as f: f.write("\n--- RESUME step=%d train_time_ms=%.1f ---\n" % ( resume_step, resume_train_time_ms)) if not os.path.exists(csvfile): with open(csvfile, "w", newline="") as f: csv.writer(f).writerow( ["step", "tokens", "train_time_ms", "step_avg_ms", "lr", "train_loss", "val_disjoint", "val_sliding", "ema_val_disjoint", "ema_val_sliding"]) else: open(logfile, "w").close() if qkv_diagnostic_file: open(qkv_diagnostic_file, "w").close() with open(csvfile, "w", newline="") as f: csv.writer(f).writerow( ["step", "tokens", "train_time_ms", "step_avg_ms", "lr", "train_loss", "val_disjoint", "val_sliding", "ema_val_disjoint", "ema_val_sliding"]) def csv_row(step, tt_ms, lr, trainloss, vd="", vs="", evd="", evs="", updates_completed=None): if master_process and csvfile: completed = step if updates_completed is None else updates_completed with open(csvfile, "a", newline="") as f: csv.writer(f).writerow( [step, completed * tokens_per_iter, "%.1f" % tt_ms, "%.2f" % (tt_ms / max(completed, 1)), "%.6g" % lr, ("%.6f" % trainloss) if trainloss is not None else "", vd, vs, evd, evs]) def sync(): if use_cuda: torch.cuda.synchronize() training_time_ms = resume_train_time_ms sync() t0 = time.perf_counter() # resume_step = next step index to execute (checkpoint saved after completing that step) for step in range(resume_step, args.num_iterations + 1): last_step = step == args.num_iterations if args.val_loss_every > 0 and (step % args.val_loss_every == 0 or last_step): if val_tokens_buf is None: raise AssertionError( "validation requested in no-validation-data mode" ) sync() training_time_ms += 1000 * (time.perf_counter() - t0) # official (baseline-protocol) number first, always vd = evaluate(raw_model, val_tokens_buf, T, args.val_batch_size, device, "disjoint") vs = evd = evs = "" if args.val_stride > 0: vs = "%.6f" % evaluate(raw_model, val_tokens_buf, T, args.val_batch_size, device, "sliding", stride=args.val_stride) if ema_active and ema_params is not None: with ema_weights(): evd = "%.6f" % evaluate(raw_model, val_tokens_buf, T, args.val_batch_size, device, "disjoint") if args.val_stride > 0: evs = "%.6f" % evaluate(raw_model, val_tokens_buf, T, args.val_batch_size, device, "sliding", stride=args.val_stride) print0("step:%d/%d | val loss %.6f%s%s" % ( step, args.num_iterations, vd, (" | sliding %s" % vs) if vs else "", (" | ema %s" % evd) if evd else "")) if master_process: if logfile: with open(logfile, "a") as f: f.write("s:%d val:%f\n" % (step, vd)) csv_row(step, training_time_ms, get_lr(min(step, args.num_iterations)), None, "%.6f" % vd, vs, evd, evs) if args.log_wandb: import wandb wandb.log({"val_loss": vd, "time": training_time_ms}, step=step * tokens_per_iter) sync() t0 = time.perf_counter() if last_step: break if args.muon_to_adam_step >= 0 and step == args.muon_to_adam_step: from muon_opt_split_qkv_cheap_v1 import handoff_muon_groups_to_adam handoff_lr = ( args.learning_rate if args.muon_to_adam_lr is None else args.muon_to_adam_lr ) switched = handoff_muon_groups_to_adam(optimizer, handoff_lr) if switched <= 0: raise RuntimeError( "Muon-to-Adam handoff found no active Muon matrix groups" ) print0( "muon_to_adam handoff at step %d: %d matrices, base_lr=%.6g" % (step, switched, handoff_lr) ) model.train() # per-step control values for wave-2 features (buffers -> no recompiles) if args.mtp_lambda > 0: decay_end = max(int(args.mtp_decay_frac * args.num_iterations), 1) lam = args.mtp_lambda * max(0.0, 1.0 - step / decay_end) raw_model.mtp_lambda_t.fill_(lam) mtp_off = mtp_offsets[step % len(mtp_offsets)] else: mtp_off = -1 use_shape = args.loss_shape > 0 if use_shape: raw_model.shape_gate.fill_( 1.0 if step >= int(args.ls_start_frac * args.num_iterations) else 0.0) checkpoint_due = ( master_process and args.save_every > 0 and (step + 1) % args.save_every == 0 ) loss_log_due = (step % args.log_every == 0) or (step % 50 == 0) # Avoid one scalar GPU add per microbatch plus a device->host sync on # steps whose diagnostic loss will not be emitted. train_loss = torch.zeros(1, device=device) if loss_log_due else None for micro_step in range(args.grad_accumulation_steps): if not single_process_fastpath: model.require_backward_grad_sync = ( micro_step == args.grad_accumulation_steps - 1 ) with ctx: _, loss = model(x, y, return_logits=False, zloss=args.zloss, use_shape=use_shape, mtp_offset=mtp_off) loss = loss / args.grad_accumulation_steps if train_loss is not None: train_loss += loss.detach() x, y = train_loader.next_batch() scaler.scale(loss).backward() lr = get_lr(step) # scale each group relative to its base (Muon spectral vs AdamW) base_main = args.learning_rate if args.learning_rate > 0 else 1.0 scale = lr / base_main for param_group in optimizer.param_groups: base = param_group.get("base_lr", args.learning_rate) param_group["lr"] = base * scale if step in muon_qkv_diagnostic_steps: # This diagnostic is symmetric across the matched arms and is # explicitly excluded from the training clock, just like the # validation-only evidence path. It never changes optimizer state. sync() training_time_ms += 1000 * (time.perf_counter() - t0) from muon_opt_split_qkv_cheap_v1 import ( muon_update, muon_update_split_qkv, ) parameter = qkv_diagnostic_parameter group = qkv_diagnostic_group if parameter.grad is None: raise RuntimeError("Muon QKV diagnostic found no accumulated gradient") momentum = optimizer.state[parameter].get("momentum_buffer") if momentum is None: momentum = torch.zeros_like(parameter) with torch.no_grad(): joint_update = muon_update( parameter.grad.clone(), momentum.clone(), beta=group["momentum"], ns_steps=group.get("ns_steps", 5), ) split_update = muon_update_split_qkv( parameter.grad.clone(), momentum.clone(), beta=group["momentum"], ns_steps=group.get("ns_steps", 5), ) uncalibrated_split_update = muon_update_split_qkv( parameter.grad.clone(), momentum.clone(), beta=group["momentum"], ns_steps=group.get("ns_steps", 5), match_joint_frobenius=False, ) joint_norm = float(joint_update.float().norm().item()) split_norm = float(split_update.float().norm().item()) uncalibrated_split_norm = float( uncalibrated_split_update.float().norm().item() ) norm_ratio = split_norm / joint_norm uncalibrated_norm_ratio = uncalibrated_split_norm / joint_norm cosine = float( torch.nn.functional.cosine_similarity( joint_update.float().reshape(1, -1), split_update.float().reshape(1, -1), ).item() ) if not all(math.isfinite(value) for value in ( joint_norm, split_norm, uncalibrated_split_norm, norm_ratio, uncalibrated_norm_ratio, cosine, )): raise RuntimeError("non-finite Muon QKV diagnostic") diagnostic = { "step": step, "parameter": "transformer.h.0.attn.c_attn.weight", "joint_ns5_frobenius": joint_norm, "split_ns5_frobenius": split_norm, "split_over_joint_frobenius": norm_ratio, "uncalibrated_split_ns5_frobenius": uncalibrated_split_norm, "uncalibrated_split_over_joint_frobenius": uncalibrated_norm_ratio, "frobenius_match_scale": split_norm / uncalibrated_split_norm, "joint_split_cosine": cosine, "ideal_polar_frobenius": math.sqrt(3 * parameter.size(1)), "excluded_from_training_time": True, } print0("Muon QKV diagnostic: %s" % json.dumps(diagnostic, sort_keys=True)) if master_process and qkv_diagnostic_file: with open(qkv_diagnostic_file, "a") as handle: handle.write(json.dumps(diagnostic, sort_keys=True) + "\n") del joint_update, split_update, uncalibrated_split_update sync() t0 = time.perf_counter() # Muon path does not use GradScaler unscale semantics; bf16 runs use # enabled=False scaler which is a no-op step. For fp16+muon, force # unscaled step if needed. if getattr(args, "optimizer", "adamw") == "muon" and args.precision != "fp16": optimizer.step() else: scaler.step(optimizer) scaler.update() optimizer.zero_grad(set_to_none=True) if ema_active: if step == ema_start_step and ema_params is None: ema_params = [p.detach().clone().float() for p in named_params] elif ema_params is not None: with torch.no_grad(): ema_update() approx_ms = None if loss_log_due or checkpoint_due: sync() approx_ms = training_time_ms + 1000 * (time.perf_counter() - t0) if loss_log_due: if not single_process_fastpath: # AVG is unsupported on gloo; world-size 1 keeps historical behavior. dist.all_reduce(train_loss, op=dist.ReduceOp.SUM) lossf = train_loss.item() / ddp_world_size print0("step:%d/%d | loss %.6f | train_time:%.2fs | step_avg:%.2fms" % ( step, args.num_iterations, lossf, approx_ms / 1000, approx_ms / (step + 1))) if master_process and logfile: with open(logfile, "a") as f: f.write("s:%d trn:%f\n" % (step, lossf)) if step % 50 == 0: csv_row( step, approx_ms, lr, lossf, updates_completed=step + 1, ) # mid-run insurance checkpoint (e.g. before final warmdown on full runs): # enables a cheap re-anneal retry if the target is narrowly missed if checkpoint_due: # Full resume payload (preemption-safe). Recipe/hyperparams unchanged. ck = dict( model=raw_model.state_dict(), optimizer=optimizer.state_dict(), ema={k: v for k, v in zip( [n for n, _ in raw_model.named_parameters()], ema_params)} if (ema_active and ema_params is not None) else None, step=step + 1, training_time_ms=approx_ms, train_loader=train_loader.state_dict(), # x/y were prefetched at the end of the final micro-step. The # cursor already points after them, so preserve this small # pending batch to avoid a skip or replay on resume. pending_x=x.detach().cpu(), pending_y=y.detach().cpu(), args=args.__dict__, trainer_sha256=TRAINER_SHA256, muon_sha256=_muon_sha256(), checkpoint_ancestor_sha256=list(checkpoint_ancestor_sha256), ) ck_path = os.path.join(run_dir, "ckpt_step%06d.pt" % (step + 1)) torch.save(ck, ck_path) checkpoint_sha256 = _sha256_file(ck_path) if checkpoint_sha256 in checkpoint_ancestor_sha256: raise RuntimeError("new checkpoint repeats an ancestor SHA-256") checkpoint_ancestor_sha256.append(checkpoint_sha256) # pointer for Modal wrapper auto-resume with open(os.path.join(run_dir, "LATEST_CKPT"), "w") as f: f.write(ck_path + "\n") print0("saved insurance checkpoint at step %d (train_time_ms=%.1f)" % ( step + 1, approx_ms)) if args.trainable_logit_prior: p = raw_model.logit_prior.detach().float() print0( "trainable_logit_prior final min=%.6f max=%.6f mean=%.6f std=%.6f" % (p.min().item(), p.max().item(), p.mean().item(), p.std().item()) ) if use_cuda: print0("peak memory consumption: %d MiB" % (torch.cuda.max_memory_allocated() // 1024 // 1024)) if master_process and args.save_final: log = dict( model=raw_model.state_dict(), ema={k: v for k, v in zip( [n for n, _ in raw_model.named_parameters()], ema_params)} if (ema_active and ema_params is not None) else None, code=code, args=args.__dict__, trainer_sha256=TRAINER_SHA256, muon_sha256=_muon_sha256(), checkpoint_ancestor_sha256=list(checkpoint_ancestor_sha256), ) torch.save(log, os.path.join(run_dir, "final.pt")) print0("saved %s" % os.path.join(run_dir, "final.pt")) destroy_process_group()