Skip to content

feat(benchmarks): add a compiled PyTorch reference bar to the release figure - #25

Merged
AlexanderSoftCode merged 4 commits into
mainfrom
feat/benchmarks-pytorch-reference
Nov 2, 2026
Merged

AlexanderSoftCode merged 4 commits into
mainfrom
feat/benchmarks-pytorch-reference

Conversation

@AlexanderSoftCode

@AlexanderSoftCode AlexanderSoftCode commented Nov 2, 2026 •

Copy link
Copy Markdown
Owner

Context

benchmark_release shows how much faster aether-ml got between versions, but not how it compares with an established framework. This PR adds PyTorch as a reference on the same CNN and GPU, drawn as a third bar on that GPU's panel.

Commits are split into the move of the shared helpers into common.py, which is mechanical, and the PyTorch script, plot changes, docs and the RX 9070 XT and RTX 3050 results.

PyTorch reference

cifar10_torch.py trains the v1 CNN with the same batch standardization, label-smoothed loss, L2 term, accuracy and AdamW step as Aether, under cifar10_breakdown.py's timing protocol. It runs in its own process because importing CuPy and PyTorch together aborts with an LLVM error, and plot.py joins the two result files. A parameter count check stops the script if the model stops matching build_model.

The whole step goes through torch.compile, with fp16 autocast and no loss scaling. It's a reference point and not a tuned ceiling.

PyTorch 2.12.0, ms/step RX 9070 XT (ROCm) RTX 3050 (CUDA)
Eager, NCHW 8.21 20.0
Compiled, NCHW (committed) 7.60 16.87
Eager, NHWC 11.7 178.7
Compiled, NHWC 12.7 16.9
Aether v1.2 4.86 21.71
Aether v1.2 vs compiled PyTorch 1.6× faster 1.3× slower

Aether v1.2 is 2.1× faster than v1.1 on the RX 9070 XT (10.2 → 4.86 ms/step) and 2.3× faster on the RTX 3050 (49.9 → 21.7 ms/step). On the RX 9070 XT it beats compiled PyTorch. On the RTX 3050 it doesn't: compiled PyTorch is faster there, and Aether is only slightly behind eager PyTorch (20.0 ms).

The layout stays NCHW on both GPUs.

  • ROCm: channels_last alone reverts to NCHW at the first conv, so the NHWC rows also need PYTORCH_MIOPEN_SUGGEST_NHWC=1, and real NHWC is much slower on this GPU.
  • CUDA: compiled NHWC ties compiled NCHW (16.9 vs 16.87 ms), so the committed bar is PyTorch's best on this GPU either way. Eager NHWC is 9× slower than eager NCHW. A profile puts about 93% of that step in cuDNN's channels-last batch-norm kernels, not the convolutions, and cudnn.benchmark doesn't change it. Compiled NHWC avoids them, which fits Inductor replacing the batch-norm calls with its own fused kernels, but I didn't test that directly.

The comparison isn't exact in three ways:

  • The PyTorch bar is wall-clock ms/step, while Aether's bars are GPU time per stage, which runs a few percent above their wall-clock total.
  • Autocast keeps the activations after each Conv2d in fp16, where aether keeps them fp32.
  • Aether's learning-rate decay is left out, because it changes the loss and not the time per step.

Shared helpers

cifar10_torch.py can't import from cifar10_breakdown.py without loading CuPy. The constants, load_train, spread and the run metadata move to common.py, which imports neither library.

Behavior changes

None. Nothing in aether/ changed. PyTorch is only needed by cifar10_torch.py and isn't added to pyproject.toml.

A torch_<gpu>.json next to a GPU's breakdown result adds the bar to that GPU's panel. A GPU without one draws as before.

Testing

  • python3 -m unittest discover tests: 1190 tests pass on NumPy and CuPy (78 skipped).
  • No new regression tests, since nothing in aether/ changed.
  • Reran cifar10_breakdown.py after the refactor on the RX 9070 XT: 4.855 ms/step, matching the saved result.
  • Ran cifar10_torch.py on the RTX 3050 on PyTorch 2.12.0+cu132 with triton-windows 3.9.0 (torch.compile has no Triton on Windows otherwise): 16.87 ms/step, median of 3 (min 16.83, max 16.88). The first compiled step took about 86 s.
  • The other RTX 3050 rows come from one-off variants of the script, not committed. Eager NHWC was run twice and gave 178.7 ms both times.
  • Regenerated all figures with plot.py and viewed the light and dark benchmark_release for layout and labels.
  • Both torch_*.json files record dirty: true, from uncommitted trees at fc06f99 (RX 9070 XT) and 29c71e3 (RTX 3050). Nothing in aether/ is used by that script.

Follow-ups (out of scope)

  • Move to ROCm 10.1 in its own PR, which would mean rebuilding CuPy and rerunning every benchmark on the new stack.
  • Bump the pyproject.toml version to 1.2.0 in a separate PR.

@AlexanderSoftCode
AlexanderSoftCode merged commit 9e62379 into main Nov 2, 2026
2 checks passed
@AlexanderSoftCode
AlexanderSoftCode deleted the feat/benchmarks-pytorch-reference branch November 3, 2026 04:34
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant