File size: 6,999 Bytes
0693474 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 | """N1: memory (serving + load peak) and draft-step time of an arm, per card.
solo one rank of the arm alone on GPU-7744b44b (TPDrafter solo_rank: the rank's own tensors and kernels, the
collectives replaced by local stand-ins of the same shapes): that card's real serving footprint and load peak,
and the rank's compute time. Run rank 0 and rank 1 (their shards differ only in the vocabulary / head / group
columns they own).
pair both ranks on the NVLink pair (PAIR window only): real collectives, both cards' footprints.
Per arm: load seconds, torch peak allocated / reserved during the load and after graph capture, nvidia-smi process
memory, GPU clocks / power before and after, graph replay p50 for B = 1-4 at k = 5, n_out 5 and 4, insert 8-128
rows. Same inputs as lane E's tp2_bench (O2's D5 bench), so the K2 control numbers are comparable.
Usage (x86host): python -m offload.dspark.native.tools.native_bench ARM solo|pair OUT.json [RANK] [ITERS]
"""
from __future__ import annotations
import json
import hashlib
import os
import subprocess
import sys
import time
from pathlib import Path
import torch
from ...worker.checkpoint import Checkpoint
from ...worker.tools.bench import time_fn
from ...worker.tools.parity_ref import make_taps
from ...worker.tools.tp2_bench import _EmbedOnly, fill_windows_tp, gpu_state, host_inp, setup_inputs
from ...worker.tp2 import check_visible
from ..tp2_native import ARMS, build
from .n2.profiling import breakdown
ROOT = "/opt/dsv41/x86host-data/dspark"
def smi_process_mib() -> list:
out = subprocess.run(["nvidia-smi", "--query-compute-apps=pid,gpu_uuid,used_memory", "--format=csv,noheader"],
capture_output=True, text=True).stdout
me = str(os.getpid())
return [l.strip() for l in out.splitlines() if l.strip().startswith(me + ",")]
def main() -> None:
arm, mode, out_path = sys.argv[1], sys.argv[2], sys.argv[3]
rank = int(sys.argv[4]) if len(sys.argv) > 4 else 0
iters = int(sys.argv[5]) if len(sys.argv) > 5 else 200
if arm not in ARMS:
raise SystemExit(f"arm {arm!r} not in {sorted(ARMS)}")
vis = check_visible(pair=(mode == "pair"))
uuids = vis.split(",")
t0 = time.time()
rep: dict = {"arm": arm, "cfg": ARMS[arm].__dict__, "mode": mode, "rank": rank if mode == "solo" else None,
"devices": uuids, "iters": iters, "gpu_before": {u: gpu_state(u) for u in uuids}}
source_root = Path(__file__).resolve().parents[2]
rep["source_sha256"] = {str(p.relative_to(source_root)): hashlib.sha256(p.read_bytes()).hexdigest()
for p in (source_root / "native/tp2_native.py", source_root / "native/kx.py",
source_root / "kernels/fusedmoe/fusedmoe.cu",
source_root / "kernels/fusedmoe/epilogue.py",
source_root / "kernels/fp8dense/dense.cu",
source_root / "kernels/fp8dense/epilogue.py")}
for name in ("native/attention.py", "native/ops.py", "worker/fused.py", "worker/tp2.py"):
p = source_root / name
rep["source_sha256"][name] = hashlib.sha256(p.read_bytes()).hexdigest() if p.exists() else None
ck = Checkpoint(ROOT)
d0 = torch.device("cuda", 0)
devs = [d0, torch.device("cuda", 1)] if mode == "pair" else [d0, d0]
for d in set(devs):
torch.empty(1, device=d) # the CUDA context and allocator exist before the counters reset
torch.cuda.reset_peak_memory_stats(d)
tl0 = time.time()
kw = {"slots": 8}
if mode == "solo":
kw["solo_rank"] = rank
m = build(arm, devs, ck, **kw)
rep["load_seconds"] = round(time.time() - tl0, 1)
rep["load_peak"] = {str(d): {"max_allocated": torch.cuda.max_memory_allocated(d),
"max_reserved": torch.cuda.max_memory_reserved(d),
"allocated": torch.cuda.memory_allocated(d)} for d in sorted(set(devs), key=str)}
rep["load_nvidia_smi"] = smi_process_mib()
if getattr(m, "load_stats", None) is not None:
rep["load_stats"] = m.load_stats.__dict__
r0 = m.ALL[0]
toks = setup_inputs(devs[r0], ROOT)
taps = make_taps(_EmbedOnly(m.ranks[r0].embed), toks, torch.Generator(device=devs[r0]).manual_seed(5))
taps_cpu, toks_cpu = taps.cpu(), toks.cpu()
fill_windows_tp(m, taps_cpu)
rep["capture_seconds"] = round(m.capture(batch_sizes=(1, 2, 3, 4), ks=(5,), insert_rows=(8, 16, 32, 64, 128),
n_outs=(4, 5), margins=False), 2)
rep["draft"] = {}
rep["insert"] = {}
with torch.cuda.device(devs[r0]):
for B in (1, 2, 3, 4):
inp = host_inp(B, toks_cpu)
for n in (5, 4):
# A short ten-call warmup does not reliably settle automatic
# boost after CPU-heavy loading. No clocks/power are changed.
settle_until = time.perf_counter() + 0.5
while time.perf_counter() < settle_until:
m.draft_graphed(inp, 5, n_out=n)
torch.cuda.synchronize()
before_case = {u: gpu_state(u) for u in uuids}
g = time_fn(lambda: m.draft_graphed(inp, 5, n_out=n), iters)
g["settle_seconds"] = 0.5
g["gpu_before"] = before_case
g["gpu_after"] = {u: gpu_state(u) for u in uuids}
rep["draft"][f"B{B}_k5_n{n}"] = g
print(f"{arm} B{B} n_out{n}: p50 {g['p50_us']} us", flush=True)
bd = breakdown(lambda: m.draft_graphed(inp, 5, n_out=n))
rep["draft"][f"B{B}_k5_n{n}"]["breakdown"] = bd
print(json.dumps({dev: row["us"] for dev, row in bd.items()}), flush=True)
for nrow in (8, 32, 128):
sr = torch.full((nrow,), 7, dtype=torch.int64).pin_memory()
pr = torch.arange(2000, 2000 + nrow).pin_memory()
tt = taps_cpu[:nrow].clone().pin_memory()
rep["insert"][f"rows{nrow}"] = time_fn(lambda: m.insert_taps_graphed(sr, pr, tt), iters)
rep["memory"] = m.memory_table()
rep["serving_nvidia_smi"] = smi_process_mib()
rep["serving_peak"] = {str(d): {"max_allocated": torch.cuda.max_memory_allocated(d),
"max_reserved": torch.cuda.max_memory_reserved(d)}
for d in sorted(set(devs), key=str)}
rep["gpu_after"] = {u: gpu_state(u) for u in uuids}
rep["seconds"] = round(time.time() - t0, 1)
with open(out_path, "w") as f:
json.dump(rep, f, indent=1, default=str)
print(json.dumps({k: rep[k] for k in ("arm", "rank", "load_seconds", "load_peak", "load_nvidia_smi",
"serving_nvidia_smi", "capture_seconds")}, indent=1, default=str))
print("WROTE", out_path, rep["seconds"], "s")
if __name__ == "__main__":
main()
|