| |
| """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 |
| 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 |
|
|
|
|
| |
| |
| |
|
|
|
|
| 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, `<path>.A` / `<path>.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 |
|
|
|
|
| |
| |
| |
|
|
|
|
| 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}/<keys>" 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}/<len {len(a)} vs {len(b)}>"] |
| 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() |
| 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 |
|
|
| |
| 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 |
| |
| |
| |
| 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": |
| 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()) |
|
|
