diff --git a/benchmarks/README.md b/benchmarks/README.md index 0f3331e..656ea9d 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -9,13 +9,21 @@ CNN: `build_model` in `cifar10_breakdown.py`, the same model as | `cifar10_breakdown.py` | ms/step, split by training-step phase, on the normal GPU path | `benchmark_release` | | `cifar10_paths.py` | ms/step and loss over one epoch, on NumPy, fallback CuPy and fused kernels | `benchmark_throughput`, `benchmark_convergence` | | `cifar10_vram.py` | live VRAM over two steps, on fallback CuPy and fused kernels | `benchmark_vram` | +| `cifar10_torch.py` | ms/step of the same CNN in PyTorch, compiled with fp16 autocast | `benchmark_release` | | `compare.py` | prints what changed between two `cifar10_breakdown.py` results | none | +| `common.py` | the constants and helpers the scripts share; imports neither CuPy nor PyTorch | none | | `plot.py` | draws the figures from saved results | all of the above | The scripts need a GPU and CuPy. They also need CIFAR-10 in `data/cifar-10`: create it with `examples/cifar10/cifar-setup.ipynb`. `plot.py` needs matplotlib instead, which isn't a dependency of the package. +`cifar10_torch.py` needs PyTorch, also not a dependency: install the build for your GPU from +pytorch.org (on a ROCm GPU, the wheel index for your ROCm version). The result records the +PyTorch version. PyTorch and CuPy can't be loaded into one process (`import cupy, torch` +aborts with an LLVM error), so the torch script runs on its own and the figure joins the two +result files. + ## Before and after a change python3 benchmarks/cifar10_breakdown.py benchmarks/results/before.json @@ -39,6 +47,13 @@ Run from the repo root on each GPU, naming the file after that GPU: PYTHONPATH=.worktrees/old python3 benchmarks/cifar10_breakdown.py benchmarks/history/v1.1.0/breakdown_rtx3050.json PYTHONPATH=.worktrees/new python3 benchmarks/cifar10_breakdown.py benchmarks/history/v1.2.0/breakdown_rtx3050.json +Then PyTorch's bar for the same GPU, from `torch_.json` in the new version's folder. It +is the same CNN in PyTorch with the whole step under `torch.compile`, fp16 autocast and no +loss scaling, so it is a reference point and not a tuned ceiling. Without the file, the figure +draws that GPU's panel without the bar: + + python3 benchmarks/cifar10_torch.py benchmarks/history/v1.2.0/torch_rtx3050.json + On a GPU other than the one the README names, also add its fused path to the convergence figure: @@ -56,6 +71,8 @@ Draw them, then remove the worktrees: git worktree remove .worktrees/new `cifar10_paths.py` runs NumPy once and each GPU path three times, which takes about 15 -minutes. The cross-version figure uses only `cifar10_breakdown.py`, so older versions need nothing -else from this folder. Each GPU gets its own panel and scale on it: comparing versions is +minutes. The cross-version figure uses only `cifar10_breakdown.py` (and `cifar10_torch.py` for the +PyTorch bar), so older versions need nothing else from this folder. The PyTorch bar is +wall-clock ms/step, while the other bars are GPU time per stage, which runs a few percent above +their wall-clock total. Each GPU gets its own panel and scale on it: comparing versions is only a fair A/B on one machine. diff --git a/benchmarks/cifar10_breakdown.py b/benchmarks/cifar10_breakdown.py index c46d442..15828f6 100644 --- a/benchmarks/cifar10_breakdown.py +++ b/benchmarks/cifar10_breakdown.py @@ -26,41 +26,27 @@ it does outside the wrapped methods grows "unattributed". """ import argparse -import datetime import json -import platform -import statistics -import subprocess -import sys import time from pathlib import Path import cupy as cp -import numpy as np +from common import ( + BASELINE_STEPS, + BATCH_SIZE, + COMPUTE_DTYPE, + WARMUP_STEPS, + base_meta, + load_train, + spread, +) import aether as ae import aether.config as config -TRAIN_DIR = Path("data") / "cifar-10" / "cifar-10_train.npz" - -BATCH_SIZE = 256 -COMPUTE_DTYPE = "float16" -WARMUP_STEPS = 20 -BASELINE_STEPS = 128 BREAKDOWN_STEPS = 100 -def load_train(): - """CIFAR-10 training images and labels, as host arrays.""" - if not TRAIN_DIR.exists(): - sys.exit( - f"[aether] {TRAIN_DIR} not found. Run from the repo root after creating it " - "with examples/cifar10/cifar-setup.ipynb." - ) - with np.load(TRAIN_DIR, allow_pickle=False) as data: - return data["X_train"], data["y_train"] - - def build_model(X_train, device="cupy"): """Same model, pipeline and precision as examples/cifar10/v1_cifar10.py.""" config.set_backend(device) @@ -227,38 +213,12 @@ def timed(*args, **kwargs): return result -def spread(samples): - return {"median": statistics.median(samples), "min": min(samples), "max": max(samples)} - - -def git_state(): - """(commit, dirty) of the checkout aether was imported from, or (None, None) - outside git, so a run against a worktree of an older tag records that tag.""" - source = Path(ae.__file__).parent - try: - commit = subprocess.run(["git", "rev-parse", "HEAD"], capture_output=True, - text=True, check=True, cwd=source).stdout.strip() - status = subprocess.run(["git", "status", "--porcelain"], capture_output=True, - text=True, check=True, cwd=source).stdout - except (OSError, subprocess.CalledProcessError): - return None, None - return commit, bool(status.strip()) - - def run_meta(): """What a result file needs to say about where it was measured.""" - commit, dirty = git_state() return { - "commit": commit, - "dirty": dirty, - "date": datetime.datetime.now().isoformat(timespec="seconds"), + **base_meta(Path(ae.__file__).parent), "gpu": cp.cuda.runtime.getDeviceProperties(0)["name"].decode(), "cupy": cp.__version__, - "python": platform.python_version(), - "batch_size": BATCH_SIZE, - "compute_dtype": COMPUTE_DTYPE, - "warmup_steps": WARMUP_STEPS, - "baseline_steps": BASELINE_STEPS, } diff --git a/benchmarks/cifar10_torch.py b/benchmarks/cifar10_torch.py new file mode 100644 index 0000000..1845d65 --- /dev/null +++ b/benchmarks/cifar10_torch.py @@ -0,0 +1,143 @@ +""" +Train the v1 CIFAR-10 CNN in PyTorch, as the reference bar on the release figure. + +Run from the repo root, in its own process: PyTorch and CuPy can't share one (see +benchmarks/README.md). PyTorch isn't a dependency of the package; install the build for your +GPU from pytorch.org. + + python3 benchmarks/cifar10_torch.py benchmarks/history/v1.2.0/torch_rx9070xt.json + +The model is build_model in cifar10_breakdown.py, written the way a PyTorch user would, with +the whole step under torch.compile: default NCHW layout, fp16 autocast and no loss scaling. +Autocast also keeps the activations after each Conv2d in fp16, where Aether keeps them fp32. +Every step does what Aether's does: standardize the batch on the GPU, forward, label-smoothed +loss plus the L2 term on the third Conv2d, accuracy, backward and an AdamW step. Aether's +learning-rate decay is left out, because it changes the loss and not the time per step. + +The timing protocol is cifar10_breakdown.py's baseline: warm up, discard one full-length run, +then time BASELINE_STEPS per repeat with no per-step sync. Compilation happens in the warm-up, +so it counts toward warmup_s and not baseline_ms. +""" +import argparse +import json +import sys +import time +from pathlib import Path + +from common import ( + BASELINE_STEPS, + BATCH_SIZE, + WARMUP_STEPS, + base_meta, + load_train, + spread, +) + +try: + import torch + import torch.nn.functional as F + from torch import nn +except ImportError: + sys.exit("[aether] PyTorch is not installed; install the build for your GPU from pytorch.org.") + +SEED = 42 # build_model's manual_seed +LABEL_SMOOTHING = 0.05 +L2 = 1e-5 # on the third Conv2d only +PARAMETERS = 141_354 # 3 Conv2d, 5 BatchNorm (gamma, beta) and the Dense, counted by hand + + +def build_model(): + def conv_block(in_channels, out_channels): + return [nn.Conv2d(in_channels, out_channels, 3, padding="same"), + nn.BatchNorm2d(out_channels, eps=1e-5, momentum=0.1), # Aether's momentum is 0.9 + nn.ReLU()] + + model = nn.Sequential( + *conv_block(3, 32), *conv_block(32, 32), nn.MaxPool2d(2, 2), nn.Dropout2d(0.1), + *conv_block(32, 64), *conv_block(64, 64), nn.MaxPool2d(2, 2), nn.Dropout2d(0.15), + *conv_block(64, 128), nn.MaxPool2d(2, 2), nn.Dropout2d(0.2), + nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(128, 10), + ).cuda() + parameters = sum(p.numel() for p in model.parameters()) + if parameters != PARAMETERS: + sys.exit(f"[aether] The PyTorch model has {parameters} parameters, not {PARAMETERS}; " + "it no longer matches build_model in cifar10_breakdown.py.") + return model + + +def make_step(X): + """A compiled training step on the GPU.""" + torch.manual_seed(SEED) + model = build_model() + l2_weight = model[8].weight # the third Conv2d + optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.01) + scaled = X.float() / 255.0 + mean, std = scaled.mean(), scaled.std(correction=0) + del scaled + correct = torch.zeros((), device=X.device, dtype=torch.int64) + + def step(x, y): + x = ((x.float() * (1.0 / 255.0) - mean) / std).permute(0, 3, 1, 2).contiguous() + with torch.autocast("cuda", dtype=torch.float16): + logits = model(x) + logits = logits.float() + loss = (F.cross_entropy(logits, y, label_smoothing=LABEL_SMOOTHING) + + L2 * l2_weight.square().sum()) + correct.add_((logits.argmax(1) == y).sum()) + optimizer.zero_grad(set_to_none=True) + loss.backward() + optimizer.step() + return loss + + return torch.compile(step) + + +def timed_train(step, X, y, num_steps): + """Wall time of num_steps shuffled full batches.""" + n = num_steps * BATCH_SIZE + torch.cuda.synchronize() + start = time.perf_counter() + order = torch.randperm(n, device=X.device) + for i in range(num_steps): + batch = order[i * BATCH_SIZE:(i + 1) * BATCH_SIZE] + step(X[batch], y[batch]) + torch.cuda.synchronize() + return time.perf_counter() - start + + +def main(): + parser = argparse.ArgumentParser(description=__doc__.splitlines()[1]) + parser.add_argument("output", type=Path, help="results file, for benchmarks/plot.py") + parser.add_argument("--repeats", type=int, default=3, + help="baseline runs after warm-up (default: 3)") + args = parser.parse_args() + if args.repeats < 1: + parser.error("--repeats must be at least 1") + if not torch.cuda.is_available(): + sys.exit("[aether] PyTorch can't see a GPU.") + + X, y = (torch.as_tensor(array).cuda() for array in load_train()) + step = make_step(X) + + warmup_s = timed_train(step, X, y, WARMUP_STEPS) + timed_train(step, X, y, BASELINE_STEPS) # the first full-length run is slow + baseline_runs = [timed_train(step, X, y, BASELINE_STEPS) * 1e3 / BASELINE_STEPS + for _ in range(args.repeats)] + + result = { + "meta": { + **base_meta(Path(__file__).resolve().parent), + "gpu": torch.cuda.get_device_name(0), + "torch": torch.__version__, + "mode": "compiled", + "repeats": args.repeats, + "warmup_s": warmup_s, + }, + "baseline_ms": spread(baseline_runs), + } + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(result, indent=2) + "\n") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/common.py b/benchmarks/common.py new file mode 100644 index 0000000..130967b --- /dev/null +++ b/benchmarks/common.py @@ -0,0 +1,63 @@ +""" +What the benchmark scripts share. It imports neither CuPy nor PyTorch, because the two +can't be loaded into one process (see benchmarks/README.md). +""" +import datetime +import platform +import statistics +import subprocess +import sys +from pathlib import Path + +import numpy as np + +TRAIN_DIR = Path("data") / "cifar-10" / "cifar-10_train.npz" + +BATCH_SIZE = 256 +COMPUTE_DTYPE = "float16" +WARMUP_STEPS = 20 +BASELINE_STEPS = 128 + + +def load_train(): + """CIFAR-10 training images and labels, as host arrays.""" + if not TRAIN_DIR.exists(): + sys.exit( + f"[aether] {TRAIN_DIR} not found. Run from the repo root after creating it " + "with examples/cifar10/cifar-setup.ipynb." + ) + with np.load(TRAIN_DIR, allow_pickle=False) as data: + return data["X_train"], data["y_train"] + + +def spread(samples): + return {"median": statistics.median(samples), "min": min(samples), "max": max(samples)} + + +def git_state(source): + """(commit, dirty) of the checkout containing `source`, or (None, None) outside git, + so a run against a worktree of an older tag records that tag.""" + try: + commit = subprocess.run(["git", "rev-parse", "HEAD"], capture_output=True, + text=True, check=True, cwd=source).stdout.strip() + status = subprocess.run(["git", "status", "--porcelain"], capture_output=True, + text=True, check=True, cwd=source).stdout + except (OSError, subprocess.CalledProcessError): + return None, None + return commit, bool(status.strip()) + + +def base_meta(source): + """What a result file says about the run, whatever library ran it. `source` is the + directory of the code that was measured.""" + commit, dirty = git_state(source) + return { + "commit": commit, + "dirty": dirty, + "date": datetime.datetime.now().isoformat(timespec="seconds"), + "python": platform.python_version(), + "batch_size": BATCH_SIZE, + "compute_dtype": COMPUTE_DTYPE, + "warmup_steps": WARMUP_STEPS, + "baseline_steps": BASELINE_STEPS, + } diff --git a/benchmarks/history/v1.2.0/torch_rtx3050.json b/benchmarks/history/v1.2.0/torch_rtx3050.json new file mode 100644 index 0000000..955c788 --- /dev/null +++ b/benchmarks/history/v1.2.0/torch_rtx3050.json @@ -0,0 +1,22 @@ +{ + "meta": { + "commit": "29c71e361d34f53efa34a84bf4551a9ac4d74ff7", + "dirty": true, + "date": "2026-11-01T21:48:18", + "python": "3.14.10", + "batch_size": 256, + "compute_dtype": "float16", + "warmup_steps": 20, + "baseline_steps": 128, + "gpu": "NVIDIA GeForce RTX 3050 6GB Laptop GPU", + "torch": "2.12.0+cu132", + "mode": "compiled", + "repeats": 3, + "warmup_s": 86.26663080000435 + }, + "baseline_ms": { + "median": 16.8658515625566, + "min": 16.82847109384511, + "max": 16.882614843780175 + } +} diff --git a/benchmarks/history/v1.2.0/torch_rx9070xt.json b/benchmarks/history/v1.2.0/torch_rx9070xt.json new file mode 100644 index 0000000..a9543ed --- /dev/null +++ b/benchmarks/history/v1.2.0/torch_rx9070xt.json @@ -0,0 +1,22 @@ +{ + "meta": { + "commit": "fc06f99038bca4e8c77283e548aeb97d323ccd1d", + "dirty": true, + "date": "2026-11-02T02:13:35", + "python": "3.14.4", + "batch_size": 256, + "compute_dtype": "float16", + "warmup_steps": 20, + "baseline_steps": 128, + "gpu": "AMD Radeon RX 9070 XT", + "torch": "2.12.0+rocm7.14.1", + "mode": "compiled", + "repeats": 3, + "warmup_s": 3.1427786790000027 + }, + "baseline_ms": { + "median": 7.6043785781223505, + "min": 7.583219632813609, + "max": 7.621538335932598 + } +} diff --git a/benchmarks/plot.py b/benchmarks/plot.py index 4e4d142..e7293ef 100644 --- a/benchmarks/plot.py +++ b/benchmarks/plot.py @@ -11,7 +11,8 @@ benchmark_release, old against new per GPU from both versions Each GPU's release panel has its own scale and compares only with itself, so every panel -is a true A/B on one machine. On the convergence figure, any paths_fused_.json adds +is a true A/B on one machine. A torch_.json in / adds PyTorch's bar to that +GPU's panel; its bar is the wall-clock ms/step, the one number it has. On the convergence figure, any paths_fused_.json adds that GPU's fused path as a dashed curve; the title and equal-loss line stay on paths.json. """ import argparse @@ -54,6 +55,7 @@ "series": {"cupy-fused": "#2a9d5c", "cupy-fallback": "#2a78d6", "numpy-cpu": "#4a3aa7", "other-gpu-fused": "#5da938"}, "stages": {"forward": "#2a9d5c", "backward": "#16643a"}, + "torch": "#ee4c2c", }, "dark": { "suffix": "-dark", @@ -68,6 +70,7 @@ "series": {"cupy-fused": "#3dbb74", "cupy-fallback": "#3987e5", "numpy-cpu": "#9085e9", "other-gpu-fused": "#81d262"}, "stages": {"forward": "#3dbb74", "backward": "#9fe8bd"}, + "torch": "#f0664a", }, } STAGE_LABELS = {"forward": "Forward", "backward": "Backward", "rest": "Everything else"} @@ -378,40 +381,56 @@ def stage_ms(breakdown): return stages -def release_panel(ax, old_version, new_version, old, new, gain, theme): +def release_panel(ax, old_version, new_version, old, new, reference, gain, theme): # The remainder is small and carries no story, so it stays neutral. colors = {**theme["stages"], "rest": theme["muted"]} + top = 2 if reference else 1 ax.grid(axis="x", color=theme["grid"], linewidth=1.0) ax.spines["left"].set_visible(False) ax.set_title(new["meta"]["gpu"], loc="left", fontsize=SZ_CAPTION, fontweight="bold", color=theme["primary"], pad=5) - for row, breakdown in ((1, old), (0, new)): + for row, breakdown in ((top, old), (top - 1, new)): left = 0.0 for stage, ms in stage_ms(breakdown).items(): ax.barh(row, ms, left=left, height=0.42, color=colors[stage], edgecolor=theme["surface"], linewidth=1.0, label=STAGE_LABELS[stage]) left += ms label = f"{fmt_ms(breakdown['baseline_ms']['median'])} ms/step" - if row == 0: + if row == top - 1: label += f" · {fmt_multiple(gain)} faster" ax.annotate(label, xy=(left, row), xytext=(7, 0), textcoords="offset points", va="center", fontsize=SZ_AXIS, color=theme["primary"], - fontweight="bold" if row == 0 else "normal") - ax.set_yticks([1, 0], [old_version, new_version], color=theme["primary"]) - ax.set_ylim(-0.5, 1.5) - ax.set_xlim(0, sum(stage_ms(old).values()) * 1.2) + fontweight="bold" if row == top - 1 else "normal") + ticks = [(top, old_version), (top - 1, new_version)] + longest = sum(stage_ms(old).values()) + if reference: + ms, new_ms = reference["baseline_ms"]["median"], new["baseline_ms"]["median"] + ratio = ms / new_ms + ax.barh(0, ms, height=0.42, color=theme["torch"]) + ax.annotate(f"{fmt_ms(ms)} ms/step · {new_version} is " + f"{fmt_multiple(ratio if ratio >= 1 else 1 / ratio)} " + f"{'faster' if ratio >= 1 else 'slower'}", + xy=(ms, 0), xytext=(7, 0), textcoords="offset points", va="center", + fontsize=SZ_AXIS, color=theme["primary"]) + ticks.append((0, "PyTorch")) + longest = max(longest, ms) + ax.set_yticks(*zip(*ticks), color=theme["primary"]) + ax.set_ylim(-0.5, top + 0.5) + ax.set_xlim(0, longest * 1.2) ax.xaxis.set_major_formatter(FuncFormatter(lambda v, _: f"{v:g} ms")) def figure_release(old_version, new_version, runs, theme): - gains = [old["baseline_ms"]["median"] / new["baseline_ms"]["median"] for old, new in runs] + gains = [old["baseline_ms"]["median"] / new["baseline_ms"]["median"] for old, new, _ in runs] headline = (f"{min(gains):.1f}–{max(gains):.1f}{TIMES}" if len(gains) > 1 else fmt_multiple(gains[0])) - first_old, first_new = runs[0] + first_old, first_new, _ = runs[0] + reference = next((ref for _, _, ref in runs if ref), None) # Everything is laid out in inches, so the panels keep their size as GPUs are added. - panel, gap, top, bottom = 0.85, 0.62, 1.32, 0.60 - height = top + bottom + panel * len(runs) + gap * (len(runs) - 1) + row_height, gap, top, bottom = 0.363, 0.50, 1.30, 0.52 + panels = [row_height * (3 if reference else 2) for _, _, reference in runs] + height = top + bottom + sum(panels) + gap * (len(runs) - 1) fig = new_figure( theme, f"{new_version} trains {headline} faster than {old_version}", @@ -419,13 +438,15 @@ def figure_release(old_version, new_version, runs, theme): + (" Each GPU has its own scale." if len(runs) > 1 else ""), f"CIFAR-10 · batch {first_new['meta']['batch_size']} · seed {SEED} · " f"median of {first_new['meta']['repeats']} runs · " - f"{first_old['meta']['commit'][:7]} vs {first_new['meta']['commit'][:7]}", + f"{first_old['meta']['commit'][:7]} vs {first_new['meta']['commit'][:7]}" + + (f" · PyTorch {reference['meta']['torch'].split('+')[0]} compiled, fp16 autocast" + if reference else ""), height=height, ) - for i, ((old, new), gain) in enumerate(zip(runs, gains)): - y = bottom + (panel + gap) * (len(runs) - 1 - i) - ax = new_axes(fig, [0.10, y / height, 0.86, panel / height], theme) - release_panel(ax, old_version, new_version, old, new, gain, theme) + for i, ((old, new, reference), gain) in enumerate(zip(runs, gains)): + y = height - top - sum(panels[:i + 1]) - gap * i + ax = new_axes(fig, [0.10, y / height, 0.86, panels[i] / height], theme) + release_panel(ax, old_version, new_version, old, new, reference, gain, theme) handles, labels = ax.get_legend_handles_labels() fig.legend(handles[:3], labels[:3], loc="upper left", bbox_to_anchor=(0.03, 1 - 0.76 / height), @@ -435,17 +456,20 @@ def figure_release(old_version, new_version, runs, theme): def load_release_runs(old_version, new_version): - """(old, new) breakdowns for every GPU, fastest first. All runs must share one pair - of commits, or the footer and the title would be describing different code.""" + """(old, new, torch) for every GPU, fastest first; torch is None on a GPU without a + torch_.json. All runs must share one pair of commits, or the footer and the title + would be describing different code.""" runs = [] for new_path in sorted((HISTORY / new_version).glob("breakdown_*.json")): old_path = HISTORY / old_version / new_path.name if not old_path.exists(): raise FileNotFoundError(f"[aether] {old_path} is missing; {new_path} has no baseline.") - runs.append((json.loads(old_path.read_text()), json.loads(new_path.read_text()))) + torch_path = new_path.with_name(new_path.name.replace("breakdown_", "torch_", 1)) + runs.append((json.loads(old_path.read_text()), json.loads(new_path.read_text()), + json.loads(torch_path.read_text()) if torch_path.exists() else None)) if not runs: raise FileNotFoundError(f"[aether] No breakdown_*.json in {HISTORY / new_version}.") - commits = {(old["meta"]["commit"], new["meta"]["commit"]) for old, new in runs} + commits = {(old["meta"]["commit"], new["meta"]["commit"]) for old, new, _ in runs} if len(commits) > 1: raise ValueError(f"[aether] Release runs measured different commits: {sorted(commits)}.") return sorted(runs, key=lambda run: run[1]["baseline_ms"]["median"]) diff --git a/notebooks/assets/benchmarks/benchmark_release-dark.png b/notebooks/assets/benchmarks/benchmark_release-dark.png index ef833d1..7b2edda 100644 Binary files a/notebooks/assets/benchmarks/benchmark_release-dark.png and b/notebooks/assets/benchmarks/benchmark_release-dark.png differ diff --git a/notebooks/assets/benchmarks/benchmark_release.png b/notebooks/assets/benchmarks/benchmark_release.png index e4dce37..d62bf55 100644 Binary files a/notebooks/assets/benchmarks/benchmark_release.png and b/notebooks/assets/benchmarks/benchmark_release.png differ