# The training run in one script: frozen architecture, map-style packed data, fp16 on T4, trapezoid LR, # and a checkpoint cycle that satisfies §3.13 by construction. # Deliberately thin glue over `transformers.Trainer` (docs/01-plan.md §5): data loading, DDP, gradient # accumulation, scaler, RNG capture and resumption are the framework's job. What is ours is the shard # reader, the LR shape, and the upload/verify/prune discipline -- the three things a generic stack cannot # know about this project. # Two resume facts that are easy to get wrong and are asserted here rather than assumed: # 1. `ignore_data_skip=True`. Trainer's own resume would replay discarded batches; we advance the data # position by slicing the permutation in shard_dataset instead. Leaving the flag at False would # double-skip and silently LOSE training data -- the exact §3.1 violation. # 2. The cursor is checked against the checkpoint, not trusted. A resumed run whose cursor, seed, # sequence length or corpus fingerprint disagrees with what is on disk stops before step 1. import argparse import json import math import os import time import torch from torch.optim.lr_scheduler import LambdaLR from transformers import (LlamaConfig, LlamaForCausalLM, Trainer, TrainerCallback, TrainingArguments, default_data_collator, set_seed) import hubckpt import shard_dataset as SD # ---- frozen at Gate 1 (docs/01-plan.md §2.3, §4). Only --layers/--hidden change, for test T1. BASE = dict(hidden=576, layers=22, heads=9, kv_heads=3, ffn=1536, vocab=49152, rope_theta=10000.0, rms_eps=1e-5, tie=True) def _kv_dividing(heads, hidden): """GQA requires num_key_value_heads to divide num_attention_heads; SDPA raises otherwise ("Number of heads in key and value must divide the number of heads in Q"). Variant shapes are built from hidden/64, which need not divide 9 -- take the largest divisor of `heads` that fits.""" per_head = max(1, hidden // 64) return max(d for d in range(1, min(per_head, heads) + 1) if heads % d == 0) def build_config(seq_len, layers=None, hidden=None, attn="eager"): """The frozen shape, and the two variants test T1 needs (20L at the same width, and a wider/shallower alternative). SwiGLU width is int(2*4h/3) rounded to a 64-multiple -- 1536 at h=576, which is the value the frozen 106,194,240 count was derived from. Getting this factor wrong silently changes the parameter count and makes T1 compare two different models. `attn` is explicit because it is a measured choice, not a default: probe p0c found eager attention beating memory-efficient SDPA at seq 1024 on Turing (docs/00-platform-notes.md §9), and Gate 3 P2 then clocked 4,071 tok/s with the transformers default -- half the rate p0c saw in a raw loop. Leaving the attention kernel to whatever the library defaults to is how a 68-hour run becomes a 27-hour run. """ h = hidden or BASE["hidden"] L = layers or BASE["layers"] # "variant" means a different WIDTH, not merely a different layer count: the frozen GQA ratio (9 Q / # 3 KV) belongs to the frozen width, and passing --hidden 576 explicitly must not quietly turn the # model into MHA-9 and make test T1 measure something else. frozen_width = (h == BASE["hidden"]) inter = BASE["ffn"] if frozen_width else (int(round(8 * h / 3)) // 64) * 64 kv = BASE["kv_heads"] if frozen_width else _kv_dividing(BASE["heads"], h) cfg = LlamaConfig(vocab_size=BASE["vocab"], hidden_size=h, intermediate_size=inter, num_hidden_layers=L, num_attention_heads=BASE["heads"], num_key_value_heads=kv, max_position_embeddings=seq_len, rope_theta=BASE["rope_theta"], rms_norm_eps=BASE["rms_eps"], tie_word_embeddings=BASE["tie"], attention_dropout=0.0, hidden_act="silu", mlp_bias=False, attn_implementation=attn, pad_token_id=BASE["vocab"] - 1, eos_token_id=0) return cfg def count_params(model): """Total, embeddings included -- the number §2 budgets against. Reported both ways: numel sum and the unique-parameter count, which differ when embeddings are tied.""" total = sum(p.numel() for p in model.parameters()) uniq = sum(p.numel() for p in set(model.parameters())) return {"sum_numel": total, "unique": uniq, "trainable": sum(p.numel() for p in model.parameters() if p.requires_grad)} # The cursor arithmetic in one place, because three separate bugs live here: # - a step consumes micro_batch * world * accum sequences, not micro_batch (DDP shards the batch); # - Trainer restores `global_step` from the checkpoint, so on resume it is ALREADY absolute. Adding the # resume point a second time inflates the cursor, and an inflated cursor makes the next run skip data # it never read -- silent data loss with no error anywhere, which is §3.1's worst case. Writing this # as a pure function of the absolute step removes the whole class of bug. # - max_steps is the TOTAL horizon (it is what the LR schedule decays against), so resuming mid-run # must keep the same steps_planned or the trapezoid changes shape. def samples_at(absolute_step, seqs_per_step): return int(absolute_step) * seqs_per_step class TrapezoidTrainer(Trainer): """Trainer with (a) a hand-rolled trapezoid schedule and (b) a data position we own.""" def __init__(self, *a, lr_shape=None, **kw): self.lr_shape = lr_shape super().__init__(*a, **kw) def create_scheduler(self, num_training_steps, optimizer=None): warm = self.lr_shape["warmup_steps"] decay_at = self.lr_shape["decay_start"] total = self.lr_shape["total_steps"] def fn(cur): if cur < warm: return (cur + 1) / max(1, warm) if cur < decay_at: return 1.0 # linear decay to exactly zero at the end: the trapezoid in docs/01-plan.md §4.3 return max(0.0, 1.0 - (cur - decay_at) / max(1, total - decay_at)) if self.lr_scheduler is None: self.lr_scheduler = LambdaLR(self.optimizer if optimizer is None else optimizer, fn) return self.lr_scheduler class HubPush(TrainerCallback): """§3.13 after every checkpoint: write the cursor in, upload, verify from the Hub, prune locally. Only rank 0 acts -- the ranks share the checkpoint directory and would race. A failed verification raises inside the callback, so the run stops AT that checkpoint rather than continuing with its only copy on a disk that can be taken away at any moment. """ def __init__(self, repo, api, token, seq_len, seed, fingerprint, seqs_per_step, prune=True, min_free_gb=6.0, repo_type="dataset", stop_after_steps=0, order_sha=None): self.repo, self.api, self.token = repo, api, token self.seq_len, self.seed, self.fingerprint = seq_len, seed, fingerprint self.seqs_per_step, self.repo_type = seqs_per_step, repo_type self.order_sha = order_sha self.prune, self.min_free_gb = prune, min_free_gb self.stop_after = int(stop_after_steps or 0) self.stopped_at = None self.n = 0 self.steps = [] self.results = [] def on_step_end(self, args, state, control, **kw): # A session that ends ON a verified checkpoint loses nothing when the platform takes the disk. # That is how the 10 % cadence is made compatible with an unknown session cap: instead of pushing # extra intermediate copies (which multiplies Hub LFS storage by the number of sessions), the run # is cut into segments whose length is a whole number of checkpoint intervals. if self.stop_after and state.global_step >= self.stop_after: self.stopped_at = state.global_step control.should_training_stop = True return control def on_save(self, args, state, control, **kw): if not state.is_world_process_zero: return control step = state.global_step d = os.path.join(args.output_dir, f"checkpoint-{step}") if not os.path.isdir(d): raise RuntimeError(f"Trainer reported a save but {d} does not exist") samples = samples_at(step, self.seqs_per_step) SD.save_cursor(d, SD.cursor_dict(self.seq_len, samples, self.seed, self.fingerprint, step=step, order_sha=self.order_sha)) t0 = time.time() # Upload and verify *without* pruning, roll the pointer and read it back, and only then delete the # local copy. The old order (prune inside push_and_prune, pointer written afterwards) meant one # failed pointer commit left a Hub-verified checkpoint that nothing pointed at and no local copy to # recover it with -- and every guard a resume runs would still pass, because weights, cursor, # fingerprint and seed are all self-consistent. Only the pointer would be wrong, and the next # session would quietly re-read 100 M tokens. res = hubckpt.push_and_prune(self.repo, d, f"ckpt/checkpoint-{step}", self.api, repo_type=self.repo_type, token=self.token, prune=False) # E-042: `prune=False` above means `push_and_prune`'s own guard never runs, so this path had to # carry it. An unverified checkpoint must stop the session *before* the pointer moves: a pointer # rolled onto bytes nobody checked, with the local copy then deleted, is a run that cannot be # resumed and does not know it. ver = res["verify"] if not ver["ok"]: raise RuntimeError( f"checkpoint {step} did not verify on the Hub " f"(missing={len(ver['missing'])} mismatched={len(ver['mismatch'])} " f"extra={len(ver['extra'])} readback_ok={ver['readback_ok']}). Keeping the local copy in " f"{d} and refusing to roll latest.json -- the next session would otherwise resume from " "bytes that were never confirmed to exist.") ptr_body = {"path_in_repo": f"ckpt/checkpoint-{step}", "step": step, "samples_consumed": samples, "tokens_consumed": samples * self.seq_len, "fingerprint": self.fingerprint, "seed": self.seed} hubckpt.write_latest_pointer(self.repo, ptr_body, self.api, repo_type=self.repo_type, token=self.token) if self.prune: hubckpt.prune_verified(d, res, self.min_free_gb) loss = next((h["loss"] for h in reversed(state.log_history) if "loss" in h), None) print(f"CKPT {step} samples={samples:,} push+verify={round(time.time()-t0,1)}s " f"ok={res['verify']['ok']} pruned={res.get('pruned')} free={res['free_after_gb']}GB " f"loss={loss}", flush=True) self.n += 1 self.steps.append(step) self.results.append(res) return control def run_benchmark(args, model, dataset, steps, world, seqs_per_step): """The same data path, optimizer and fp16 recipe as the real run, with no Trainer in the loop. Why this exists: P2 measured 5,105 tok/s for 22L/seq-1024 where probe p0c measured 11,062 tok/s for a near-identical shape with a hand-rolled loop. The gap is either (a) the attention kernel, (b) per-step framework overhead, or (c) a measurement artefact, and which one it is decides whether the main run needs 27 h or 68 h. Guessing wrong in either direction is expensive, so measure both paths. """ import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.data import DataLoader local_rank = int(os.environ.get("LOCAL_RANK", "0")) device = torch.device("cuda", local_rank) torch.cuda.set_device(device) if world > 1 and not dist.is_initialized(): dist.init_process_group("nccl") model = model.to(device) if args.grad_ckpt: model.gradient_checkpointing_enable() model.config.use_cache = False net = DDP(model, device_ids=[local_rank]) if world > 1 else model opt = torch.optim.AdamW(net.parameters(), lr=args.lr, betas=(0.9, 0.95), weight_decay=0.1, eps=1e-8) try: scaler = torch.amp.GradScaler("cuda", enabled=True) # torch >= 2.4 except TypeError: scaler = torch.cuda.amp.GradScaler(enabled=True) # older signature # Same reader, same shuffle, no accumulation: the comparison against the Trainer cells is then # per-step cost with the framework removed, and nothing else differs. dl = DataLoader(dataset, batch_size=args.micro_batch, shuffle=False, num_workers=2, pin_memory=False, drop_last=True) t0 = time.time() toks = 0 for i, batch in enumerate(dl): if i >= steps: break batch = {k: v.to(device, non_blocking=True) for k, v in batch.items()} opt.zero_grad(set_to_none=True) with torch.autocast("cuda", dtype=torch.float16): out = net(**batch) scaler.scale(out.loss).backward() scaler.unscale_(opt) torch.nn.utils.clip_grad_norm_(net.parameters(), 1.0) scaler.step(opt) scaler.update() toks += batch["input_ids"].numel() if i == 0: torch.cuda.synchronize() t0 = time.time() # exclude CUDA context + first-trace cost from the rate toks = 0 sec = time.time() - t0 res = {"benchmark_steps": i + 1 if steps else 0, "tok_per_s": round(toks / max(1e-6, sec)), "seconds": round(sec, 1), "tokens": toks, "peak_gpu_gb": round(torch.cuda.max_memory_allocated() / 1e9, 2), "world": world, "seq_len": args.seq_len, "micro_batch": args.micro_batch, "attn": args.attn, "grad_ckpt": bool(args.grad_ckpt), "params": sum(p.numel() for p in model.parameters())} print("BENCH_JSON " + json.dumps(res, sort_keys=True), flush=True) if world > 1: dist.barrier() dist.destroy_process_group() return res def main(): ap = argparse.ArgumentParser() ap.add_argument("--root", default="/kaggle/working/mixroot") ap.add_argument("--out", default="/kaggle/working/run") ap.add_argument("--hub-repo", default="", help="dataset repo for checkpoints; empty = no upload") ap.add_argument("--hub-repo-type", default="dataset", choices=["dataset", "model"]) # D-011: these ARE the frozen main-run values (measured 7,732-7,828 tok/s at 1024/eager on 2xT4), so # that launching the run does not depend on remembering a flag string. ap.add_argument("--seq-len", type=int, default=1024) ap.add_argument("--attn", default="eager", choices=["eager", "sdpa", "flash_attention_2"], help="p0c measured eager beating SDPA on Turing; Gate 3 P2/A-B measured +52 %% for " "eager at seq 1024 and an OOM at 2048. flash_attention_2 needs sm80+") ap.add_argument("--grad-ckpt", default=True, action=argparse.BooleanOptionalAction, help="on by default: it is what keeps peak memory at ~7 GB instead of 12 GB, and the " "measured rate cost at 1024 was nil (7,828 with GC at m8 vs 7,732 without at m4)") ap.add_argument("--layers", type=int, default=0) ap.add_argument("--hidden", type=int, default=0) ap.add_argument("--micro-batch", type=int, default=4, help="4 x 1024 x 32 accum x 2 cards = 262,144 tokens/step, the frozen batch size") ap.add_argument("--accum", type=int, default=32) ap.add_argument("--lr", type=float, default=6e-4) ap.add_argument("--optim", default="adamw_torch", help="only mathematically-equivalent implementations are fair game at freeze time " "(adamw_torch / adamw_torch_fused / adamw_apex_fused); 8-bit variants change the " "update, so they are excluded by rule, not by preference") ap.add_argument("--gc-nonreentrant", default=False, action=argparse.BooleanOptionalAction, help="pass use_reentrant=False to gradient checkpointing: same recomputation, " "different autograd path, and reports on whether it is faster under DDP conflict") ap.add_argument("--torch-compile", default=False, action=argparse.BooleanOptionalAction, help="Trainer's torch_compile hook; fused kernels, no change to the update rule") ap.add_argument("--tokens", type=int, default=1_000_000_000) ap.add_argument("--max-steps", type=int, default=0, help="0 = derive from --tokens; >0 overrides") ap.add_argument("--warmup-frac", type=float, default=0.02) ap.add_argument("--decay-frac", type=float, default=0.80) ap.add_argument("--seed", type=int, default=20260919) ap.add_argument("--data-seed", type=int, default=20260919) ap.add_argument("--resume", default="", help="checkpoint dir; empty = auto from Hub latest") ap.add_argument("--val-tokens", type=int, default=2_000_000) ap.add_argument("--log-every", type=int, default=10) ap.add_argument("--push-every-steps", type=int, default=0, help="0 = only at 10%% boundaries") ap.add_argument("--stop-after-steps", type=int, default=0, help="absolute global step at which THIS session stops training, so it ends on a " "verified checkpoint instead of being killed between two of them. Must be a " "whole multiple of the checkpoint interval; 0 = run to the horizon.") ap.add_argument("--prune", action="store_true", help="delete local checkpoint after hub verify") ap.add_argument("--dry-run", action="store_true", help="build everything, train 0 steps") ap.add_argument("--benchmark", type=int, default=0, help="N steps with no Trainer in the loop (isolates framework per-step cost)") args = ap.parse_args() rank = int(os.environ.get("RANK", "0")) world = int(os.environ.get("WORLD_SIZE", "1")) say = (lambda *m: print(*m, flush=True)) if rank == 0 else (lambda *m: None) man_path = os.path.join(args.root, "manifest.json") if not os.path.exists(man_path): raise SystemExit(f"no manifest at {man_path}: build the mix first (Phase 2 is not optional)") man = json.load(open(man_path)) if man["total_tokens"] < args.tokens: say(f"WARNING mix has {man['total_tokens']:,} tokens < target {args.tokens:,}") store = SD.PackedTokenStore(args.root, man) fingerprint = SD.mix_fingerprint(args.root, man) # One definition of the step, used by the cursor, the schedule and the token count: a step consumes # micro_batch * world sequences (DDP shards the batch), so tokens/step follows from it. seqs_per_step = args.micro_batch * world * args.accum tokens_per_step = seqs_per_step * args.seq_len steps_planned = args.max_steps or int(args.tokens / tokens_per_step) say(f"mix: {store.total_tokens:,} tokens in {len(store.shards)} shards fp={fingerprint} " f"| {tokens_per_step:,} tok/step x {steps_planned} steps = " f"{tokens_per_step * steps_planned:,} tokens") # ---- resume point: the cursor is the run's identity as much as the weights are. # `--resume auto` is what a fresh instance with an empty disk passes: it asks the Hub where the run # stopped and pulls that checkpoint down. §3.1 forbids recovering from anything left on an instance. if args.resume == "auto": if not args.hub_repo: raise SystemExit("--resume auto needs --hub-repo") import ounce100m_credentials ounce100m_credentials.install() from huggingface_hub import HfApi tk = os.environ.get("HF_TOKEN") ptr = hubckpt.latest_pointer(args.hub_repo, repo_type=args.hub_repo_type, token=tk) if ptr.get("not_found"): # A 404 is normally a fact -- "this run has never checkpointed" -- and step 0 is then right. # But a missing pointer beside existing checkpoints means the pointer upload failed after a # good checkpoint, and starting from zero would silently retrain the corpus from the top, # which §3.1 forbids and E-018 is about. So list the repo before believing the 404. try: listed = sorted(k for k in hubckpt.hub_listing(args.hub_repo, args.hub_repo_type, token=tk) if k.startswith("ckpt/checkpoint-")) except hubckpt.RepoMissing: # The checkpoint repo has never been created: nothing has ever been pushed, so step 0 is # the only correct answer. The first session of the run always takes this branch, and # treating a 404 as "listing failed" would refuse it (a generic `except Exception` below # makes the list non-empty and the refusal fires). say(f"{args.hub_repo} does not exist yet -- nothing has ever been pushed") listed = [] except Exception as e: listed = [f""] if listed: raise SystemExit(f"--resume auto found no latest.json on {args.hub_repo}, but the repo " f"holds {len(listed)} checkpoint path(s) ({listed[:3]}). The pointer is " "missing, not the run -- refusing to start from step 0.") say(f"no latest.json and no checkpoints on {args.hub_repo} -> starting from step 0") args.resume = "" elif "error" in ptr: # Anything else (401, timeout, DNS) means the lookup failed, not that there is nothing there. # Falling through to step 0 here would silently retrain the first stretch of the run, which is # exactly the §3.1 violation this whole branch exists to prevent (E-018's shape). raise SystemExit(f"--resume auto could not read the hub pointer: {ptr['error']}") else: # A pointer is a resume point only if nothing newer was ever pushed. This scan catches the one # case write_latest_pointer's own read-back cannot see: a pointer left behind by a rollback, a # deletion, or a session that died between its checkpoint push and its pointer roll. Resuming # from the stale pointer would re-train an interval with weights and a cursor that agree with # each other and disagree with the Hub. import re as _re steps_on_hub = set() for attempt in (1, 2): try: steps_on_hub = {int(m.group(1)) for k in hubckpt.hub_listing(args.hub_repo, args.hub_repo_type, token=tk) if (m := _re.match(r"ckpt/checkpoint-(\d+)/", k))} break except Exception as e: say(f" checkpoint scan attempt {attempt}: {type(e).__name__}: {str(e)[:140]}") time.sleep(15) else: raise SystemExit("--resume auto could not enumerate the checkpoint repo, so it cannot " f"prove latest.json (step {ptr.get('step')}) is the newest state") if steps_on_hub and max(steps_on_hub) > int(ptr.get("step") or 0): raise SystemExit(f"latest.json points at step {ptr.get('step')} but {args.hub_repo} also " f"holds checkpoint-{max(steps_on_hub)}: the pointer is behind the run. " "Resuming here would re-train that interval, so refusing.") # Both ranks need the checkpoint on disk, and this block runs before TrainingArguments has # created a process group -- so there is no barrier to synchronise on. Downloading into a # per-rank directory instead removes the race entirely (two writers, one path), at the cost of # a second 1.7 GB pull that happens in parallel. Each rank then reads its own verified copy and # derives the same start_sample/start_step from the same cursor.json bytes. dest = os.path.join(args.out, f"resume_r{rank}", os.path.basename(ptr["path_in_repo"])) say(f"auto-resume: hub says step {ptr.get('step')} at {ptr['path_in_repo']}") got = hubckpt.download_checkpoint(args.hub_repo, ptr["path_in_repo"], dest, HfApi(), repo_type=args.hub_repo_type, token=tk) say(f"downloaded {got['files']} files, {got['bytes']/1e6:.0f} MB in {got['seconds']}s") if got["bytes"] < 1_000_000: raise SystemExit("auto-resume downloaded a checkpoint that is not real weights") args.resume = got["dir"] say(f"resuming from {args.resume}") cur = SD.load_cursor(args.resume) if args.resume else None if args.resume and cur is None: # Weights at step N with the data stream back at sample 0 would retrain the opening of the corpus # and report a continuous loss curve. Refuse instead. raise SystemExit(f"resume refused: {args.resume} has no cursor.json -- a checkpoint without a " "recorded data position cannot be resumed, only restarted (§3.1)") if args.resume: say(f"resume from {args.resume} cursor={cur}") start_sample = (cur or {}).get("samples_consumed", 0) start_step = (cur or {}).get("step") or 0 # Identity of the visit order, derived from the corpus size and the seed rather than trusted to the # seed: numpy's permutation is only stable within a version, and Kaggle updates its image underneath # long runs. ~50 ms once per session, and it is what makes "resumed onto a different order" loud. order_sha = SD.perm_sha(store.total_tokens // args.seq_len, args.data_seed) if cur: for k, v in (("seq_len", args.seq_len), ("shuffle_seed", args.data_seed), ("shuffle_perm_sha", order_sha), ("dataset_files_sha", fingerprint)): if cur.get(k) != v: raise SystemExit(f"resume refused: cursor {k}={cur.get(k)} != current {v}. " "A resumed run must read the same data in the same order (§3.1).") if start_sample % seqs_per_step: raise SystemExit(f"resume refused: cursor samples_consumed={start_sample} is not a whole " f"number of steps (seqs/step={seqs_per_step}) -- batch geometry changed") if cur.get("step") is not None and cur["step"] * seqs_per_step != start_sample: raise SystemExit(f"resume refused: cursor step={cur['step']} and samples_consumed=" f"{start_sample} disagree at seqs/step={seqs_per_step}") if start_step * seqs_per_step != start_sample: start_step = start_sample // seqs_per_step train_ds = SD.make_dataset(store, args.seq_len, args.data_seed, start_sample=start_sample) assert getattr(train_ds, "order_sha", None) == order_sha, "two derivations of the order disagree" # Validation comes from the mix's held-out val/ shards, never from the training stream: scoring on # tokens the model already consumed would be a number nobody could reproduce (and §3.12). vs = [s["file"] for s in man.get("val_shards", [])] if not vs: raise SystemExit("manifest lists no val_shards -- Gate 2 requires a held-out split") vstore = SD.PackedTokenStore(args.root, man, files=vs) n_val = min(max(1, args.val_tokens // args.seq_len), vstore.total_tokens // args.seq_len) val_ds = SD.make_dataset(vstore, args.seq_len, args.data_seed, shuffle=False, count=n_val) need = (steps_planned - start_step) * seqs_per_step if len(train_ds) < need: # Reading past the end of the stream makes the Trainer begin epoch 2 at the top of the same # permutation: the run would then report the full token target while quietly re-training on the # first windows. The mix is built with ~12 % headroom over the target for exactly this reason, so # arriving here means that assumption stopped being true (review finding), and the honest response # is to stop before step 1 rather than to discover it in the final report. raise SystemExit( f"dataset too small for the plan: {len(train_ds):,} windows available from sample " f"{start_sample:,}, {need:,} needed for {steps_planned} steps x {seqs_per_step} sequences " f"({need * args.seq_len:,} tokens). Lower --tokens to the 0.9 B fallback (D-012) or rebuild " "the mix larger -- never let the epoch wrap.") say(f"data: {len(train_ds):,} training windows available, {need:,} needed " f"({100.0 * len(train_ds) / need - 100:.1f} % margin), {n_val:,} held-out val windows") cfg = build_config(args.seq_len, args.layers or None, args.hidden or None, attn=args.attn) set_seed(args.seed) # before init: both ranks must start from identical weights, not two draws model = LlamaForCausalLM(cfg).to(dtype=torch.float32) pc = count_params(model) say("params:", json.dumps(pc)) if not (90_000_000 <= pc["sum_numel"] <= 110_000_000) and not args.layers and not args.hidden: raise SystemExit(f"parameter count {pc['sum_numel']:,} is outside §2's 90-110M budget") if args.benchmark: # Both ranks iterate the same samples here (no DistributedSampler): this path measures step cost, # not a training run, and it must never be used to produce a checkpoint. say(f"benchmark mode: {args.benchmark} steps, no Trainer, micro={args.micro_batch} " f"seq={args.seq_len} attn={args.attn} world={world}") run_benchmark(args, model, train_ds, args.benchmark, world, seqs_per_step) return # 10 checkpoints per run is §2's cadence; --push-every-steps overrides for the preflight. ckpt_every = args.push_every_steps or max(1, steps_planned // 10) if args.stop_after_steps: if args.stop_after_steps % ckpt_every: raise SystemExit(f"--stop-after-steps {args.stop_after_steps} is not a whole multiple of the " f"checkpoint interval {ckpt_every}: the session would end between two " f"checkpoints, which is exactly what this flag exists to prevent") if args.stop_after_steps > steps_planned: raise SystemExit(f"--stop-after-steps {args.stop_after_steps} exceeds the horizon " f"{steps_planned}; use 0 to run to the end") api = None if args.hub_repo: import ounce100m_credentials ounce100m_credentials.install() from huggingface_hub import HfApi api = HfApi() # rank 0 only: two ranks committing the same .gitattributes to a brand-new repo is a write race # on the one path that must not fail (only rank 0 ever uploads checkpoints, so rank 1 needs none) if rank == 0: say("ckpt repo:", json.dumps(hubckpt.ensure_repo(args.hub_repo, api, os.environ.get("HF_TOKEN"), args.hub_repo_type))) say(f"checkpoint hub target {args.hub_repo} (prune={args.prune}, every {ckpt_every} steps)") tacc = TrainingArguments( output_dir=args.out, per_device_train_batch_size=args.micro_batch, # Same geometry for the validation pass: with the library default (8) the end-of-run eval measures # a peak that has nothing to do with the training configuration, which is how an OOM in the no- # checkpointing cell could be mistaken for a passing one (review finding). per_device_eval_batch_size=args.micro_batch, gradient_accumulation_steps=args.accum, learning_rate=args.lr, weight_decay=0.1, adam_beta1=0.9, adam_beta2=0.95, adam_epsilon=1e-8, max_grad_norm=1.0, lr_scheduler_type="constant", # replaced by TrapezoidTrainer.create_scheduler warmup_ratio=0.0, # the trapezoid owns warmup num_train_epochs=1, max_steps=0 if args.dry_run else steps_planned, fp16=True, bf16=False, # measured decision, not a default: A-B found eager attention 52 % faster at seq 1024 and OOM at # 2048; checkpointing is what makes eager at 2048 possible at all, at the cost of a recompute pass gradient_checkpointing=args.grad_ckpt, # Same recomputation, different autograd path -- offered as a cells-vs-baseline question, not a # belief. Reports on use_reentrant=False under DDP contradict each other across versions. gradient_checkpointing_kwargs={"use_reentrant": not args.gc_nonreentrant} if args.grad_ckpt else None, optim=args.optim, torch_compile=args.torch_compile, ddp_find_unused_parameters=False, # Rank 0 spends minutes inside on_save uploading ~1.7 GB, listing the repo and downloading one file # back to hash it, while rank 1 waits at the next collective. The library's 1,800 s default would # abort the run on a slow-egress Kaggle box -- precisely during the checkpoint that the other rank # was protecting. Four hours makes the collective patient without hiding a real deadlock. ddp_timeout=14400, dataloader_num_workers=2, dataloader_pin_memory=False, remove_unused_columns=False, # our dict keys must reach the model untouched ignore_data_skip=True, # see header note 1 -- we advance position ourselves save_strategy="no" if args.dry_run else "steps", save_steps=ckpt_every, save_total_limit=1, # §3.13: the Hub is the store, the instance is scratch logging_steps=args.log_every, report_to=[], seed=args.seed, data_seed=args.seed, accelerator_config={"dispatch_batches": False}, average_tokens_across_devices=False, ) tr = TrapezoidTrainer( model=model, args=tacc, train_dataset=train_ds, eval_dataset=val_ds, data_collator=default_data_collator, lr_shape={ "warmup_steps": max(100, int(steps_planned * args.warmup_frac)), "decay_start": int(steps_planned * args.decay_frac), "total_steps": steps_planned}, ) hub_cb = None if api is not None: hub_cb = HubPush(args.hub_repo, api, os.environ.get("HF_TOKEN"), args.seq_len, args.data_seed, fingerprint, seqs_per_step, prune=args.prune, repo_type=args.hub_repo_type, stop_after_steps=args.stop_after_steps, order_sha=order_sha) tr.add_callback(hub_cb) # Precision must be ASSERTED, not configured: bf16 "works" silently on Turing and fp32 master # weights can be lost in a refactor (docs/00-platform-notes.md §8, memory/ERRORS.md E-006). The # scaler's home moved between transformers versions, so look in both places instead of trusting one # attribute path (E-010: prefer version-free surfaces). scaler = getattr(tr, "scaler", None) or getattr(getattr(tr, "accelerator", None), "scaler", None) eff_attn = (getattr(model.config, "_attn_implementation", None) or getattr(model.config, "attn_implementation", None)) say(f"precision: fp16={tacc.fp16} bf16={tacc.bf16} scaler={type(scaler).__name__} " f"master_dtype={next(model.parameters()).dtype} requested_attn={args.attn} " f"effective_attn={eff_attn}") assert tacc.fp16 and not tacc.bf16, "not the frozen T4 fp16 recipe" assert next(model.parameters()).dtype == torch.float32, "master weights are not fp32" prior = None if args.resume: tsp = os.path.join(args.resume, "trainer_state.json") if os.path.exists(tsp): h0 = [h for h in json.load(open(tsp)).get("log_history", []) if "loss" in h] prior = h0[-1]["loss"] if h0 else None say(f"prior loss at resume: {prior} -- loss continuity is judged against this number") if args.dry_run: say(json.dumps({"dry_run": True, "train_samples": len(train_ds), "val_samples": len(val_ds), "steps_planned": steps_planned, "ckpt_every": ckpt_every, "params": pc, "fingerprint": fingerprint, "scaler": type(scaler).__name__})) return t0 = time.time() tr.train(resume_from_checkpoint=args.resume or None) elapsed = time.time() - t0 step = tr.state.global_step # Did this segment really end on a pushed checkpoint? `stopped_at` is set when the step counter reaches # the target, but performing the save at that step is the library's decision. If it ever ordered the # stop ahead of the save, the session would exit 0 having pushed nothing, and every relaunch would # redo the same interval while `latest` stayed put -- which looks like a stuck run, not a broken one. if hub_cb is not None and args.stop_after_steps and rank == 0: # Rank 0 only, because every fact below is a rank-0 fact: `on_save` returns early on the other # ranks, so `hub_cb.steps` is empty on rank 1 by construction and this guard -- written to catch # a segment that ended without a pushed checkpoint -- killed *every* correctly checkpointed # segment on rank 1 with `pushed: []`. That SystemExit, not a divergent collective, is what # probe v1/v4/v5/v6/v7 saw as "rank 1 exitcode 1 with error_file ": rank 1 died on its own # assertion while rank 0 sat in the next all_reduce until torchrun SIGTERMed it (E-044). if args.stop_after_steps <= start_step: raise SystemExit(f"--stop-after-steps {args.stop_after_steps} is not beyond the resume point " f"step {start_step}: this session would train nothing") if step < args.stop_after_steps: raise SystemExit(f"asked to stop at step {args.stop_after_steps}, stopped at {step}") if args.stop_after_steps < steps_planned and args.stop_after_steps not in hub_cb.steps: raise SystemExit(f"segment ended at step {step} but no checkpoint was pushed and verified " f"for it (pushed: {sorted(hub_cb.steps)}). The next session would resume " "from the previous pointer and re-read this interval.") samples = samples_at(step, seqs_per_step) consumed = samples * args.seq_len consumed_this_run = max(0, step - start_step) * tokens_per_step hist = [h for h in tr.state.log_history if "loss" in h] final_loss = hist[-1]["loss"] if hist else None # Every rank, with `print` and not `say`: probe v6 died on the leg that stopped early and the only # evidence torchrun passed up was "rank 1 exitcode 1, error_file ", because the lines that would # have localised it are rank-0-only. These four marks say which statement in the post-training # sequence a rank never reached. print(f"[rank {rank}] post-train: step={step}/{steps_planned} log_history_loss_entries={len(hist)} " f"final_loss={final_loss} start_step={start_step}", flush=True) say(f"TRAIN DONE step={step}/{steps_planned} elapsed={elapsed/60:.1f} min " f"tokens_consumed={consumed:,} last_loss={final_loss}") # E-006, again: a run that produced NaN still reports a confident tokens/sec. Never let that pass. if final_loss is None or not math.isfinite(final_loss): raise SystemExit(f"final loss is not finite ({final_loss}) -- not exporting anything") # A segment that stopped short of the horizon does not run the validation pass, on every rank. # # Written for what probes v1/v4 saw ("rank 1 exitcode 1, error_file , right after a forced stop"), # which was attributed to `evaluate()` being entered from a diverged collective state. E-044 later found # the actual cause -- the rank-0-only push assertion firing on rank 1 -- so this skip is not what makes # the forced-stop leg work. It stays, on its own merits: the held-out PPL of a 762-step segment is not # worth a forward pass over 2 M tokens inside a billed session, and §5's per-wake monitoring is served by # the training loss. The report's PPL comes from the horizon leg, which is where it is measured. forced_stop = bool(hub_cb and hub_cb.stopped_at and step < steps_planned) print(f"[rank {rank}] past the stop guard: forced_stop={forced_stop} " f"stopped_at={getattr(hub_cb, 'stopped_at', None)}", flush=True) ppl = None val_err = None if forced_stop: say(f"validation skipped: segment ended at step {step} of {steps_planned} (evaluate() after a " "forced stop diverges between ranks -- probe 1 v1/v4)") else: try: ev = tr.evaluate() ppl = math.exp(min(20.0, ev["eval_loss"])) say(f"validation loss={ev['eval_loss']:.4f} ppl={ppl:.2f} over {n_val} held-out windows") except Exception as e: # `say()` is a no-op on rank 1, and that silence is what made a rank-1 evaluation crash unreadable # in probe 1 v1: the handler ran on both ranks, only rank 0 spoke, and the divergence looked like an # unexplained exit 1 with `error_file: `. Report on every rank, and carry the failure into the # record so a missing PPL can never be read as a passing validation pass. val_err = f"{type(e).__name__}: {str(e)[:300]}" print(f"[rank {rank}] validation eval failed: {val_err}", flush=True) # Sampled here, on every rank, before the rank-0-only export block: a collective reduction inside # `if rank == 0:` would hang, and reading rank 0's allocator alone is what made the probe's memory # verdict systematically optimistic -- it hides the other card's peak, the two CUDA contexts, and the # reserved-but-unallocated blocks that are the real ceiling on a 14.56 GiB T4 (review finding B4). print(f"[rank {rank}] peak_stats: entering (this is a collective)", flush=True) mem = peak_stats() print(f"[rank {rank}] peak_stats: done, max_reserved={mem['max_reserved_gb']} GiB", flush=True) say("memory:", json.dumps(mem, sort_keys=True)) # Everything below writes files, so it is rank 0's alone (E-029). Two ranks building the same export # directory is a write race: `shutil.copy(src, out)` with `out` not yet created writes a regular FILE # called `final` where rank 0 is making a directory of the same name, and the observed result was # `training_args.bin` vanishing mid-hash -- which killed a fully trained, fully checkpointed leg on # its last line. The variable is also no longer called `out`, which until now meant both the run # directory and the export inside it. exp = os.path.join(args.out, "final") if rank == 0: tr.save_model(exp) # The tokenizer travels with the run so `final/` is loadable by AutoTokenizer on its own (§8). If # it is not in the mix root, say so in the summary instead of skipping quietly and being surprised. shipped = [] import shutil as _sh os.makedirs(exp, exist_ok=True) for fn in ("tokenizer.json", "tokenizer_config.json", "special_tokens_map.json", "vocab.json", "merges.txt"): src = os.path.join(args.root, fn) if os.path.exists(src): _sh.copy(src, exp) shipped.append(fn) say(f"tokenizer files shipped with the export: {shipped or 'NONE -- Phase 5 must add them'}") SD.save_cursor(exp, SD.cursor_dict(args.seq_len, samples, args.data_seed, fingerprint, step=step, order_sha=order_sha)) json.dump({"step": step, "start_step": start_step, "tokens_consumed": consumed, "tokens_this_run": consumed_this_run, "session_seconds": round(elapsed, 1), "params": pc, "params_total": pc["sum_numel"], "fingerprint": fingerprint, "shuffle_perm_sha": order_sha, "final_loss": final_loss, "prior_loss": prior, "samples_consumed": samples, "steps_planned": steps_planned, # The recipe, written down. The Phase 5 card reads this file to state hyperparameters, # and it had none of them: every claim about sequence length, batch shape, throughput or # attention would have rendered as `None` because those lived only in the stdout RUN_JSON # line, a different schema (review finding A2). `seconds` was also renamed: it is THIS # session's elapsed, and a five-session run published as one number is wrong by ~5x (A3). "seq_len": args.seq_len, "micro_batch": args.micro_batch, "accum": args.accum, "world": world, "tokens_per_step": tokens_per_step, "tok_per_s": round(consumed_this_run / max(1, elapsed)), "attn": args.attn, "effective_attn": eff_attn, "grad_ckpt": bool(args.grad_ckpt), "optim": args.optim, "lr": args.lr, "warmup_frac": args.warmup_frac, "decay_frac": args.decay_frac, "seed": args.seed, "data_seed": args.data_seed, "val_ppl": ppl, "val_loss_error": val_err, "val_skipped": forced_stop, "val_windows": n_val, "peak_gpu_gb": mem["rank0_allocated_gb"], "peak_reserved_gb": mem["rank0_reserved_gb"], "peak_gpu_gb_max_rank": mem["max_allocated_gb"], "peak_reserved_gb_max_rank": mem["max_reserved_gb"]}, open(os.path.join(exp, "run_summary.json"), "w"), indent=1, sort_keys=True) print(f"[rank {rank}] past the export and push block", flush=True) segment = forced_stop # the same predicate, computed before the validation decision above if api is not None and not segment and rank == 0: r = hubckpt.push_and_prune(args.hub_repo, exp, "final", api, repo_type=args.hub_repo_type, token=os.environ.get("HF_TOKEN"), prune=args.prune) # .get, not [k]: `pruned` and `free_after_gb` are only present on some paths through # push_and_prune, and a KeyError here would throw away a completed run (E-023). say("final push:", json.dumps({k: r.get(k) for k in ("verify", "pruned", "free_before_gb", "free_after_gb", "prune_blocked_reason")}, default=str)[:900]) # The run is over and `latest.json` has to say so. `on_save` only ever rolls it at a 381-step # boundary, so without this the pointer stops at 3,810, the launcher's `start_step >= horizon` # check never fires, and every relaunch after completion trains four more steps and pushes again. if step >= steps_planned: hubckpt.write_latest_pointer(args.hub_repo, { "path_in_repo": "final", "step": step, "samples_consumed": samples, "tokens_consumed": consumed, "fingerprint": fingerprint, "seed": args.data_seed}, api, repo_type=args.hub_repo_type, token=os.environ.get("HF_TOKEN")) say(f"latest.json rolled to the completed run: final at step {step}") elif segment: # The session stopped on a checkpoint interval, and HubPush already pushed AND verified that # checkpoint and rolled latest.json to it. Re-pushing it as "final" would add a 1.7 GB LFS # revision per session for a copy the Hub already holds. say(f"segment boundary reached at step {step} of {steps_planned}: the checkpoint HubPush just " "verified is the resume point, so no separate final push") say("RUN_JSON " + json.dumps({"final_step": step, "val_ppl": ppl, "tokens_consumed": consumed, "segment_stop": segment, "steps_planned": steps_planned, "tok_per_s": round(consumed_this_run / max(1, elapsed)), "final_loss": final_loss, "prior_loss": prior, "micro_batch": args.micro_batch, "seq_len": args.seq_len, "accum": args.accum, "world": world, "params": pc["sum_numel"], "attn": args.attn, "grad_ckpt": bool(args.grad_ckpt), "optim": args.optim, "gc_nonreentrant": bool(args.gc_nonreentrant), "torch_compile": bool(args.torch_compile), "effective_attn": eff_attn, "peak_gpu_gb": mem["rank0_allocated_gb"], "peak_reserved_gb": mem["rank0_reserved_gb"], "peak_gpu_gb_max_rank": mem["max_allocated_gb"], "peak_reserved_gb_max_rank": mem["max_reserved_gb"], # So a caller cannot report success for a validation pass that never # happened: a null val_ppl with no error here would be a bug, and a null # with an error is a measured failure. "val_error": val_err, "val_skipped": forced_stop, "val_windows": n_val}, sort_keys=True)) def peak_stats(): """Peak CUDA memory, from every rank, with reserved bytes and the cross-rank maximum. An OOM mid-run is a lost session, so this is part of every preflight cell's evidence rather than something to discover at step 300. `max_memory_allocated` on rank 0 alone understates what the card is actually holding: the allocator reserves more than it hands out, and the other rank's peak is a separate number on a separate GPU that the run's own limit has to respect. """ out = {"cuda": torch.cuda.is_available(), "rank0_allocated_gb": None, "rank0_reserved_gb": None, "max_allocated_gb": None, "max_reserved_gb": None, "device": None} if not out["cuda"]: return out try: alloc = torch.cuda.max_memory_allocated() res = torch.cuda.max_memory_reserved() out["rank0_allocated_gb"] = round(alloc / 1e9, 2) out["rank0_reserved_gb"] = round(res / 1e9, 2) out["device"] = torch.cuda.get_device_name(0) t = torch.tensor([float(alloc), float(res)], device=alloc and "cuda" or "cpu") if torch.distributed.is_available() and torch.distributed.is_initialized(): torch.distributed.all_reduce(t, op=torch.distributed.ReduceOp.MAX) out["max_allocated_gb"] = round(float(t[0]) / 1e9, 2) out["max_reserved_gb"] = round(float(t[1]) / 1e9, 2) else: out["max_allocated_gb"] = out["rank0_allocated_gb"] out["max_reserved_gb"] = out["rank0_reserved_gb"] except Exception as e: out["error"] = f"{type(e).__name__}: {str(e)[:120]}" return out if __name__ == "__main__": main()