Skip to content

Training a torch MLP

An ordinary PyTorch training script, converted by ppy convert with no hand editing. It standardizes 20,000 rows of features in a Python loop, then runs 100 steps of a two-layer MLP. The preprocessing loop went from about 70 ms to under 1 ms, and the training step became one ATen region.

Run it

python  train.ppy
ppy run train.ppy

What it prints

python train.ppy

# device: cpu
# native prep: False
# aten region: False
prep      66.7 ms   checksum=-21282.454900
train    142.6 ms   loss 1.0250 -> 1.0019

ppy run train.ppy

# device: cpu
# native prep: True
# aten region: True
prep       1.1 ms   checksum=-21282.454900
train    320.7 ms   loss 1.0250 -> 1.0019

Where the speedup is

def standardize(raw: Buffer[float], out: Buffer[float], rows: int, cols: int) -> float:

standardize reads and writes by index only, so --promote-buffers declared its parameters Buffer[float] and rewrote the values feeding them into array.array. The loop lowers to native code writing into memory the caller owns. The 20,000×16 standardization that dominated the script's Python time disappears from the profile.

Where it is not

forward_loss is a single function of curated tensor operations (matmul, add, relu, sub, mul, mean), so the torch plugin compiles it into one C++ ATen region. Every at:: call still goes through the dispatcher, and .backward() sees the same graph.

What the region removes is the Python round trip between operators. Over a 20,000-row batch the matmul is the time and the round trips are not, so the step runs at PyTorch's speed. On an accelerator, kernel launch latency dominates the same way.

What the input had to get right

  • standardize indexes rather than slices. A slice copies, so a sliced parameter cannot be borrowed, and the converter says so: remark[R3003]: raw is sliced, which copies.
  • descend narrows parameter.grad before using it. It is Tensor | None until a backward pass fills it, and the unnarrowed version is a latent crash that ppy check reports.
  • forward_loss is one function. Split across two, every operator pays a Python round trip and no region forms.

Compared with PyTorch as it is usually written

The benchmark covers the standardization of 20,000×16 rows and one forward pass of the step's operators. The programs are in compare/:

Milliseconds, best of five (the forward pass best of two hundred), over five processes; eight PyTorch threads.

PPy standardizes in the loop the script was written with, over Buffer[float]. PyTorch is how the same preprocessing is written for PyTorch, vectorized over the batch, with roll for the interaction term:

def standardize(raw: Buffer[float], out: Buffer[float], rows: int, cols: int) -> float:
    total: float = 0.0
    for row in range(rows):
        base: int = row * cols
        target: int = row * cols * 2
        sum_: float = 0.0
        for i in range(cols):
            sum_ += raw[base + i]
        mean: float = sum_ / cols
        ...

def standardize(raw: torch.Tensor) -> tuple[torch.Tensor, float]:
    mean = raw.mean(dim=1, keepdim=True)
    deviation = torch.sqrt(((raw - mean) ** 2).mean(dim=1, keepdim=True)) + 1e-8
    z = (raw - mean) / deviation
    interaction = z * z.roll(-1, dims=1)
    out = torch.cat([z, interaction], dim=1)
    return out, float(out.sum())

forward_loss is the same six operators on every side: one ATen region under PPy, eager under PyTorch, and under torch.compile in the third column.

PPy ppy run CPython, the same file PyTorch, vectorized torch.compile
standardize, 20000 rows 0.83 ± 0.02 62.74 ± 0.20 0.72 ± 0.04 0.72 ± 0.08
forward pass, per call 0.11 ± 0.01 0.11 ± 0.00 0.11 ± 0.00 0.11 ± 0.00
  • Standardize. The loop as the script wrote it, compiled, is level with the vectorized rewrite. Both are one pass over the rows, and the rewrite costs seven tensor temporaries the loop never makes. Under python the same loop is the cost the conversion removed.
  • Forward pass. This is PyTorch's on every side: the matmul over 20,000 rows is the time, and the Python round trips between six operators are not. The region neither gains nor loses there, and torch.compile's guards are a few microseconds on top.

Intel Core Ultra 9 386H; PyTorch 2.14.0 (CPU) on CPython 3.13.13, PPy against the same PyTorch on CPython 3.14.5, from a checkout on a native filesystem.

Where the code comes from

Generated, not hand-written: train.ppy is exactly what ppy convert train.py --promote-buffers writes, and examples/verify_conversions.py checks that on every run.

