Training Small Digit Adder Transformer on my Laptop

I explore how very small ML models might in some cases be better trained on CPU than GPU. I use the recent AdderBoard as a source for models and inspiration. In one of these I was able to achieve an approximately 8-9x per-iteration speedup on my laptop (CPU, custom C++) compared with the PyTorch A100 runs below: 17.9/2.2 for the compiled CUDA-Graphs run and 20.1/2.2 for the baseline.

These are very small transformer models designed to reliably compute integer addition. I suspected these models were small enough that framework and launch overheads might dominate on GPU (for the training version of this competition). So I picked the best trained adder listed when I ran the experiment, followed the instructions to reproduce training, and built an optimized version of its training loop.

My interpretation is that the speedup is primarily from removing framework/dispatch/allocation overhead on the CPU side and avoiding GPU kernel-launch overhead for tiny kernels. Cache residency and parallelism likely contribute too.

The code for my C++ implementation of the training lives in this repository.

Methodology: What Is an Iteration?

Per-iteration timing in this post is intended to cover:

• Batch/data generation (when done inside the training loop)
• Forward pass
• Loss computation
• Backward pass
• Optimizer step

Excluded from per-iteration timing:

• Validation
• Checkpointing
• Logging/progress I/O

Warmup / startup handling:

• Startup is reported separately in the tables as Excl. Startup Time and Excl. Startup Iters
• Iteration 0 is treated as startup in the GPU table (so 999 steady-state iters)
• CPU C++ warmup was not separately instrumented; reported value is from repeated loop timing

Quick Look at GPU Performance (original)

I spun up a GCP A100 instance and followed the README to run training from scratch. I also traced execution to understand overhead and utilization.

Since the model is so small there is significant kernel-launch overhead, seen as gaps between meaningful compute on the trace.

Omitting traced results, per-iteration latency is about 20ms.

Optimize with torch.compile and CUDA Graphs

Given the launch overhead, I looked for a quick win with CUDA Graphs. CUDA Graphs record a sequence of launches and replay them, which often reduces launch/dispatch overhead for fixed-shape workloads. It helps here, but not dramatically.

Interpretation: total wall time can get worse while steady-state per-iteration improves, because torch.compile adds upfront compilation overhead.

If you care about time-to-first-1000-steps, baseline can still win. If you care about long steady-state runs, compile/graphs can win.

Quick Look at CPU Performance (original)

As alluded to earlier, I suspected this model was tiny enough that GPU overhead (kernel launch and framework plumbing) could offset raw hardware advantage for this specific workload. I also ran PyTorch on CPU as a rough sanity check.

Environment details for CPU sanity check:

• CPU model and exact runtime settings were not recorded for that run
• I only noted that the instance was likely around 12 cores
• torch.get_num_threads(), OMP_NUM_THREADS, MKL_NUM_THREADS, and BLAS backend were not captured

Laptop (inlined C++)

The GPU beats untuned PyTorch CPU on that system, but this model is so small that it should fit in cache, so a specialized C++ implementation can push overhead down a lot. I asked Codex to write a single-file C++ training loop with model-size limits hardcoded as constexpr, which gives the compiler room for aggressive optimization (unrolling/vectorization).

This performed surprisingly well on my laptop. The implementation also has OpenMP parallelism. One detail worth noting is that -fopenmp-simd alone does not enable OpenMP threading in GCC; the full build must also link with -fopenmp.

Raw text for the exact compile + run commands used:

CXXFLAGS="-Ofast -march=native -mtune=native -flto -ffast-math -funroll-loops -fopenmp-simd -fstrict-aliasing -fno-trapping-math -fno-math-errno -falign-functions=32 -falign-loops=32" make
./train_311p --steps 500000 --lr 0.02 --lr-stages 162000:0.001,262000:0.0003,312000:0.0001,400000:0.00005

Fairness / Controls / Caveats

What I held constant (to the extent reproduced from the training script):

• Model architecture and size
• Intended core training setup from the referenced script (batch size, sequence length, vocabulary/tokenization, optimizer, learning rate, seed, dtype)

What I did not fully control (or did not record well enough):

• CPU thread affinity/pinning
• BLAS backend and thread-runtime details (MKL/oneDNN/OpenMP runtime behavior)
• GPU clock/power settings
• Exact synchronization points around timing on GPU
• Exact parity of startup/warmup handling across all runs
• Exact numerical equivalence between PyTorch and the -Ofast C++ implementation

This comparison is not apples-to-apples hardware benchmarking (PyTorch on A100 vs hand-written C++ on a laptop). I still think it is useful as a case study: for tiny models, overheads can dominate enough that a specialized CPU implementation wins on per-iteration latency.

A realistic GPU path here is likely better operator fusion and fewer launches, rather than assuming one giant shared-memory “megakernel” is always practical.

Validation

During training I withhold the eight sums shown below and validate against them. The delayed transition to exact answers resembles the grokking behavior reported by the original author. This is a small fixed validation set, but the behavior is striking.

• 100k checkpoint: relied on a “2-repetition” failure mode, leading to high variance and an average error of 7.22%.
• 300k checkpoint: developed a “9-rounding” strategy. While it still fails exact arithmetic, approximation is much tighter, bringing average error to 2.97%.
• 500k checkpoint: exact match on this eight-example validation set.

Conclusion

When a problem is this small, overhead terms can dominate and move the result in unintuitive ways. Here the trace points to framework overhead and frequent tiny GPU launches, while specialization, vectorization, cache residency, and threading all help the C++ implementation.

Apparent superlinear effects can also show up when a per-thread working set moves into faster caches, although I did not isolate that effect here.

I still think there is room to recover more GPU performance here with better fusion and fewer launches; a monolithic training megakernel is one extreme, but not the only path.