the last question

tinygrad PPO on CPU and Metal

Aug 23, 2026

One LunarLander training epoch collects 20,000 transitions from four environments. It then makes 20 passes over those transitions, in minibatches of 10,000. That is 5,000 vector-environment steps and 40 optimizer steps. Every fifth epoch, including epoch zero, it also evaluates the policy on 100 fresh episodes.

During collection, the actor processes four observations at a time. During training, the actor and critic process 10,000 observations, then compute gradients and update their parameters.

Training and evaluation

The experiment runner records the source revision, architecture, package versions, device, and settings. These runs used an Apple M4, Python 3.11, PyTorch 2.8.0, and tinygrad 5b60a09. Each actor and critic has four hidden layers of width 128.

Implementation Transitions at stop Seconds to stop Saved-policy evaluation
PyTorch 520,000 45.60 233.57
tinygrad 520,000 579.42 218.68

I ran each trainer once without setting a seed. The clock starts before imports and stops when a scheduled stochastic evaluation reaches 200. It includes initialization, training, and periodic evaluation. After stopping the clock, I saved and loaded the policy, then evaluated it on 100 fresh episodes. The last column gives that mean return. The tinygrad compiler cache directory was reused.

Evaluation samples actions from the learned Gaussian, clips them to [-1, 1], and averages complete episode returns. Each of the four evaluation environments contributes 25 episodes. Taking the first 100 episodes to finish would overrepresent faster-finishing environments. Training episodes can span several policy updates. Evaluation holds one policy fixed for all 100 episodes.

An earlier tinygrad run used stale cached actor weights. I discarded it. The regression tests compare cached actions with the actor after optimizer steps and checkpoint loading. The learning records include the failed attempt.

Where an epoch spends its time

I ran each trainer for three epochs, one process at a time, with BEAM=1 and the existing compiler cache. PyTorch requested four CPU threads. tinygrad’s CPU thread limit was ten. The protocol and raw records include the commands and settings.

The third epoch:

Implementation Collection Full PPO update Epoch total
PyTorch 0.821 s 0.423 s 1.245 s
tinygrad 9.761 s 4.277 s 14.038 s

Collection includes actor inference, action sampling, environment stepping and IPC, and buffer writes. The update includes old-policy probabilities, critic values, GAE, shuffling, minibatch construction, and all 40 optimizer steps. The epoch timer also covers progress display and scheduled evaluation, which runs every fifth epoch starting at zero.

The first tinygrad epoch took 26.63 seconds: 10.67 collecting, 10.63 updating, and 5.33 outside those phases, including its scheduled evaluation. Each completed run records all three epochs, starting with the initial JIT calls. Both probes ended after 60,000 transitions with negative evaluation scores.

Collection took 9.76 of the third epoch’s 14.04 seconds. Its actor inference, environment, and buffer costs still need separate timings.

A four-thread tinygrad attempt failed while replaying cached beam-search options during its first update. I reran with a ten-thread limit using the same cache directory. The four-thread traceback records the failure.

Controlled optimizer step

I also timed a single optimizer step [0] on a fixed synthetic batch, with matching initial weights and inputs already realized. tinygrad takes 234.56 ms and PyTorch takes 10.58 ms.

Recorded configuration:

The runner starts five fresh process pairs and alternates which backend runs first. Each process reports its median; the tables report the median of those five medians. Ratios are calculated within each pair, then reduced to their median. Each timed operation synchronizes before the clock stops.

Both backends first take three untimed optimizer steps. tinygrad uses these for initial execution, capture, and replay. Five more warmup steps precede ten timed steps. Weights and optimizer state carry over through all 18 steps.

Operation PyTorch tinygrad Paired ratio
One optimizer step, batch 10,000 10.582 ms 234.562 ms 22.1×

The step includes forward, loss, backward, gradient clipping, and Adam.

Which operations account for the step?

The hidden layers multiply (10000×128) activations by (128×128) weights. Backpropagation also forms weight gradients by multiplying (128×10000) transposed activations by (10000×128) output gradients. I measured those shapes separately, plus the four-observation shape used during collection. Each matrix test uses 30 samples after ten warmup operations.

Operation PyTorch tinygrad Paired ratio
(4×128) @ (128×128) 0.0015 ms 0.0844 ms 55.9×
(10000×128) @ (128×128) 0.2260 ms 6.2470 ms 27.7×
(128×10000) @ (10000×128) 0.2328 ms 10.2067 ms 43.7×

