#!/usr/bin/env python """Qwen3-8B LoRA SFT memory / throughput / resume probe on 4x A10G (P2-A). .venv/bin/python scripts/qwen3_sft_probe.py prepare --lens 4096 8192 16384 [--prompt-len 1024] torchrun --standalone --nproc_per_node=4 scripts/qwen3_sft_probe.py train CONFIG [dotted.overrides=...] .venv/bin/python scripts/qwen3_sft_probe.py report RUN_DIR [--accum A --world W --seq-len L] .venv/bin/python scripts/qwen3_sft_probe.py batches RUN_DIR What it measures, and why it is not just `scripts/sft.py`. `train` runs the SAME functions `scripts/sft.py` runs (`build_model`, `run_gates`, `train`: arch dispatch, LoRA wiring, the Trainer, the checkpoint format), with two additions the production entry deliberately does not carry: a `Tracker.log` wrapper that stamps every history record with `rank`, wall time and the CUDA allocator's peak, and a heavy-job guard identical to sft.py's. The numbers it reports are therefore the numbers a real launch would see. DATA is REAL TOKENS, not random ids: the formal statements and informal prefixes of the Goedel reference repo's miniF2F + MOBench files (604 problems), tokenized by the checkpoint's own tokenizer in the legacy venv, concatenated into one stream, and cut into rows of exactly `len` tokens at distinct offsets. Loss values on these rows are meaningless as a quality signal (rows are arbitrary windows); compute, memory and the sharding/resume mechanics are exactly those of real rows of that length. Every micro-batch is one row of the target length, so every micro-step costs the same, and tokens/s is `len / micro_step_time`. Import order is load-bearing: this module imports only the stdlib at top level. """ from __future__ import annotations import argparse import json import os import statistics import subprocess import sys REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) sys.path.insert(0, REPO) LEGACY_PY = "/apps/development/rlhf/nano-mistral-o1/.venv/bin/python" GOEDEL_REPO = "/apps/development/rlhf/reference-repos/Goedel-Prover-V2" SNAPSHOT = ("/apps/development/rlhf/hf_cache/hub/models--Goedel-LM--Goedel-Prover-V2-8B/snapshots/" "dfd02e6271a58375dfbf3ece0175277cf6b6a89a") OUT_ROOT = os.path.join(REPO, ".tmp", "p2", "p2a") STREAM = os.path.join(OUT_ROOT, "probe_stream.npy") _TOKENIZE = r''' import json, sys import numpy as np from transformers import AutoTokenizer ckpt, files, out = sys.argv[1], sys.argv[2].split(","), sys.argv[3] tok = AutoTokenizer.from_pretrained(ckpt) ids = [] for fn in files: for line in open(fn): d = json.loads(line) text = (d.get("informal_prefix") or "") + (d.get("formal_statement") or "") + "\n" ids.extend(tok(text, add_special_tokens=False)["input_ids"] + [tok.eos_token_id]) np.save(out, np.asarray(ids, dtype=np.int32)) print(json.dumps({"tokens": len(ids), "vocab": len(tok)})) ''' def cut_rows(stream, length: int, n_rows: int, seed: int = 0): """`n_rows` windows of exactly `length` tokens at distinct, seeded offsets (wrapping the stream when it is shorter than needed). Pure numpy; unit-tested.""" import numpy as np if length < 2: raise ValueError("length must be >= 2") n = len(stream) rng = np.random.default_rng(seed) offs = rng.choice(n, size=n_rows, replace=n_rows > n) doubled = np.concatenate([stream] * (length // n + 2)) return [doubled[o:o + length] for o in offs] def make_tiny_snapshot(out_dir: str, vocab: int = 2048, layers: int = 4, hidden: int = 256, seed: int = 0) -> dict: """A random-weight, structurally-faithful Qwen3 snapshot (GQA 8:4, head_dim 32, untied head) with the Goedel tokenizer files, for wiring smokes on the 4-GPU path in seconds instead of minutes. NOT a model of anything: it exists so FSDP2 + LoRA + checkpoint + resume bugs surface before the 8B load. Returns the config overrides a TOML needs to agree with it.""" import torch from align.model import qwen3 as q3 from align.model.qwen3_checkpoint import save_hf_qwen3 cfg = q3.Qwen3Config(name="qwen3-tiny", vocab_size=vocab, hidden_size=hidden, intermediate_size=2 * hidden, n_layers=layers, n_heads=8, n_kv_heads=4, head_dim=hidden // 8, max_position_embeddings=4096) torch.manual_seed(seed) m = q3.Qwen3ForCausalLM(cfg, attn_implementation="sdpa", device="cpu") snap = os.path.join(out_dir, "snapshots", "tiny") save_hf_qwen3(m, snap, dtype=torch.bfloat16, tokenizer_dir=SNAPSHOT) return {"paths.model_snapshot": os.path.join(out_dir, "snapshots"), "model.n_params": q3.count_params(cfg), "model.n_layers": layers, "model.vocab_size": vocab, "model.hidden_size": hidden} def prepare(lens: list[int], rows: int, prompt_len: int, out_root: str, vocab_mod: int | None = None) -> dict: import numpy as np from align.data.store import ShardWriter os.makedirs(out_root, exist_ok=True) if not os.path.isfile(STREAM): env = {"PATH": os.environ.get("PATH", ""), "HOME": os.environ.get("HOME", ""), "HF_HOME": "/apps/development/rlhf/hf_cache", "HF_HUB_OFFLINE": "1", "TRANSFORMERS_OFFLINE": "1", "TOKENIZERS_PARALLELISM": "false", "TMPDIR": OUT_ROOT} files = ",".join(os.path.join(GOEDEL_REPO, "dataset", f) for f in ("minif2f.jsonl", "MOBench.jsonl")) p = subprocess.run([LEGACY_PY, "-c", _TOKENIZE, SNAPSHOT, files, STREAM], capture_output=True, text=True, timeout=600, env=env, cwd=OUT_ROOT) if p.returncode != 0: raise RuntimeError(p.stderr[-2000:]) stream = np.load(STREAM) if vocab_mod: stream = stream % vocab_mod # tiny-vocab smokes only: ids must index the tiny embedding out = {"stream_tokens": int(len(stream))} for L in lens: d = os.path.join(out_root, f"L{L}", "sft", "train") if prompt_len >= L: raise ValueError(f"prompt_len {prompt_len} >= row length {L}") with ShardWriter(d, "sft-train-0000") as w: for i, row in enumerate(cut_rows(stream, L, rows, seed=L)): w.add_row(row.tolist(), problem_id=f"probe-{L}-{i}", subject="lean", boundary_positions=[prompt_len], prompt_len=prompt_len) out[f"L{L}"] = {"dir": d, "rows": rows} return out # --------------------------------------------------------------------------- # analysis (pure; unit-tested) # --------------------------------------------------------------------------- def read_jsonl(path: str) -> list[dict]: if not os.path.exists(path): return [] lines = [l for l in open(path).read().split("\n") if l.strip()] out = [] for i, l in enumerate(lines): try: out.append(json.loads(l)) except json.JSONDecodeError: if i != len(lines) - 1: raise return out def rank_histories(run_dir: str, world: int) -> dict[int, list[dict]]: """rank -> history records (rank 0 in `run_dir`, rank r in `run_dir/rank{r}`).""" return {r: read_jsonl(os.path.join(run_dir if r == 0 else os.path.join(run_dir, f"rank{r}"), "history.jsonl")) for r in range(world)} def rank_disjointness(hist: dict[int, list[dict]]) -> dict: """Do the ranks see DIFFERENT rows? `batch` is the list of dataset row indices of a micro-step. Returns the per-rank row sets and whether they are pairwise disjoint, plus per-micro-step disjointness (the contract: at micro-step m no two ranks share a row).""" seen = {r: [tuple(rec["batch"]) for rec in recs] for r, recs in hist.items()} sets = {r: {i for b in bs for i in b} for r, bs in seen.items()} ranks = sorted(sets) pairwise = all(not (sets[a] & sets[b]) for i, a in enumerate(ranks) for b in ranks[i + 1:]) n_micro = min(len(v) for v in seen.values()) per_micro = all(len({i for r in ranks for i in seen[r][m]}) == sum(len(seen[r][m]) for r in ranks) for m in range(n_micro)) return {"rows_per_rank": {r: len(s) for r, s in sets.items()}, "pairwise_disjoint": pairwise, "per_micro_step_disjoint": per_micro, "micro_steps_compared": n_micro, "first_batches": {r: [list(b) for b in seen[r][:4]] for r in ranks}} def step_stats(recs: list[dict], seq_len: int, accum: int, world: int, drop_first: int = 2, save_every: int | None = None) -> dict: """Median micro-step time from consecutive `t_wall` stamps of ONE rank (skipping the first `drop_first` intervals: allocator warmup, first-use kernels, and any interval containing a checkpoint save), peak memory over the run, and tokens/s = world * seq_len / micro_step_time (every rank processes one row of `seq_len` tokens per micro-step).""" recs = sorted(recs, key=lambda r: r["micro_step"]) dts = [b["t_wall"] - a["t_wall"] for a, b in zip(recs, recs[1:]) if b["micro_step"] == a["micro_step"] + 1 and a["micro_step"] >= drop_first and not (save_every and a["micro_step"] % save_every == 0)] if not dts: raise ValueError("no usable step intervals") med = statistics.median(dts) return { "n_intervals": len(dts), "micro_median_s": med, "micro_min_s": min(dts), "micro_max_s": max(dts), "optimizer_step_s": med * accum, "tokens_per_s_total": world * seq_len / med, "tokens_per_s_per_gpu": seq_len / med, "alloc_peak_gb": max(r["mem/alloc_peak_gb"] for r in recs), "reserved_peak_gb": max(r["mem/reserved_peak_gb"] for r in recs), "loss_first": recs[0]["loss"], "loss_last": recs[-1]["loss"], "grad_norm_last": next((r["grad_norm"] for r in reversed(recs) if r["grad_norm"] == r["grad_norm"]), float("nan")), } def dedupe_by_micro(recs: list[dict]) -> dict[int, dict]: """micro_step -> record, the LAST occurrence winning (a relaunch may re-log the boundary it resumed at).""" return {int(r["micro_step"]): r for r in recs} def compare_resume(ref: dict[int, list[dict]], sup: dict[int, list[dict]], stop_micro: int, floor_mult: float = 2.0, floor_pad: float = 5e-3) -> dict: """Does a SIGTERM'd-and-resumed run continue the uninterrupted one? Three separate questions, per rank, because a loss curve alone cannot answer the first two: * DATA: the resumed segment (micro > stop) saw EXACTLY the batches the reference saw at the same micro-steps (the stream offset round-trips), and nothing was replayed or skipped (micro-steps are contiguous across the stop). * CONTINUITY: |dloss| of the resumed segment vs the reference stays within `floor_mult * (pre-stop |dloss|) + floor_pad`. The pre-stop segment is the SAME computation in both runs, so its |dloss| measures pure run-to-run nondeterminism (NCCL reduction order, SDPA backward, 8-bit rounding); the resumed segment must not be worse than that floor by more than the factor. A vacuous comparison (one side empty) raises. * PROVENANCE is checked separately and exactly by the caller (the relaunch log names the checkpoint it loaded; this function does not claim it). """ out = {} for r in sorted(ref): a, b = dedupe_by_micro(ref[r]), dedupe_by_micro(sup[r]) common = sorted(set(a) & set(b)) pre = [abs(a[m]["loss"] - b[m]["loss"]) for m in common if m <= stop_micro] post = [abs(a[m]["loss"] - b[m]["loss"]) for m in common if m > stop_micro] if not pre or not post: raise ValueError(f"rank {r}: need micro-steps on both sides of the stop ({stop_micro}); pre={len(pre)} post={len(post)}") batches_equal = all(a[m]["batch"] == b[m]["batch"] for m in common) contiguous = sorted(b) == list(range(min(b), max(b) + 1)) tol = floor_mult * max(pre) + floor_pad out[r] = {"n_pre": len(pre), "n_post": len(post), "batches_equal_all_common": batches_equal, "micro_steps_contiguous": contiguous, "pre_stop_max_dloss": max(pre), "resumed_max_dloss": max(post), "tolerance": tol, "continuity_passed": max(post) <= tol, "last_micro_ref": max(a), "last_micro_resumed": max(b)} out["passed"] = all(v["batches_equal_all_common"] and v["micro_steps_contiguous"] and v["continuity_passed"] and v["last_micro_ref"] == v["last_micro_resumed"] for k, v in out.items()) return out def adapter_diff(lora_a: dict, lora_b: dict) -> dict: """Relative difference of two adapter states (`lora_state_dict` format, `.A` / `.B`). The strongest resume check a loss curve cannot give: the FINAL adapter of an interrupted + resumed run against an uninterrupted run of the same schedule. `rel_B` is ||B_a - B_b|| / ||B_b||, global over all layers. B starts at EXACTLY zero, so ||B|| is the whole distance trained; a resume that restarts the adapter from the bare base (B = 0) after k of n optimizer steps lands near rel_B ~ (1 - k/n)-ish, orders of magnitude above run-to-run nondeterminism (flash-attention backward atomics) which stays at the 1e-3 level. """ import torch if set(lora_a) != set(lora_b): raise KeyError(f"adapter key sets differ: {sorted(set(lora_a) ^ set(lora_b))[:4]}") num = {"A": 0.0, "B": 0.0} den = {"A": 0.0, "B": 0.0} for k in lora_a: kind = k.rsplit(".", 1)[1] a, b = lora_a[k].double(), lora_b[k].double() if a.shape != b.shape: raise ValueError(f"{k}: {tuple(a.shape)} vs {tuple(b.shape)}") num[kind] += float((a - b).pow(2).sum()) den[kind] += float(b.pow(2).sum()) if den["B"] == 0.0: raise ValueError("reference adapter has B == 0: nothing was trained, the comparison is vacuous") return {"rel_A": (num["A"] / den["A"]) ** 0.5, "rel_B": (num["B"] / den["B"]) ** 0.5, "norm_B": den["B"] ** 0.5, "n_tensors": len(lora_a)} def adapter_diff_main(a) -> int: import torch ld = lambda d: torch.load(os.path.join(d, "optimizer_rank0.pt"), weights_only=True, map_location="cpu")["lora"] res = adapter_diff(ld(a.ckpt_a), ld(a.ckpt_b)) print(json.dumps(res, indent=2)) return 0 # --------------------------------------------------------------------------- # train (torchrun target) # --------------------------------------------------------------------------- def tree_equal(a, b, path: str = "") -> list[str]: """Paths at which two nested state structures (dict / list / tensor / scalar) are NOT bitwise equal. Tensors compare by dtype, shape and `torch.equal`; everything else by `==`. Empty list = identical. Pure; unit-tested.""" import torch if isinstance(a, dict) and isinstance(b, dict): out = [f"{path}/" for _ in [0] if set(a) != set(b)] for k in sorted(set(a) & set(b), key=str): out += tree_equal(a[k], b[k], f"{path}/{k}") return out if isinstance(a, (list, tuple)) and isinstance(b, (list, tuple)): if len(a) != len(b): return [f"{path}/"] return [d for i, (x, y) in enumerate(zip(a, b)) for d in tree_equal(x, y, f"{path}[{i}]")] if isinstance(a, torch.Tensor) and isinstance(b, torch.Tensor): same = a.dtype == b.dtype and a.shape == b.shape and torch.equal(a.cpu(), b.cpu()) return [] if same else [path] return [] if a == b else [path] def _verify_resume_only(trainer): """Stand-in for `Trainer.train` under --verify-resume-only: after `resume_from_latest()` the live sharded state (the gathered adapter, the 8-bit moments, the step counters) must equal what the checkpoint holds, BITWISE. No step is taken, so there is no GPU nondeterminism to excuse a difference -- this isolates the scatter path (`distribute_tensor` / local-shard copies) from training noise, which the adapter-diff of two uninterrupted runs (rel_B up to ~0.4) cannot do.""" import torch from align.utils.dist import rank R = rank() ck = trainer.latest_checkpoint() if ck is None: raise RuntimeError("--verify-resume-only needs an existing checkpoint to resume from") saved = torch.load(os.path.join(ck, f"optimizer_rank{R}.pt"), weights_only=True, map_location="cpu") live = trainer.train_module.optim_state_dict() # collective: every rank gathers the adapter diffs = tree_equal(saved, live) meta = torch.load(os.path.join(ck, "resume.pt"), weights_only=True, map_location="cpu") counters = {"micro_step": (meta["micro_step"], trainer.micro_step), "optimizer_step": (meta["optimizer_step"], trainer.optimizer_step)} bad_counters = {k: v for k, v in counters.items() if v[0] != v[1]} res = {"rank": R, "checkpoint": ck, "state_tensors_compared": len(tree_leaves(saved)), "unequal_paths": diffs[:10], "n_unequal": len(diffs), "counters": counters, "bad_counters": bad_counters, "passed": not diffs and not bad_counters} with open(os.path.join(trainer.run_dir, f"verify_resume_rank{R}.json"), "w") as f: json.dump(res, f, indent=2, default=str) print(json.dumps(res, default=str), flush=True) trainer.stop_reason = "verify_only" trainer.verify_passed = bool(res["passed"]) return {"optimizer_step": trainer.optimizer_step, "verify_only": True, **{"passed": res["passed"]}} def final_exit_code(trainer, verify_only: bool, any_rank=None) -> int: """Process exit status. Under --verify-resume-only a FAILED (or never-recorded) bitwise verdict is nonzero on EVERY rank (the verdict is MAX-reduced first), so a wrapper / CI cannot read success from a run whose per-rank verdict only went to verify_resume_rank{R}.json and stdout.""" if verify_only: if any_rank is None: from align.utils.dist import any_rank if any_rank(not getattr(trainer, "verify_passed", False)): return 1 return trainer.stop_exit_code or 0 def tree_leaves(x) -> list: import torch if isinstance(x, dict): return [l for v in x.values() for l in tree_leaves(v)] if isinstance(x, (list, tuple)): return [l for v in x for l in tree_leaves(v)] return [x] if isinstance(x, torch.Tensor) else [] def train_main(argv: list[str]) -> int: import importlib.util # scripts/sft.py by path: the same build_model / run_gates / train the production entry runs. spec = importlib.util.spec_from_file_location("sft_script", os.path.join(REPO, "scripts", "sft.py")) sft_script = importlib.util.module_from_spec(spec) spec.loader.exec_module(sft_script) from align.config import load_config from align.train.stage_gates import GateLedger, eval_source_hash, parse_args from align.utils.guard import apply_resource_guards verify_only = "--verify-resume-only" in argv argv = [x for x in argv if x != "--verify-resume-only"] args = parse_args(argv) cfg = load_config(args.config, args.overrides) apply_resource_guards(cpu_threads=8, heavy=True, label="p2a-qwen3-sft-probe", rss_cfg=cfg) import logging import time logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s") from align.utils.dist import init_distributed run_dir = os.path.join(cfg.paths.checkpoint_dir, "sft") init_distributed(run_dir=run_dir) import torch from align.train.stage_gates import assert_distinct_devices from align.utils.dist import barrier, rank assert_distinct_devices() R = rank() os.makedirs(run_dir, exist_ok=True) with open(os.path.join(run_dir, f"rank{R}.pid"), "w") as f: f.write(str(os.getpid())) from align.utils import tracking orig_log = tracking.Tracker.log def log(self, step, metrics=None, **kw): torch.cuda.synchronize() m = {**(dict(metrics) if metrics else {}), **kw} m.update({"rank": R, "t_wall": time.time(), "mem/alloc_peak_gb": torch.cuda.max_memory_allocated() / 2**30, "mem/reserved_peak_gb": torch.cuda.max_memory_reserved() / 2**30}) orig_log(self, step, m) tracking.Tracker.log = log ledger = GateLedger(run_dir, eval_source_hash=eval_source_hash()) banner = sft_script.run_gates(cfg, ledger) from align.registry import render_banner if R == 0: print(render_banner(banner, title=f"align SFT probe :: {cfg.run_name}"), flush=True) model = sft_script.build_model(cfg) torch.cuda.synchronize() post_wrap = torch.cuda.memory_allocated() / 2**30 # The load phase held a full 16.4 GB bf16 replica per rank before FSDP sharded it; the caching # allocator keeps that block reserved. Without this, `reserved_peak` reports the LOAD, not the # training step (measured: 18.875 GiB at both seq 4096 and 8192). torch.cuda.empty_cache() torch.cuda.reset_peak_memory_stats() if verify_only: from align.train import trainer as trainer_mod trainer_mod.Trainer.train = _verify_resume_only trainer = sft_script.train(cfg, ledger, model, banner) barrier() rec = {"rank": R, "post_wrap_alloc_gb": post_wrap, "alloc_peak_gb": torch.cuda.max_memory_allocated() / 2**30, "reserved_peak_gb": torch.cuda.max_memory_reserved() / 2**30, "stop_reason": trainer.stop_reason, "micro_step": trainer.micro_step, "optimizer_step": trainer.optimizer_step, "load_stats": getattr(model, "_stage_load_stats", {})} with open(os.path.join(run_dir, f"final_rank{R}.json"), "w") as f: json.dump(rec, f, indent=2) return final_exit_code(trainer, verify_only) def report_main(a) -> int: run = os.path.join(a.run_dir, "sft") if os.path.isdir(os.path.join(a.run_dir, "sft")) else a.run_dir hist = rank_histories(run, a.world) out = {"disjoint": rank_disjointness(hist), "per_rank": {r: step_stats(h, a.seq_len, a.accum, a.world, save_every=a.save_every) for r, h in hist.items()}} print(json.dumps(out, indent=2)) return 0 def compare_main(a) -> int: pick = lambda d: os.path.join(d, "sft") if os.path.isdir(os.path.join(d, "sft")) else d res = compare_resume(rank_histories(pick(a.ref), a.world), rank_histories(pick(a.sup), a.world), a.stop_micro) print(json.dumps(res, indent=2)) return 0 if res["passed"] else 1 def batches_main(a) -> int: run = os.path.join(a.run_dir, "sft") if os.path.isdir(os.path.join(a.run_dir, "sft")) else a.run_dir hist = rank_histories(run, a.world) for r, recs in hist.items(): print(f"rank {r}: " + " ".join(f"m{rec['micro_step']}:{rec['batch']}" for rec in recs[:12])) print(json.dumps(rank_disjointness(hist), indent=2)) return 0 def main(argv=None) -> int: argv = list(sys.argv[1:] if argv is None else argv) if argv and argv[0] == "train": # config + dotted overrides are parsed by the repo's own parser return train_main(argv[1:]) ap = argparse.ArgumentParser(description=__doc__.split("\n")[0]) sub = ap.add_subparsers(dest="cmd", required=True) p = sub.add_parser("prepare") p.add_argument("--lens", type=int, nargs="+", required=True) p.add_argument("--rows", type=int, default=256) p.add_argument("--prompt-len", type=int, default=1024) p.add_argument("--out-root", default=os.path.join(OUT_ROOT, "shards")) p.add_argument("--vocab-mod", type=int, default=None) t = sub.add_parser("make-tiny") t.add_argument("--out-dir", default=os.path.join(OUT_ROOT, "tiny_model")) for name in ("report", "batches"): r = sub.add_parser(name) r.add_argument("run_dir") r.add_argument("--world", type=int, default=4) r.add_argument("--accum", type=int, default=1) r.add_argument("--seq-len", type=int, default=4096) r.add_argument("--save-every", type=int, default=None) d = sub.add_parser("adapter-diff") d.add_argument("ckpt_a") d.add_argument("ckpt_b") c = sub.add_parser("compare") c.add_argument("ref") c.add_argument("sup") c.add_argument("--stop-micro", type=int, required=True) c.add_argument("--world", type=int, default=4) a = ap.parse_args(argv) if a.cmd == "adapter-diff": return adapter_diff_main(a) if a.cmd == "compare": return compare_main(a) if a.cmd == "make-tiny": print(json.dumps(make_tiny_snapshot(a.out_dir), indent=2)) return 0 if a.cmd == "prepare": print(json.dumps(prepare(a.lens, a.rows, a.prompt_len, a.out_root, a.vocab_mod), indent=2)) return 0 return report_main(a) if a.cmd == "report" else batches_main(a) if __name__ == "__main__": raise SystemExit(main())