The problem with why its slower is that the task is too simple and thus the overhead overshadows the speed up. This is because polynomial_autograd.py is simple (basic arithmetic). Here checkout some more complex operations and look at the speed up
import time
import torch
import sys
DEVICE = sys.argv[1] if len(sys.argv) > 1 else "cpu"
WARMUP = 3
ITERS = 10
def bench(name, fn, warmup=WARMUP, iters=ITERS):
for _ in range(warmup):
fn()
if DEVICE == "mps":
torch.mps.synchronize()
times = []
for _ in range(iters):
t0 = time.perf_counter()
fn()
if DEVICE == "mps":
torch.mps.synchronize()
t1 = time.perf_counter()
times.append(t1 - t0)
mean = sum(times) / len(times)
std = (sum((t - mean) ** 2 for t in times) / len(times)) ** 0.5
print(f" {name:40s} {mean*1000:8.2f} ms +/- {std*1000:.2f} ms")
return mean
print(f"DeVICE: {DEVICE} | PyTorch {torch.__version__}\n", flush=True)
# 1. Matrix multiply (large, compute-bound)
N = 8192
a = torch.randn(N, N, device=DEVICE, dtype=torch.float32)
b = torch.randn(N, N, device=DEVICE, dtype=torch.float32)
bench(f"matmul {N}x{N}", lambda: torch.mm(a, b))
# 2. Convolution (batched, spatial)
B, C, H, W = 128, 64, 56, 56
K = 128
conv = torch.nn.Conv2d(C, K, 3, padding=1, device=DEVICE)
inp = torch.randn(B, C, H, W, device=DEVICE)
bench(f"conv2d {B}x{C}x{H}x{W} -> {K}", lambda: conv(inp))
# 3. Batched MV with 512x1024x1024 random mats
M_SIZE = 1024
BATCH = 512
matrices = torch.randn(BATCH, M_SIZE, M_SIZE, device=DEVICE)
vectors = torch.randn(BATCH, M_SIZE, 1, device=DEVICE)
bench(f"bmm {BATCH}x{M_SIZE}x{M_SIZE}", lambda: torch.bmm(matrices, vectors))
# 4. SoftMAX AND layer norm over large fake sequence
SEQ, DIM = 4096, 2048
x = torch.randn(SEQ, DIM, device=DEVICE)
bench(f"softmax {SEQ}x{DIM}", lambda: torch.softmax(x, dim=-1))
bench(f"layernorm {SEQ}x{DIM}", lambda: torch.nn.functional.layer_norm(x, (DIM,)))
# 5. Small CNN training loop
C_IN, C_OUT = 64, 64
model = torch.nn.Sequential(
torch.nn.Conv2d(C_IN, C_OUT, 3, padding=1, device=DEVICE),
torch.nn.ReLU(),
torch.nn.Conv2d(C_OUT, C_OUT, 3, padding=1, device=DEVICE),
torch.nn.ReLU(),
torch.nn.Conv2d(C_OUT, C_OUT, 3, padding=1, device=DEVICE),
).to(DEVICE)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
data = torch.randn(64, C_IN, 32, 32, device=DEVICE)
target = torch.randn(64, C_OUT, 32, 32, device=DEVICE)
def training_step():
optimizer.zero_grad()
out = model(data)
loss = torch.nn.functional.mse_loss(out, target)
loss.backward()
optimizer.step()
WARMUP_SMALL = 20
ITERS_SMALL = 50
for _ in range(WARMUP_SMALL):
training_step()
if DEVICE == "mps":
torch.mps.synchronize()
times = []
for _ in range(ITERS_SMALL):
t0 = time.perf_counter()
training_step()
if DEVICE == "mps":
torch.mps.synchronize()
t1 = time.perf_counter()
times.append(t1 - t0)
mean = sum(times) / len(times)
std = (sum((t - mean) ** 2 for t in times) / len(times)) ** 0.5
print(f" {'training step (3-layer CNN)':40s} {mean*1000:8.2f} ms +/{std*1000:.2f} ms")
Note: we use time.perf_counter instead of time.time as it is more reliable (not affected by external system things)
| Benchmark | CPU | MPS | Speedup |
|---|---|---|---|
| matmul 8192×8192 | 1366 ms | 696 ms | 1.96× |
| conv2d 128×64×56×56 | 179 ms | 16.6 ms | 10.8× |
| bmm 512×1024×1024 | 48.2 ms | 26.3 ms | 1.83× |
| softmax 4096×2048 | 5.4 ms | 1.2 ms | 4.6× |
| layernorm 4096×2048 | 2.4 ms | 1.1 ms | 2.3× |
| training step (3-layer CNN) | 216 ms | 23.9 ms | 9.0× |
see how on the CNN it gets around 9x. This was done on a macbook m1 air. The warmup is there so the timings do not include one time Metal shader compilation… So therefore, most places (especially for real heavy work)
note: you’ll probably see the speed up with changing x = torch.linspace(-1, 1, 2000, dtype=dtype) to torch.linspace(-math.pi, math.pi, 2000) in the example
https://raw.githubusercontent.com/pytorch/tutorials/refs/heads/main/beginner_source/examples_autograd/polynomial_autograd.py