Read on: PyTorch regions · Training with JAX · Plugins: PyTorch

21_training_torch/train.ppy

import array
import math
import time
from collections.abc import Sequence
from typing import Literal

import ppy
import torch
from ppy import Buffer


def standardize(raw: Buffer[float], out: Buffer[float], rows: int, cols: int) -> float:
    total: float = 0.0
    for row in range(rows):
        base: int = row * cols
        target: int = row * cols * 2

        sum_: float = 0.0
        for i in range(cols):
            sum_ += raw[base + i]
        mean: float = sum_ / cols

        spread: float = 0.0
        for i in range(cols):
            spread += (raw[base + i] - mean) * (raw[base + i] - mean)
        deviation: float = math.sqrt(spread / cols) + 1e-8

        for i in range(cols):
            value: float = (raw[base + i] - mean) / deviation
            out[target + i] = value
            total += value
        for i in range(cols):
            interaction: float = out[target + i] * out[target + (i + 1) % cols]
            out[target + cols + i] = interaction
            total += interaction
    return total


@ppy.pure
def forward_loss(
    x: torch.Tensor,
    y: torch.Tensor,
    w1: torch.Tensor,
    b1: torch.Tensor,
    w2: torch.Tensor,
    b2: torch.Tensor,
) -> torch.Tensor:
    hidden: torch.Tensor = torch.relu(torch.add(torch.matmul(x, w1), b1))
    predicted: torch.Tensor = torch.add(torch.matmul(hidden, w2), b2)
    residual: torch.Tensor = torch.sub(predicted, y)
    return torch.mean(torch.mul(residual, residual))


def descend(params: Sequence[torch.Tensor], rate: float) -> None:
    with torch.no_grad():
        for parameter in params:
            gradient = parameter.grad
            if gradient is None:
                continue
            parameter -= rate * gradient
            parameter.grad = None


def train_step(
    x: torch.Tensor,
    y: torch.Tensor,
    params: Sequence[torch.Tensor],
    rate: float,
) -> float:
    loss: torch.Tensor = forward_loss(x, y, params[0], params[1], params[2], params[3])
    loss.backward()
    descend(params, rate)
    return loss.item()


@ppy.pure
def preferred_device() -> Literal['cuda', 'cpu']:
    if torch.cuda.is_available():
        return "cuda"
    return "cpu"


def main() -> None:
    rows: int = 20000
    cols: int = 16
    device: str = preferred_device()

    torch.manual_seed(0)
    raw = array.array("d", torch.randn(rows * cols).tolist())
    features = array.array("d", [0.0] * (rows * cols * 2))

    started: float = time.perf_counter()
    checksum: float = standardize(raw, features, rows, cols)
    prep_ms: float = (time.perf_counter() - started) * 1000.0

    x: torch.Tensor = torch.tensor(features, dtype=torch.float32).reshape(rows, cols * 2).to(device)
    y: torch.Tensor = torch.randn(rows, 1).to(device)
    params: list[torch.Tensor] = [
        torch.randn(cols * 2, 32, device=device, requires_grad=True),
        torch.zeros(32, device=device, requires_grad=True),
        torch.randn(32, 1, device=device, requires_grad=True),
        torch.zeros(1, device=device, requires_grad=True),
    ]
    with torch.no_grad():
        params[0] *= 0.1
        params[2] *= 0.1

    started = time.perf_counter()
    first: float = 0.0
    last: float = 0.0
    for index in range(100):
        last = train_step(x, y, params, 0.02)
        if not index:
            first = last
    train_ms: float = (time.perf_counter() - started) * 1000.0

    print(f"# device: {device}")
    print(f"# native prep: {ppy.native.compiled(standardize)}")
    print(f"# aten region: {getattr(forward_loss, '__ppy_region__', False)}")
    print(f"prep  {prep_ms:8.1f} ms   checksum={checksum:.6f}")
    print(f"train {train_ms:8.1f} ms   loss {first:.4f} -> {last:.4f}")


main()

Counterpart programs

The programs the comparison above measured, each written the way its tool expects. The PPy one is first.

train_bench.ppy (PPy) 
"""The trainer of `train.ppy`, timed: the standardization loop, and one forward pass of the
step's operators. The PPY side; the same file under `python` is the CPython column."""

import array
import math
import time
from collections.abc import Sequence

import ppy
import torch
from ppy import Buffer


