| """N1: CUDA-graph replay == eager for an arm (emulated TP2 on one card, or the pair): drafts and confidences for |
| B = 1-4 at k = 5, n_out 5 and 4, after O2's bench windows. Usage: native_graphcheck ARM emu|pair OUT.json |
| """ |
| from __future__ import annotations |
|
|
| import json |
| import sys |
| from dataclasses import replace |
|
|
| import torch |
|
|
| from ...worker.checkpoint import Checkpoint |
| from ...worker.tools.parity_ref import make_taps |
| from ...worker.tools.tp2_bench import _EmbedOnly, fill_windows_tp, host_inp, setup_inputs |
| from ...worker.tp2 import check_visible |
| from ..tp2_native import ARMS, build |
|
|
| ROOT = "/opt/dsv41/x86host-data/dspark" |
|
|
|
|
| def main() -> None: |
| arm, mode, out = sys.argv[1], sys.argv[2], sys.argv[3] |
| check_visible(pair=(mode == "pair")) |
| d0 = torch.device("cuda", 0) |
| devs = [d0, torch.device("cuda", 1)] if mode == "pair" else [d0, d0] |
| cfg = ARMS[arm] |
| if len(sys.argv) > 4: |
| cfg = replace(cfg, dense_profile=sys.argv[4]) |
| m = build(cfg, devs, Checkpoint(ROOT), slots=8) |
| toks = setup_inputs(d0, ROOT) |
| taps = make_taps(_EmbedOnly(m.ranks[0].embed), toks, torch.Generator(device=d0).manual_seed(5)) |
| fill_windows_tp(m, taps.cpu()) |
| m.capture(batch_sizes=(1, 2, 3, 4), ks=(5,), insert_rows=(8,), n_outs=(4, 5), margins=False) |
| rep = {"arm": arm, "cfg": cfg.__dict__, "mode": mode, "cases": {}} |
| ok = True |
| for B in (1, 2, 3, 4): |
| inp = host_inp(B, toks.cpu()) |
| for n in (5, 4): |
| e_ids, e_conf, _, e_ids1 = m.draft(inp, 5, n_out=n) |
| g_ids, g_conf, _ = m.draft_graphed(inp, 5, n_out=n) |
| same = bool(torch.equal(e_ids.cpu(), g_ids[:B].cpu())) and bool(torch.equal(e_conf.cpu(), g_conf[:B].cpu())) |
| ranks_agree = e_ids1 is None or bool(torch.equal(e_ids.cpu(), e_ids1.cpu())) |
| rep["cases"][f"B{B}_n{n}"] = {"graph_equals_eager": same, "rank_ids_agree": ranks_agree} |
| ok &= same and ranks_agree |
| for st in m.ranks[0].stages: |
| if hasattr(st.nbank, "check"): |
| st.nbank.check() |
| rep["all_ok"] = ok |
| with open(out, "w") as f: |
| json.dump(rep, f, indent=1) |
| print(json.dumps(rep, indent=1)) |
| if not ok: |
| raise SystemExit("graph replay differs from eager") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
