William Lin
552d7dfadc
- fastvideo-kernel/wan22_vae_decode: the Wan2.2 (TI2V-5B) VAE decoder on sm_100a / sm_103a at 704x1248, built on first use against fastvideo-kernel's CUTLASS submodule. It frees intermediates early and carries 4 latent frames per pass: a 5 s video peaks at 20.3 GiB, down from 94.6 GiB in the kernel's first version, with the same output. - WanVAEConfig.decode_kernel routes AutoencoderKLWan.decode through it where it applies (bf16 with fp32 accumulation, about 7x faster than the fp32 decode on a B300, PSNR 61-63 dB, SSIM 0.998 against it); decode_kernel_parallel splits the decode over the sequence-parallel group (2, 4 or 8 GPUs) and broadcasts the video, with the ranks agreeing on whether every one of them could build the kernel. - PAN2's DiT runs fused Triton kernels for Q/K RMSNorm + RoPE and LayerNorm + modulation (pan2_fusions, set_fused_kernels). - PAN2PipelineConfig turns all three on (vae_decode_kernel, vae_decode_kernel_parallel, dit_fused_kernels). - VideoGenerator copies a decoded RGB video's frames to the host once instead of once per frame. - The PAN2 presets sample at guidance 2.5, the reference launcher's default since 2026-10-23.
772 lines
37 KiB
Python
772 lines
37 KiB
Python
"""Correctness, timing and library-kernel audit of a solution: the Wan2.2 VAE decode of latents into video, on one, two
|
|
or four GB200 GPUs of one node. The task's CI runs this same file (see ci.py).
|
|
|
|
The task (spec.py has the details)
|
|
--------
|
|
A solution decodes latents z [1, 48, T, 44, 78] (bf16) into the video [1, 3, 4T-3, 704, 1248] (bf16, in [-1, 1]) with
|
|
the Wan2.2 VAE's decoder, as diffusers' AutoencoderKLWan decodes them (reference.py in the kernel repo). It is a
|
|
directory (default solution/, see loader.py) compiled into one extension module that exports
|
|
|
|
run(z, weights, rank, world, pg, out) positional, destination passing
|
|
prepare(weights, rank, world, pg) -> list[Tensor] optional
|
|
|
|
called in one process per GPU: z the full latents on every rank; weights the decoder's 110 bf16 tensors on the rank's
|
|
GPU, in the order of spec.WEIGHTS (or what prepare returned); pg torch.distributed's NCCL ProcessGroup over the ranks;
|
|
out the video on rank 0, from torch.empty and never zeroed (every element must be written), and an empty tensor on the
|
|
other ranks. When run has returned on every rank, rank 0's out holds the video. prepare, if exported, runs once per
|
|
process before any timed call and is not timed; its result replaces the weights in every run.
|
|
|
|
The workloads (spec.WORKLOADS) are clips decoded at a world size, named <clip>x<world>. A clip is the first T latent
|
|
frames of one of three real 20 s clips, named by its video frames: f1, f5, f17 (T = 1, 2, 5) and f121, f241, f481
|
|
(T = 31, 61, 121). The long clips at 1, 2 and 4 GPUs are the nine timed workloads (f121x1 ... f481x4), whose geomean
|
|
is the objective; the short clips at 1 and 4 GPUs (f1x1 ... f17x4) are checked and never timed.
|
|
|
|
Running it
|
|
----------
|
|
In a job with as many GPUs as the largest world size it runs, with the task's data at $DATA (default /data), from the
|
|
repo root:
|
|
|
|
python tools/bench.py every workload this job has the GPUs for: check, audit, time
|
|
python tools/bench.py --workload f5x1 f17x1 a subset; the short clips take seconds
|
|
python tools/bench.py --workload f121x1 --no-time correctness only
|
|
python tools/bench.py --kernels also list the audited calls' kernels by GPU time
|
|
python tools/bench.py --solution other another solution directory
|
|
|
|
A run
|
|
-----
|
|
The solution is built once, by this process, exactly as it is deployed (loader.py; a development copy, tools/, takes a
|
|
build its node already made of the same sources, flags and toolchain from the node's cache, unless --no-cache). Then,
|
|
per world size in turn (1, 2, 4), torch.multiprocessing spawns one process per rank, rank r on GPU r of the job's
|
|
CUDA_VISIBLE_DEVICES. Each initializes NCCL over 127.0.0.1 with its device (the default group, the solution's pg) and
|
|
a gloo group for the harness's own barriers, imports the build (in its working directory, as deployed), moves the
|
|
weights ($DATA/decoder_weights_bf16.pt, checked against spec.WEIGHTS) to its GPU and calls prepare, if the module
|
|
exports it. The ranks then run that world size's workloads in spec.WORKLOADS' order. A rank that raises, dies or
|
|
hangs (WORKLOAD_SEC per workload, SESSION_START_SEC before the first) fails the workload it was at: all the ranks'
|
|
processes are stopped, and the workloads after it run in new ones. This process runs no solution code.
|
|
|
|
Every call gets z at a fresh address (a copy of the clip's latents) and out at a fresh address, 256-byte aligned and
|
|
filled with NaN on rank 0, so that an element run leaves unwritten fails the check. run must not write z, nor the
|
|
weights (compared after each workload with what prepare left), nor start threads: each fails the workload.
|
|
|
|
Correctness
|
|
-----------
|
|
The first call of a workload, its warm-up, is checked on rank 0 against the clip's fp32 decode ($DATA/ref_<clip>.pt):
|
|
every element finite and within [-1, 1]; and with d = out - ref in fp32, max |d| <= min(2, 3 x bf16_max_abs) and
|
|
rms(d) <= 3 x bf16_rms, the error of the release's bf16 decode of that clip ($DATA/calibration.json; spec.py,
|
|
Correctness). The line shows max |d| and rms(d), each with its allowance and the share of it used, and the PSNR,
|
|
10 log10(4 / mse). Every later call's output (the audited call's, the trials') must agree with the checked one:
|
|
lose none of the elements it wrote (none left NaN or overflowed), and be within the same allowance of it everywhere
|
|
else. A race, or a state carried from call to call, fails the workload.
|
|
|
|
Timing (the timed workloads)
|
|
------
|
|
After the warm-up come the trials: 3 when the warm-up took under 1.5 s, else 2. A trial is one call timed by wall
|
|
clock between barriers on all ranks: synchronize the device, barrier, start, run, synchronize, barrier, stop. The
|
|
time of record is rank 0's; the latency of record the mean of the trials, and the fastest trial is shown too. While
|
|
the trials run, this process reads the SM clock of the world's GPUs through NVML every CLOCK_PERIOD_S (a rank inside
|
|
run may hold its GIL for the whole call, so the ranks do not read it themselves); the mean of rank 0's GPU is shown
|
|
with the latency, and the range over the world's GPUs where they differ by 2% or more: the clocks are not locked
|
|
alike on every GPU, and this says at what clock the times were taken. Which processes hold the world's GPUs is
|
|
asked (nvidia-smi) before every trial and after the last: a workload timed while another process held one of them is
|
|
marked SHARED, since its times do not compare with those of another job.
|
|
|
|
Audit
|
|
-----
|
|
At world 1, one more call per timed clip runs under CUPTI, and each kernel is attributed to the library that launched
|
|
it (see audit.py). Library kernels (cuBLAS/cuBLASLt, cuDNN, ATen) are allowed as placeholders: their count per call
|
|
and their share of the call's GPU time show where the solution still depends on them. NCCL's kernels are
|
|
communication, reported apart and not counted as placeholders. --kernels lists the audited calls' kernels.
|
|
|
|
Bound
|
|
-----
|
|
Each timed workload's latency is shown with its bound (bounds.py with spec.bounds: the decode's FLOPs spread evenly
|
|
over the world's GPUs at the 16-bit tensor-core peak, or the video and the weights at the memory bandwidth, whichever
|
|
is slower) and the gap, latency / bound. The peak assumes the SM clock GUILD_SM_CLOCK_MHZ where the environment sets
|
|
it (the guild's: 2062 MHz, the GB200's highest), else the GPU's highest.
|
|
|
|
In CI, and from tools/
|
|
----------------------
|
|
ci.py runs this file in child processes: --save DIR has rank 0 save each workload's checked output there (node-local
|
|
scratch) instead of checking it, so that CI judges it in a process that runs no solution code; --no-check skips the
|
|
check (the timing runs), --no-audit the audit, and --json writes the report. The kernel repo holds copies of this
|
|
file, spec.py, loader.py, audit.py and bounds.py in tools/, which run as they do here: they read the data at $DATA
|
|
(default /data) and the solution from solution/ under the current directory.
|
|
"""
|
|
|
|
import argparse
|
|
import collections
|
|
import ctypes
|
|
import functools
|
|
import json
|
|
import math
|
|
import os
|
|
import random
|
|
import socket
|
|
import statistics
|
|
import subprocess
|
|
import sys
|
|
import threading
|
|
import time
|
|
import traceback
|
|
from datetime import timedelta
|
|
from multiprocessing import connection
|
|
|
|
# The caching allocator stays as torch configures it by default: without expandable segments, whose memory
|
|
# cudaIpcGetMemHandle cannot export, since a solution may share its buffers between ranks by CUDA IPC.
|
|
import torch
|
|
import torch.distributed as dist
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
import audit # noqa: E402
|
|
import bounds # noqa: E402
|
|
import loader # noqa: E402
|
|
import spec # noqa: E402
|
|
|
|
# Trials per timed workload, fewer when one call takes long.
|
|
TRIALS, LONG_TRIALS, LONG_CALL_S = 3, 2, 1.5
|
|
# How often this process reads rank 0's SM clock while ranks run.
|
|
CLOCK_PERIOD_S = 0.02
|
|
NVML_CLOCK_SM = 1 # the nvmlClockType_t of the SM clock
|
|
# Limits that stop a runaway: a session's start (NCCL, the module, the weights, prepare), and one workload (its calls,
|
|
# its check or its save); a call of the release's bf16 decode takes 15 s at the longest clip.
|
|
SESSION_START_SEC = 900
|
|
WORKLOAD_SEC = 1800
|
|
STOP_GRACE_SEC = 15
|
|
# Elements of a video compared at once, in fp32 (1 GiB per operand).
|
|
CHUNK_ELEMENTS = 1 << 28
|
|
|
|
|
|
def gpu_holders():
|
|
"""{GPU UUID (lower case, without "GPU-"): ids of the processes that hold it} now, as nvidia-smi gives them, or
|
|
None when it cannot say."""
|
|
try:
|
|
got = subprocess.run(["nvidia-smi", "--query-compute-apps=gpu_uuid,pid", "--format=csv,noheader"],
|
|
capture_output=True, text=True, timeout=30)
|
|
except (OSError, subprocess.TimeoutExpired):
|
|
return None
|
|
if got.returncode != 0:
|
|
return None
|
|
out = collections.defaultdict(set)
|
|
for line in got.stdout.splitlines():
|
|
gpu, _, pid = line.partition(",")
|
|
if pid.strip().isdigit():
|
|
out[gpu.strip().lower().removeprefix("gpu-")].add(int(pid))
|
|
return out
|
|
|
|
|
|
def others(questions, uuids, ours):
|
|
"""How many processes besides this run's held one of these GPUs at any of the questions (gpu_holders' answers), or
|
|
None when one could not be asked. This run's processes are known by their ids where nvidia-smi gives the ids this
|
|
process sees (one PID namespace, the usual case: then a rank that also holds another rank's GPU is still one of
|
|
ours); otherwise each GPU is taken to hold one of ours, the rank on it."""
|
|
if not questions or any(q is None for q in questions):
|
|
return None
|
|
total = 0
|
|
for uuid in uuids:
|
|
seen = set().union(*(q.get(uuid, set()) for q in questions))
|
|
total += len(seen - ours) if seen & ours else max(0, len(seen) - 1)
|
|
return total
|
|
|
|
|
|
@functools.cache
|
|
def sm_clock(uuid):
|
|
"""A function that reads the SM clock of the GPU with this UUID now, in MHz (None when a read fails), or None where
|
|
NVML cannot be loaded or does not know the GPU. NVML in this process takes microseconds per read."""
|
|
try:
|
|
nvml = ctypes.CDLL("libnvidia-ml.so.1")
|
|
if nvml.nvmlInit_v2() != 0:
|
|
return None
|
|
handle, mhz = ctypes.c_void_p(), ctypes.c_uint()
|
|
if not uuid or nvml.nvmlDeviceGetHandleByUUID(f"GPU-{uuid}".encode(), ctypes.byref(handle)) != 0:
|
|
return None
|
|
except (OSError, AttributeError):
|
|
return None
|
|
|
|
def read():
|
|
ok = nvml.nvmlDeviceGetClockInfo(handle, NVML_CLOCK_SM, ctypes.byref(mhz)) == 0
|
|
return mhz.value if ok and mhz.value else None
|
|
|
|
return read if read() else None
|
|
|
|
|
|
class Clock:
|
|
"""The SM clocks of the world's GPUs (rank order), read every CLOCK_PERIOD_S by a thread of this process, as
|
|
(monotonic time, [MHz or None per GPU])."""
|
|
|
|
def __init__(self, uuids):
|
|
self.reads, self.samples, self.done = [sm_clock(u) for u in uuids], [], threading.Event()
|
|
self.thread = threading.Thread(target=self._run, daemon=True) if any(self.reads) else None
|
|
if self.thread:
|
|
self.thread.start()
|
|
|
|
def _run(self):
|
|
while not self.done.wait(CLOCK_PERIOD_S):
|
|
self.samples.append((time.monotonic(), [read() if read else None for read in self.reads]))
|
|
|
|
def stop(self):
|
|
self.done.set()
|
|
if self.thread:
|
|
self.thread.join()
|
|
|
|
def means(self, windows):
|
|
"""Per GPU, the mean of its reads within these (start, stop) windows of time.monotonic(), or None without
|
|
any."""
|
|
inside = [mhz for t, mhz in list(self.samples) if any(a <= t <= b for a, b in windows)]
|
|
out = []
|
|
for i in range(len(self.reads)):
|
|
got = [m[i] for m in inside if m[i]]
|
|
out.append(statistics.mean(got) if got else None)
|
|
return out
|
|
|
|
|
|
class Pool:
|
|
"""A workload's buffers for one tensor: each call's at a fresh address, shifted by multiples of 256 bytes."""
|
|
|
|
ALIGN, MAX_SHIFT = 256, 2048
|
|
|
|
def __init__(self, shape, dtype, device, calls, seed):
|
|
rng = random.Random(seed)
|
|
self.offsets = [0]
|
|
for _ in range(calls - 1):
|
|
self.offsets.append(self.offsets[-1] + rng.randint(1, self.MAX_SHIFT // self.ALIGN))
|
|
self.shape, self.numel = tuple(shape), math.prod(shape)
|
|
self.block = max(1, self.ALIGN // torch.empty(0, dtype=dtype).element_size())
|
|
self.buffer = torch.empty(self.numel + self.offsets[-1] * self.block, dtype=dtype, device=device)
|
|
self.calls = 0
|
|
|
|
def next(self):
|
|
start = self.offsets[self.calls] * self.block
|
|
self.calls += 1
|
|
return self.buffer[start:start + self.numel].view(self.shape)
|
|
|
|
|
|
def fingerprint(tensors):
|
|
"""Each tensor's bytes summed as integers: whether run() wrote the weights."""
|
|
out = []
|
|
for t in tensors:
|
|
raw = t.detach().contiguous().view(-1).view(torch.uint8)
|
|
try:
|
|
raw = raw.view(torch.int32)
|
|
except RuntimeError:
|
|
pass
|
|
out.append(int(raw.sum(dtype=torch.int64)))
|
|
return out
|
|
|
|
|
|
def judge(out, ref, allowed):
|
|
"""A video against the reference, or a later call's video against the checked one's: {"passed", "nonfinite",
|
|
"outside" (elements beyond [-1, 1]), "max_abs", "rms", "psnr_db", "max_allowed", "rms_allowed"}, d = out - ref where
|
|
both are finite (rms and PSNR over those elements, None where there is none), in fp32 on out's GPU (the current one
|
|
for a video on the CPU), in chunks of frames, the squares summed in fp64."""
|
|
device = out.device if out.is_cuda else torch.device("cuda", torch.cuda.current_device())
|
|
step = max(1, CHUNK_ELEMENTS // max(1, out[:, :, :1].numel()))
|
|
nonfinite = outside = compared = unmatched = 0
|
|
max_abs = squares = 0.0
|
|
for f in range(0, out.shape[2], step):
|
|
x = out[:, :, f:f + step].to(device, torch.float32)
|
|
y = ref[:, :, f:f + step].to(device, torch.float32)
|
|
finite = torch.isfinite(x)
|
|
nonfinite += int((~finite).sum())
|
|
outside += int((x.abs() > 1).sum())
|
|
# Where the checked call itself left an element unfinite, a later call is not compared there (the check said).
|
|
both = finite & torch.isfinite(y)
|
|
compared += int(both.sum())
|
|
unmatched += int((~finite & torch.isfinite(y)).sum())
|
|
d = torch.where(both, x - y, 0.0)
|
|
max_abs = max(max_abs, float(d.abs().max()))
|
|
squares += float(d.square().sum(dtype=torch.float64))
|
|
del x, y, finite, both, d
|
|
mse = squares / compared if compared else None
|
|
rms = math.sqrt(mse) if compared else None
|
|
passed = (nonfinite == 0 and outside == 0 and compared == out.numel() and max_abs <= allowed["max_abs"]
|
|
and rms <= allowed["rms"])
|
|
return {"passed": passed, "nonfinite": nonfinite, "unmatched": unmatched, "outside": outside,
|
|
"max_abs": max_abs if compared else None, "rms": rms, "psnr_db": 10 * math.log10(4 / mse) if mse else None,
|
|
"max_allowed": allowed["max_abs"], "rms_allowed": allowed["rms"]}
|
|
|
|
|
|
def verdict(c):
|
|
"""A check (judge's result) in a few words."""
|
|
parts = []
|
|
if c["nonfinite"]:
|
|
parts.append(f"{c['nonfinite']} elements not finite (NaN: never written, or overflowed)")
|
|
if c["outside"]:
|
|
parts.append(f"{c['outside']} elements outside [-1, 1]")
|
|
if c["max_abs"] is None:
|
|
parts.append("no element to compare")
|
|
else:
|
|
parts.append(f"max |d| {c['max_abs']:.4g} of {c['max_allowed']:.4g} (x{c['max_abs'] / c['max_allowed']:.2f})")
|
|
parts.append(f"rms {c['rms']:.4g} of {c['rms_allowed']:.4g} (x{c['rms'] / c['rms_allowed']:.2f})")
|
|
parts.append(f"PSNR {c['psnr_db']:.1f} dB" if c["psnr_db"] is not None else "identical")
|
|
return f"{'ok ' if c['passed'] else 'FAIL'} ({', '.join(parts)})"
|
|
|
|
|
|
def run_workload(w, sol, weights, marks, rank, world, pg, harness, ranks, cfg):
|
|
"""One workload on one rank, in step with the other ranks (see the module docstring); rank 0's result."""
|
|
entry = spec.WORKLOADS[w]
|
|
clip, device = entry["clip"], torch.device("cuda", rank)
|
|
timed = entry["timed"] and cfg["time"]
|
|
audited = entry["timed"] and cfg["audit"] and world == 1
|
|
calls = 1 + audited + (TRIALS if timed else 0)
|
|
seed = list(spec.WORKLOADS).index(w)
|
|
z0 = spec.make_inputs(w, device)[0]
|
|
zs = Pool(z0.shape, z0.dtype, device, calls, seed)
|
|
outs = Pool(spec.output_shape(w), torch.bfloat16, device, calls, seed + 1) if rank == 0 else None
|
|
allowed = spec.allowance(clip)
|
|
res = {"clip": clip, "world": world, "timed": entry["timed"]}
|
|
problems = []
|
|
threads = threading.active_count()
|
|
|
|
def arguments():
|
|
z = zs.next()
|
|
z.copy_(z0)
|
|
out = outs.next().fill_(math.nan) if outs else torch.empty(0, dtype=torch.bfloat16, device=device)
|
|
return z, out
|
|
|
|
def unchanged(z):
|
|
if not torch.equal(z, z0) and "run() wrote its input z" not in problems:
|
|
problems.append("run() wrote its input z")
|
|
|
|
def call():
|
|
z, out = arguments()
|
|
torch.cuda.synchronize(device)
|
|
dist.barrier(group=harness)
|
|
start = time.monotonic()
|
|
sol.fn(z, weights, rank, world, pg, out)
|
|
torch.cuda.synchronize(device)
|
|
dist.barrier(group=harness)
|
|
stop = time.monotonic()
|
|
unchanged(z)
|
|
return out, start, stop
|
|
|
|
# The warm-up, which is the call checked.
|
|
out, start, stop = call()
|
|
warm = stop - start
|
|
first, worst = None, 0.0
|
|
if rank == 0:
|
|
if cfg["save"]:
|
|
path = os.path.join(cfg["save"], f"{w}.pt")
|
|
torch.save(out.cpu(), path)
|
|
res["saved"] = path
|
|
elif cfg["check"]:
|
|
res["check"] = judge(out, spec.load_reference(clip), allowed)
|
|
if calls > 1:
|
|
first = out.clone()
|
|
|
|
def agree(out, what):
|
|
nonlocal worst
|
|
# It agrees when it loses no element the checked call wrote, and is within the allowance of it elsewhere;
|
|
# the checked call's own faults are its check's to report.
|
|
c = judge(out, first, allowed)
|
|
worst = max(worst, c["max_abs"] or 0.0)
|
|
close = c["max_abs"] is None or (c["max_abs"] <= allowed["max_abs"] and c["rms"] <= allowed["rms"])
|
|
if c["unmatched"] or not close:
|
|
lost = f"{c['unmatched']} elements not finite where it was; " if c["unmatched"] else ""
|
|
problems.append(f"{what}'s output does not agree with the checked call's: {lost}"
|
|
f"{verdict(c).removeprefix('FAIL').removeprefix('ok').strip()}")
|
|
|
|
# World 1 alone: one call under CUPTI, which attributes every kernel.
|
|
if audited:
|
|
z, out = arguments()
|
|
ops = audit.trace(sol.fn, z, weights, rank, world, pg, out)
|
|
unchanged(z)
|
|
agree(out, "the audited call")
|
|
res |= {"kernels": audit.kernels(ops), "launches": sum(1 for op in ops if op.kind == "kernel"),
|
|
"lib": sum(1 for op in ops if audit.placeholder(op.library)), "lib_share": audit.library_share(ops),
|
|
"comm": sum(1 for op in ops if op.library in audit.COMMUNICATION), "comm_share": audit.comm_share(ops)}
|
|
del ops
|
|
|
|
if timed:
|
|
n = [TRIALS if warm < LONG_CALL_S else LONG_TRIALS]
|
|
dist.broadcast_object_list(n, src=0, group=harness)
|
|
questions, trials, windows = [], [], []
|
|
for k in range(n[0]):
|
|
if rank == 0:
|
|
questions.append(gpu_holders())
|
|
out, start, stop = call()
|
|
trials.append(1e3 * (stop - start))
|
|
windows.append((start, stop))
|
|
if rank == 0:
|
|
agree(out, f"trial {k + 1}")
|
|
if rank == 0:
|
|
questions.append(gpu_holders())
|
|
res |= {"trials": trials, "ms": statistics.mean(trials), "min_ms": min(trials), "windows": windows}
|
|
held = others(questions, [uuid for _, uuid in ranks], {pid for pid, _ in ranks} | {os.getppid()})
|
|
if held:
|
|
res["shared"] = held
|
|
if first is not None:
|
|
res["agree_max_abs"] = worst
|
|
if fingerprint(weights) != marks:
|
|
problems.append("run() wrote its weights (what prepare returned is read-only to run: keep scratch elsewhere)")
|
|
if threading.active_count() > threads:
|
|
problems.append(f"the solution started threads ({threading.active_count() - threads} more)")
|
|
|
|
# What every rank found, so that rank 0 reports it.
|
|
found = [None] * world
|
|
dist.all_gather_object(found, problems, group=harness)
|
|
text = [p if world == 1 else f"rank {r}: {p}" for r, ps in enumerate(found) for p in ps]
|
|
if text and rank == 0:
|
|
res["error"] = "; ".join(text)
|
|
return res
|
|
|
|
|
|
def worker(rank, world, port, cfg, conn):
|
|
"""The process of one rank (see A run in the module docstring): it reports to the process that spawned it through
|
|
conn, ("ready", rank, uuid, prepare seconds), ("start", workload) and ("result", workload, result) on rank 0,
|
|
("done", rank), or ("error", rank, workload or None, phase, traceback), after which it exits at once."""
|
|
phase, current = "start", None
|
|
try:
|
|
torch.cuda.set_device(rank)
|
|
device = torch.device("cuda", rank)
|
|
timeout = timedelta(seconds=WORKLOAD_SEC)
|
|
dist.init_process_group("nccl", init_method=f"tcp://127.0.0.1:{port}", rank=rank, world_size=world,
|
|
device_id=device, timeout=timeout)
|
|
harness = dist.new_group(backend="gloo", timeout=timeout)
|
|
pg = dist.group.WORLD
|
|
phase = "load"
|
|
sol = loader.import_built(cfg["workdir"], cfg["spec"])
|
|
phase = "weights"
|
|
weights = spec.load_weights(device)
|
|
phase = "prepare"
|
|
prepare_s = None
|
|
if sol.prepare is not None:
|
|
torch.cuda.synchronize(device)
|
|
start = time.monotonic()
|
|
got = sol.prepare(weights, rank, world, pg)
|
|
torch.cuda.synchronize(device)
|
|
prepare_s = time.monotonic() - start
|
|
if not isinstance(got, (list, tuple)) or not all(isinstance(t, torch.Tensor) for t in got):
|
|
raise TypeError(f"prepare returned {type(got).__name__}, not a list of tensors")
|
|
weights = list(got)
|
|
uuid = str(torch.cuda.get_device_properties(device).uuid).lower().removeprefix("gpu-")
|
|
ranks = [None] * world
|
|
dist.all_gather_object(ranks, (os.getpid(), uuid), group=harness)
|
|
conn.send(("ready", rank, uuid, prepare_s))
|
|
phase = "workload"
|
|
marks = fingerprint(weights)
|
|
for w in cfg["workloads"]:
|
|
current = w
|
|
if rank == 0:
|
|
conn.send(("start", w))
|
|
res = run_workload(w, sol, weights, marks, rank, world, pg, harness, ranks, cfg)
|
|
if rank == 0:
|
|
conn.send(("result", w, res))
|
|
current = None
|
|
torch.cuda.empty_cache()
|
|
phase = "end"
|
|
dist.barrier(group=harness)
|
|
conn.send(("done", rank))
|
|
dist.destroy_process_group()
|
|
except BaseException:
|
|
try:
|
|
conn.send(("error", rank, current, phase, traceback.format_exc(limit=8)))
|
|
finally:
|
|
sys.stdout.flush()
|
|
sys.stderr.flush()
|
|
# At once: NCCL's teardown would wait for the other ranks, which may wait for this one.
|
|
os._exit(1)
|
|
|
|
|
|
def free_port():
|
|
with socket.socket() as s:
|
|
s.bind(("127.0.0.1", 0))
|
|
return s.getsockname()[1]
|
|
|
|
|
|
def stop(procs, grace=0.0):
|
|
"""Ends rank processes: those that do not exit within `grace` seconds get SIGTERM, then SIGKILL after
|
|
STOP_GRACE_SEC."""
|
|
deadline = time.monotonic() + grace
|
|
for p in procs:
|
|
p.join(max(0.0, deadline - time.monotonic()))
|
|
for p in procs:
|
|
if p.is_alive():
|
|
p.terminate()
|
|
deadline = time.monotonic() + STOP_GRACE_SEC
|
|
for p in procs:
|
|
p.join(max(0.1, deadline - time.monotonic()))
|
|
if p.is_alive():
|
|
p.kill()
|
|
p.join(10)
|
|
|
|
|
|
def exit_text(code):
|
|
if code is None:
|
|
return "did not exit"
|
|
if code < 0:
|
|
import signal
|
|
|
|
try:
|
|
return f"was killed by {signal.Signals(-code).name}"
|
|
except ValueError:
|
|
return f"was killed by signal {-code}"
|
|
return f"exited with code {code}"
|
|
|
|
|
|
def session(world, workloads, cfg, report, show):
|
|
"""One set of rank processes running these workloads at one world size, in order; results go into report as they
|
|
arrive, and show(w) prints each. Returns the workloads it did not reach: after a failure at one workload (which gets
|
|
the error), those after it; [] when every one ran, or when the failure came before any (every one gets it)."""
|
|
ctx = torch.multiprocessing.get_context("spawn")
|
|
port = free_port()
|
|
pipes = [ctx.Pipe(duplex=False) for _ in range(world)]
|
|
procs = [ctx.Process(target=worker, args=(r, world, port, dict(cfg, workloads=workloads), pipes[r][1]),
|
|
daemon=True) for r in range(world)]
|
|
for p in procs:
|
|
p.start()
|
|
for _, send in pipes:
|
|
send.close()
|
|
receivers = {pipes[r][0]: r for r in range(world)}
|
|
done, errors, finished, uuids = set(), [], [], {}
|
|
current, clock, failure = None, None, None
|
|
deadline = time.monotonic() + SESSION_START_SEC
|
|
started = time.monotonic()
|
|
|
|
def handle(msg):
|
|
nonlocal current, clock, deadline
|
|
kind = msg[0]
|
|
if kind == "ready":
|
|
_, rank, uuid, prepare_s = msg
|
|
uuids[rank] = uuid
|
|
if rank == 0:
|
|
report.setdefault("prepare_s", {})[str(world)] = prepare_s
|
|
if len(uuids) == world:
|
|
clock = Clock([uuids[r] for r in range(world)])
|
|
ready = f", prepare {report['prepare_s'][str(world)]:.1f} s" if prepare_s is not None else ""
|
|
print(f"world {world}: {world} rank{'s' if world > 1 else ''} up in {time.monotonic() - started:.0f} s"
|
|
f"{ready}", flush=True)
|
|
elif kind == "start":
|
|
current = msg[1]
|
|
deadline = time.monotonic() + WORKLOAD_SEC
|
|
elif kind == "result":
|
|
_, w, res = msg
|
|
windows = res.pop("windows", None)
|
|
means = clock.means(windows) if clock and windows else []
|
|
if means and means[0]:
|
|
res["clock_mhz"] = means[0]
|
|
if world > 1 and means and all(means):
|
|
res["world_clock_mhz"] = [min(means), max(means)]
|
|
report["workloads"][w] = res
|
|
finished.append(w)
|
|
current = None
|
|
deadline = time.monotonic() + WORKLOAD_SEC
|
|
show(w)
|
|
elif kind == "error":
|
|
errors.append(msg[1:])
|
|
elif kind == "done":
|
|
done.add(msg[1])
|
|
|
|
try:
|
|
while len(done) < world:
|
|
# A rank that exited cleanly may still have its last messages in its pipe: it fails once that is read.
|
|
bad = [(r, p.exitcode) for r, p in enumerate(procs) if r not in done and p.exitcode is not None
|
|
and (p.exitcode != 0 or pipes[r][0] not in receivers)]
|
|
if bad:
|
|
failure = "; ".join(f"rank {r} {exit_text(code)}" for r, code in bad)
|
|
break
|
|
left = deadline - time.monotonic()
|
|
if left <= 0:
|
|
failure = (f"timed out after {WORKLOAD_SEC} s" if current else
|
|
f"the ranks did not start within {SESSION_START_SEC} s")
|
|
break
|
|
waits = list(receivers) + [p.sentinel for p in procs if p.exitcode is None]
|
|
for obj in connection.wait(waits, timeout=min(left, 5.0)):
|
|
if obj in receivers:
|
|
try:
|
|
handle(obj.recv())
|
|
except (EOFError, OSError):
|
|
del receivers[obj]
|
|
# What the other ranks still say, such as the error that made one exit.
|
|
drain = time.monotonic() + 2.0
|
|
while failure and receivers and time.monotonic() < drain:
|
|
for obj in connection.wait(list(receivers), timeout=max(0.0, drain - time.monotonic())):
|
|
try:
|
|
handle(obj.recv())
|
|
except (EOFError, OSError):
|
|
del receivers[obj]
|
|
finally:
|
|
stop(procs, 0.0 if failure else STOP_GRACE_SEC)
|
|
if clock:
|
|
clock.stop()
|
|
if not failure:
|
|
return []
|
|
|
|
# The first error a rank reported says what failed; else how a rank ended.
|
|
if errors:
|
|
rank, at, phase, text = errors[0]
|
|
if phase == "load":
|
|
report["load_error"] = text.strip().splitlines()[-1]
|
|
text = (f"rank {rank}: " if world > 1 else "") + text.strip()
|
|
else:
|
|
text = failure
|
|
rest = [w for w in workloads if w not in finished]
|
|
at = current or (errors and errors[0][1])
|
|
if at not in rest:
|
|
at = rest[0] if rest and finished else None
|
|
if not at or report.get("load_error"):
|
|
for w in rest:
|
|
report["workloads"][w] = {"error": text, "world": world, "clip": spec.WORKLOADS[w]["clip"],
|
|
"timed": spec.WORKLOADS[w]["timed"]}
|
|
show(w)
|
|
return []
|
|
report["workloads"][at] = {"error": text, "world": world, "clip": spec.WORKLOADS[at]["clip"],
|
|
"timed": spec.WORKLOADS[at]["timed"]}
|
|
show(at)
|
|
return rest[rest.index(at) + 1:]
|
|
|
|
|
|
def kernel_lines(res, limit=25):
|
|
"""The audited call's kernels by GPU time, one line each."""
|
|
rows = res.get("kernels", [])
|
|
total = sum(r[3] for r in rows) or 1
|
|
lines = [f" {'[' + lib + ']' if lib else '':<9} {ns / 1e6:10.2f} ms {100 * ns / total:5.1f}% x{n:<6} "
|
|
f"{name if name.startswith('(') else audit.short_name(name, 110)}" for lib, name, n, ns in rows[:limit]]
|
|
if len(rows) > limit:
|
|
rest = rows[limit:]
|
|
lines.append(f" {'':<9} {sum(r[3] for r in rest) / 1e6:10.2f} ms "
|
|
f"{100 * sum(r[3] for r in rest) / total:5.1f}% {len(rest)} more kinds")
|
|
return lines
|
|
|
|
|
|
def main():
|
|
p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
|
p.add_argument("--solution", default="solution", help="the solution directory (default: solution)")
|
|
p.add_argument("--workload", nargs="*", help="workloads to run (default: every one this job has the GPUs for)")
|
|
p.add_argument("--no-time", action="store_true", help="correctness only: no audit, no trials")
|
|
p.add_argument("--no-check", action="store_true", help="do not check the outputs (CI's timing runs)")
|
|
p.add_argument("--no-audit", action="store_true", help="no call under CUPTI")
|
|
p.add_argument("--kernels", action="store_true", help="list the audited calls' kernels by GPU time")
|
|
p.add_argument("--build-dir", help="write and build the solution here; reused when it holds the same build")
|
|
p.add_argument("--no-cache", action="store_true", help="build fresh, without the node's cache of builds")
|
|
p.add_argument("--save", help="CI: save each workload's checked output here instead of checking it")
|
|
p.add_argument("--json", help="write the results to this file")
|
|
a = p.parse_args()
|
|
unknown = [w for w in a.workload or [] if w not in spec.WORKLOADS]
|
|
if unknown:
|
|
p.error(f"unknown workloads {unknown}; known: {list(spec.WORKLOADS)}")
|
|
# Absolute paths, the data's too: the ranks run in the solution's working directory.
|
|
solution = os.path.abspath(a.solution)
|
|
if os.path.isfile(solution):
|
|
solution = os.path.dirname(solution)
|
|
a.build_dir, a.save, a.json = (x and os.path.abspath(x) for x in (a.build_dir, a.save, a.json))
|
|
os.environ["DATA"] = os.path.abspath(spec.data_dir())
|
|
report = {"workloads": {}}
|
|
|
|
def finish():
|
|
if a.json:
|
|
with open(a.json, "w") as f:
|
|
json.dump(report, f)
|
|
|
|
rules = audit.rule_table(getattr(spec, "RULES", None))
|
|
report["rules"] = audit.check(solution, rules)
|
|
for v in report["rules"]:
|
|
print(f"rule violated: {v}")
|
|
print(f"stack rules: {'ok' if not report['rules'] else 'FAIL'} ({len(rules)} rules)")
|
|
|
|
# The workloads: those named, or every one this job has the GPUs for.
|
|
gpus = torch.cuda.device_count()
|
|
chosen = a.workload or [w for w, e in spec.WORKLOADS.items() if e["world"] <= gpus]
|
|
skipped = [w for w in spec.WORKLOADS if w not in chosen and not a.workload]
|
|
if skipped:
|
|
print(f"skipped {', '.join(skipped)}: this job has {gpus} GPU{'s' if gpus != 1 else ''}")
|
|
order = [w for w in spec.WORKLOADS if w in chosen]
|
|
timed = [w for w in order if spec.WORKLOADS[w]["timed"] and not a.no_time]
|
|
limits, dev, peak = bounds.workload_bounds(spec, timed)
|
|
if limits:
|
|
print(bounds.describe(dev, peak))
|
|
|
|
try:
|
|
workdir, build_spec, seconds = loader.stage(solution, a.build_dir,
|
|
cache=None if a.no_cache else loader.cache_dir())
|
|
except Exception as e:
|
|
report["load_error"] = str(e) if isinstance(e, loader.SolutionError) else traceback.format_exc(limit=6)
|
|
print(f"FAILED to build {solution}:\n{report['load_error']}", flush=True)
|
|
finish()
|
|
sys.exit(1)
|
|
report |= {"language": build_spec["languages"][0], "build_s": seconds}
|
|
built = f"built in {seconds:.1f} s" if seconds is not None else "a build reused"
|
|
print(f"{build_spec['entry_point']} ({', '.join(build_spec['languages'])}): {built}", flush=True)
|
|
|
|
width = max(len(w) for w in spec.WORKLOADS)
|
|
|
|
def show(w):
|
|
res = report["workloads"][w]
|
|
line = f"{w:<{width}}"
|
|
if "ms" in res:
|
|
at = f", SM clock {res['clock_mhz']:.0f} MHz" if "clock_mhz" in res else ""
|
|
low, high = res.get("world_clock_mhz", (None, None))
|
|
if low and "clock_mhz" in res and high - low >= 0.02 * res["clock_mhz"]:
|
|
at += f", the world's GPUs {low:.0f}-{high:.0f} MHz"
|
|
line += f" {res['ms']:9.1f} ms (fastest trial {res['min_ms']:.1f}{at})"
|
|
if w in limits:
|
|
res |= {"bound_ms": limits[w], "gap": res["ms"] / limits[w]}
|
|
line += f" bound {limits[w]:.1f} ms, gap x{res['gap']:.2f}"
|
|
if "check" in res:
|
|
line += " " + verdict(res["check"])
|
|
elif "saved" in res:
|
|
line += " saved"
|
|
if not res.get("timed", True):
|
|
line += " (correctness only)"
|
|
if "lib" in res:
|
|
line += f" library kernels {res['lib']} of {res['launches']} ({100 * res['lib_share']:.0f}% of GPU time)"
|
|
if res.get("comm"):
|
|
line += f", NCCL {res['comm']} ({100 * res['comm_share']:.0f}%)"
|
|
if "shared" in res:
|
|
line += f" SHARED: {res['shared']} other process(es) held a GPU while it was timed"
|
|
if "error" in res:
|
|
line += " ERROR\n" + res["error"]
|
|
print(line, flush=True)
|
|
if a.kernels and "kernels" in res:
|
|
print("\n".join(kernel_lines(res)), flush=True)
|
|
finish()
|
|
|
|
cfg = {"workdir": workdir, "spec": build_spec, "time": not a.no_time, "check": not a.no_check,
|
|
"audit": not a.no_audit and (not a.no_time or a.kernels), "save": a.save}
|
|
for world in sorted({spec.WORKLOADS[w]["world"] for w in order}):
|
|
todo = [w for w in order if spec.WORKLOADS[w]["world"] == world]
|
|
if world > gpus or report.get("load_error"):
|
|
why = report.get("load_error") or f"needs {world} GPUs; this job has {gpus}"
|
|
for w in todo:
|
|
report["workloads"][w] = {"error": why, "world": world, "clip": spec.WORKLOADS[w]["clip"],
|
|
"timed": spec.WORKLOADS[w]["timed"]}
|
|
show(w)
|
|
continue
|
|
while todo:
|
|
todo = session(world, todo, cfg, report, show)
|
|
|
|
results = [report["workloads"][w] for w in order if w in report["workloads"]]
|
|
times = [report["workloads"].get(w, {}).get("ms") for w in timed]
|
|
if timed and all(times):
|
|
gaps = [report["workloads"][w]["gap"] for w in timed if "gap" in report["workloads"][w]]
|
|
gap = f", gap to the bound x{bounds.geomean(gaps):.2f}" if len(gaps) == len(timed) else ""
|
|
clocks = [report["workloads"][w].get("clock_mhz") for w in timed]
|
|
at = f", mean SM clock {statistics.mean(clocks):.0f} MHz" if all(clocks) else ""
|
|
print(f"geomean {bounds.geomean(times):.1f} ms over {len(timed)} timed workloads (mean of trials){gap}{at}")
|
|
shared = [w for w in order if "shared" in report["workloads"].get(w, {})]
|
|
if shared:
|
|
print(f"SHARED GPU: another process held one while {', '.join(shared)} "
|
|
f"{'was' if len(shared) == 1 else 'were'} timed, so those times do not compare with another job's")
|
|
checked = [r for r in results if "check" in r]
|
|
if checked:
|
|
ok = sum(1 for r in checked if r["check"]["passed"] and "error" not in r)
|
|
print(f"correctness: ok at {ok} of {len(checked)} workloads checked")
|
|
failed = [w for w in order if "error" in report["workloads"].get(w, {})]
|
|
if failed:
|
|
print(f"errors at {', '.join(failed)}")
|
|
audited = [r for r in results if "kernels" in r]
|
|
if audited:
|
|
libs = collections.Counter()
|
|
for r in audited:
|
|
for lib, _, n, _ in r["kernels"]:
|
|
if audit.placeholder(lib):
|
|
libs[lib] += n
|
|
detail = ", ".join(f"{k} {v}" for k, v in libs.most_common())
|
|
share = statistics.mean(r["lib_share"] for r in audited)
|
|
print(f"library kernels: {sum(libs.values())} launches in {len(audited)} audited calls"
|
|
f"{' (' + detail + ')' if detail else ''}, {100 * share:.0f}% of GPU time on average")
|
|
finish()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|