def standardize(raw: Buffer[float], out: Buffer[float], rows: int, cols: int) -> float:
    total: float = 0.0
    for row in range(rows):
        base: int = row * cols
        target: int = row * cols * 2

        sum_: float = 0.0
        for i in range(cols):
            sum_ += raw[base + i]
        mean: float = sum_ / cols

        spread: float = 0.0
        for i in range(cols):
            spread += (raw[base + i] - mean) * (raw[base + i] - mean)
        deviation: float = math.sqrt(spread / cols) + 1e-8

        for i in range(cols):
            value: float = (raw[base + i] - mean) / deviation
            out[target + i] = value
            total += value
        for i in range(cols):
            interaction: float = out[target + i] * out[target + (i + 1) % cols]
            out[target + cols + i] = interaction
            total += interaction
    return total


@ppy.pure
def forward_loss(
    x: torch.Tensor,
    y: torch.Tensor,
    w1: torch.Tensor,
    b1: torch.Tensor,
    w2: torch.Tensor,
    b2: torch.Tensor,
) -> torch.Tensor:
    hidden: torch.Tensor = torch.relu(torch.add(torch.matmul(x, w1), b1))
    predicted: torch.Tensor = torch.add(torch.matmul(hidden, w2), b2)
    residual: torch.Tensor = torch.sub(predicted, y)
    return torch.mean(torch.mul(residual, residual))


def descend(params: Sequence[torch.Tensor], rate: float) -> None:
    with torch.no_grad():
        for parameter in params:
            gradient = parameter.grad
            if gradient is None:
                continue
            parameter -= rate * gradient
            parameter.grad = None


def train_step(
    x: torch.Tensor,
    y: torch.Tensor,
    params: Sequence[torch.Tensor],
    rate: float,
) -> float:
    loss: torch.Tensor = forward_loss(x, y, params[0], params[1], params[2], params[3])
    loss.backward()
    descend(params, rate)
    return loss.item()


def main() -> None:
    rows: int = 20000
    cols: int = 16
    torch.manual_seed(0)
    torch.set_num_threads(8)
    raw = array.array("d", torch.randn(rows * cols).tolist())
    features = array.array("d", [0.0] * (rows * cols * 2))
    standardize(raw, features, 10, cols)
    best_prep = 1e9
    for _ in range(5):
        started: float = time.perf_counter()
        checksum: float = standardize(raw, features, rows, cols)
        best_prep = min(best_prep, (time.perf_counter() - started) * 1000.0)
    x: torch.Tensor = torch.tensor(features, dtype=torch.float32).reshape(rows, cols * 2)
    y: torch.Tensor = torch.randn(rows, 1)
    torch.manual_seed(1)
    params: list[torch.Tensor] = [
        torch.randn(cols * 2, 32, requires_grad=True),
        torch.zeros(32, requires_grad=True),
        torch.randn(32, 1, requires_grad=True),
        torch.zeros(1, requires_grad=True),
    ]
    with torch.no_grad():
        params[0] *= 0.1
        params[2] *= 0.1
    last: float = train_step(x, y, params, 0.02)
    best_forward = 1e9
    with torch.no_grad():
        for _ in range(200):
            started = time.perf_counter()
            loss: float = forward_loss(x, y, params[0], params[1], params[2], params[3]).item()
            best_forward = min(best_forward, (time.perf_counter() - started) * 1000.0)
    print(f"# standardize, 20000 rows: {best_prep:.2f} ms")
    print(f"# forward pass, per call: {best_forward:.4f} ms")
    print(f"{checksum:.4f} {last:.4f} {loss:.4f}")


main()
train_compile.py (Python) 
"""The same trainer with `forward_loss` under `torch.compile`: Inductor fuses the operators."""

import time

import torch


def standardize(raw: torch.Tensor) -> tuple[torch.Tensor, float]:
    mean = raw.mean(dim=1, keepdim=True)
    deviation = torch.sqrt(((raw - mean) ** 2).mean(dim=1, keepdim=True)) + 1e-8
    z = (raw - mean) / deviation
    interaction = z * z.roll(-1, dims=1)
    out = torch.cat([z, interaction], dim=1)
    return out, float(out.sum())


@torch.compile
def forward_loss(x, y, w1, b1, w2, b2):
    hidden = torch.relu(torch.add(torch.matmul(x, w1), b1))
    predicted = torch.add(torch.matmul(hidden, w2), b2)
    residual = torch.sub(predicted, y)
    return torch.mean(torch.mul(residual, residual))