tinygrad also reports direct generated-kernel time:

Operation Synchronized TinyJit Direct kernel
(4×128) @ (128×128) 0.0844 ms 0.0212 ms
(10000×128) @ (128×128) 6.2470 ms 6.1975 ms
(128×10000) @ (10000×128) 10.2067 ms 10.0797 ms

For the two larger shapes, the direct-kernel timing is within two percent of the TinyJit median.

PyTorch calls APL_sgemm through Apple Accelerate. tinygrad executes generated CPU kernels.

The matrix checks use a float64 reference; the largest absolute error is 0.001536. The optimizer-step check compares the first three losses, whose largest difference is 2.265e-6. The saved parameter comparison shows a maximum absolute weight difference of 0.001357 after three steps and 0.01617 after all 18 steps.

One update through tinygrad

The loss and Adam updates begin as a lazy graph. A separate capture of this workload contains 1,250 UOps and 68 requested outputs immediately before realization.

Realization schedules and compiles that graph into 150 program calls from 65 unique rendered programs. TinyJit stores this ordered call list and replays it with new input buffers.

The first two calls are the actor and critic input layers. Each takes 10,000 states, an 8×128 weight matrix, and a 128-value bias. tinygrad renders each layer as one C function that contains the matrix reduction and bias addition.

Each captured call records its buffer shapes, input and output slots, launch size, applied schedule options, and rendered source. The complete pipeline trace contains all 150 calls and 65 sources.

Kernel attribution

I enabled HCQ2=1 to collect CPU kernel timestamps in VIZ. HCQ2 uses a different execution path from the timing run above.

One marked HCQ2 step contains 150 kernels. I summed their durations from the timeline, giving 234.46 ms:

Kernel family Time Calls Kernel share
Forward and input-gradient matrices 139.22 ms 12 59.39%
Weight-gradient matrices 77.43 ms 6 33.03%
All other kernels 17.81 ms 132 7.58%

The 18 matrix calls account for 216.65 ms, or 92.42% of the kernel timeline. The other 132 calls account for 17.81 ms.

Generated tinygrad CPU kernel

Clang emits vector loads and fmla v*.4s instructions for the inspected float4 kernel.

CPU programs

A runtime sample of the Apple BLAS worker contains SME fmopa instructions. Clang lowers the tinygrad C function to NEON fmla instructions.

Direct CPU function test

A direct C runner gives both functions identical 64-byte-aligned buffers for (10000×128) @ (128×128). It times calls to the extracted tinygrad C function or Apple’s cblas_sgemm. It uses ten warmup calls and 50 timed calls in five alternating fresh processes.

Function Median Throughput
Apple cblas_sgemm 0.229 ms 1428.2 GFLOPS
tinygrad generated C, beam schedule 2.753 ms 119.0 GFLOPS
tinygrad generated C, default schedule 6.272 ms 52.2 GFLOPS

The paired default-to-BLAS ratio is 27.6 times. Beam makes the generated function 2.25 times faster. The beam-selected function takes 12.2 times longer than BLAS. These timings cover whole routines; instruction choice, packing, and processor use are still combined.

-O2, -O3, and -O3 -ffast-math give similar tinygrad times in this test.

BEAM=1 selects a different schedule before timing. The default forward kernel uses a 4×4 output tile. Beam selects 5×16.

A five-pair optimizer-step comparison measures 261.94 ms for the default schedule and 100.68 ms for beam, with a paired speedup of 2.60 times. [1]

Clang lowers the default inner loop to 16 NEON fmla instructions and the beam inner loop to 20.

Metal programs

The same synchronized matrix test ran on PyTorch MPS and tinygrad Metal.

Path Median Five process medians
PyTorch MPS 0.4273 ms 0.4156, 0.9094, 0.4036, 0.4370, 0.4273 ms
tinygrad Metal 0.9002 ms 0.9002, 1.0712, 0.4602, 0.8184, 1.0281 ms

The paired ratios range from 1.14 to 2.41. The tinygrad Metal function uses simdgroup matrix operations. PyTorch delegates to MPSGraph. The results include all five process pairs.

Evidence







Notes

[0] Schulman et al., Proximal Policy Optimization Algorithms (2017). Section 3, equation (7), defines the clipped policy objective. ↩︎

[1] Goto and van de Geijn, Anatomy of High-Performance Matrix Multiplication (2008). Sections 4 and 6 discuss data packing, memory access, and register blocking. ↩︎