def train_step(x, y, params, rate):
    loss = forward_loss(x, y, params[0], params[1], params[2], params[3])
    loss.backward()
    with torch.no_grad():
        for weight in params:
            if weight.grad is not None:
                weight -= rate * weight.grad
                weight.grad = None
    return loss.item()


def main():
    rows, cols = 20000, 16
    torch.manual_seed(0)
    torch.set_num_threads(8)
    raw = torch.randn(rows * cols).to(torch.float64).reshape(rows, cols)
    standardize(raw[:10])
    best_prep = 1e9
    for _ in range(5):
        started = time.perf_counter()
        features, checksum = standardize(raw)
        best_prep = min(best_prep, (time.perf_counter() - started) * 1000.0)
    x = features.to(torch.float32)
    y = torch.randn(rows, 1)
    torch.manual_seed(1)
    params = [
        torch.randn(cols * 2, 32, requires_grad=True),
        torch.zeros(32, requires_grad=True),
        torch.randn(32, 1, requires_grad=True),
        torch.zeros(1, requires_grad=True),
    ]
    with torch.no_grad():
        params[0] *= 0.1
        params[2] *= 0.1
    last = train_step(x, y, params, 0.02)
    best_forward = 1e9
    with torch.no_grad():
        for _ in range(200):
            started = time.perf_counter()
            loss = forward_loss(x, y, params[0], params[1], params[2], params[3]).item()
            best_forward = min(best_forward, (time.perf_counter() - started) * 1000.0)
    print(f"# standardize, 20000 rows: {best_prep:.2f} ms")
    print(f"# forward pass, per call: {best_forward:.4f} ms")
    print(f"{checksum:.4f} {last:.4f} {loss:.4f}")


main()
train_eager.py (Python) 
"""The same trainer as PyTorch is written: the standardization vectorized over the batch,
the forward pass eager, one Python round trip per operator."""

import time

import torch


def standardize(raw: torch.Tensor) -> tuple[torch.Tensor, float]:
    mean = raw.mean(dim=1, keepdim=True)
    deviation = torch.sqrt(((raw - mean) ** 2).mean(dim=1, keepdim=True)) + 1e-8
    z = (raw - mean) / deviation
    interaction = z * z.roll(-1, dims=1)
    out = torch.cat([z, interaction], dim=1)
    return out, float(out.sum())


def forward_loss(x, y, w1, b1, w2, b2):
    hidden = torch.relu(torch.add(torch.matmul(x, w1), b1))
    predicted = torch.add(torch.matmul(hidden, w2), b2)
    residual = torch.sub(predicted, y)
    return torch.mean(torch.mul(residual, residual))


def train_step(x, y, params, rate):
    loss = forward_loss(x, y, params[0], params[1], params[2], params[3])
    loss.backward()
    with torch.no_grad():
        for weight in params:
            if weight.grad is not None:
                weight -= rate * weight.grad
                weight.grad = None
    return loss.item()


def main():
    rows, cols = 20000, 16
    torch.manual_seed(0)
    torch.set_num_threads(8)
    raw = torch.randn(rows * cols).to(torch.float64).reshape(rows, cols)
    standardize(raw[:10])
    best_prep = 1e9
    for _ in range(5):
        started = time.perf_counter()
        features, checksum = standardize(raw)
        best_prep = min(best_prep, (time.perf_counter() - started) * 1000.0)
    x = features.to(torch.float32)
    y = torch.randn(rows, 1)
    torch.manual_seed(1)
    params = [
        torch.randn(cols * 2, 32, requires_grad=True),
        torch.zeros(32, requires_grad=True),
        torch.randn(32, 1, requires_grad=True),
        torch.zeros(1, requires_grad=True),
    ]
    with torch.no_grad():
        params[0] *= 0.1
        params[2] *= 0.1
    last = train_step(x, y, params, 0.02)
    best_forward = 1e9
    with torch.no_grad():
        for _ in range(200):
            started = time.perf_counter()
            loss = forward_loss(x, y, params[0], params[1], params[2], params[3]).item()
            best_forward = min(best_forward, (time.perf_counter() - started) * 1000.0)
    print(f"# standardize, 20000 rows: {best_prep:.2f} ms")
    print(f"# forward pass, per call: {best_forward:.4f} ms")
    print(f"{checksum:.4f} {last:.4f} {loss:.4f}")


main()

Source: examples/21_training_torch.