A trainer under torchrun and accelerate launch

A plain PyTorch trainer whose kernels are .ppy modules, started by python, torchrun, or accelerate launch. import ppy is the whole integration: each rank finds the kernels' native build in the project cache, and preprocessing that took 110 ms per rank in Python takes 1.6 ms. There is no bootstrap, no launcher of its own, and no ppy run.

import ppy  # first: the import hook, serving .ppy modules from their native build
import torch

import features  # features.ppy: the preprocessing loops, native
import model  # model.ppy: the model as one ATen region, native

Run it

ppy build --warm .
python train.py
torchrun --standalone --nproc_per_node=2 train.py
accelerate launch --multi_gpu --num_processes 2 train.py
PPY_IMPORT=python python train.py
ppy check .
python features.ppy && ppy run features.ppy
python model.ppy    && ppy run model.ppy

What it prints

ppy build --warm ., ppy check .

(prints nothing; exits 0)

python train.py

# rank 0/1 device=cpu native=True region=True loader=GeneratedLoader
rank 0: prep       1.9 ms   checksum=-21015.470416 outside=4103
rank 0: train  16513.7 ms   loss 1.0412 -> 1.0107

torchrun --standalone --nproc_per_node=2 train.py

# rank 0/2 device=cpu native=True region=True loader=GeneratedLoader
rank 0: prep       1.9 ms   checksum=-21015.470416 outside=4103
rank 0: train    783.5 ms   loss 1.0412 -> 1.0123
# rank 1/2 device=cpu native=True region=True loader=GeneratedLoader
rank 1: prep       1.9 ms   checksum=-20883.104755 outside=4005
rank 1: train    781.9 ms   loss 1.0486 -> 1.0239

accelerate launch --multi_gpu --num_processes 2 train.py

not run here: accelerate is not installed

PPY_IMPORT=python python train.py

# rank 0/1 device=cpu native=False region=False loader=PPySourceLoader
rank 0: prep     125.0 ms   checksum=-21015.470416 outside=4103
rank 0: train  12706.6 ms   loss 1.0412 -> 1.0107

python features.ppy && ppy run features.ppy

checksum=-2083.056718 outside=134 counts=[452, 3853, 11555, 16790, 16661, 11022, 3174, 359]
checksum=-2083.056718 outside=134 counts=[452, 3853, 11555, 16790, 16661, 11022, 3174, 359]

python model.ppy && ppy run model.ppy

# aten region: False
loss=4.584410
# aten region: True
loss=4.584410

accelerate is not a dependency of this repository. --multi_gpu is what makes it launch several processes, and the script runs them on the CPU when there is no CUDA device.

How the launcher starts it

The program does not know how it was started. train.py reads the environment every launcher agrees on (RANK, WORLD_SIZE, LOCAL_RANK) and runs the same under python, torchrun, and accelerate launch. The launcher starts ordinary interpreters, and each one's import ppy finds the kernels' native build.

The two kernels

features.ppy holds the per-batch arithmetic as loops over borrowed buffers: standardizing rows, appending interactions, and bucketing the result. That is where the time goes on the Python side of a trainer, and where the native path wins.

model.ppy imports torch. Its forward_loss is one function of curated tensor operations, compiled into an ATen region that ships beside the manifest and loads without the compiler in the process. .backward() sees the same graph, because every at:: call still goes through the dispatcher.

Warm the cache before the launch

The first process to import a kernel builds it into .ppy-cache/, and every process after that finds the build. Under a launcher the ranks start together. Without a warm cache each rank builds the same artifact and the first to finish is kept, which is correct but paid for N times.

ppy build --warm . before the launch builds every kernel once, and no rank builds anything. A cold first build of the torch region takes tens of seconds, once.

Results

CPU, 2 ranks, each on 20,000 rows × 16 columns; 100 steps of a 32-unit MLP.

per rank	PPY_IMPORT=python	import ppy
preprocessing (standardize + bucketize)	109.9 ms	1.6 ms
100 training steps (forward_loss region)	399 ms	250–600 ms

The preprocessing is the point. The region removes four Python round trips per step, which 21_training_torch measures at about 20% in isolation. Here the step is dominated by the tensor work, and the run-to-run noise of a two-rank CPU launch is wider than the gain. On an accelerator the region changes nothing measurable.

• PPY_IMPORT=python runs the same files as plain Python.
• PPY_QUIET=1 silences the per-rank notes.

Read on: Interop · CLI: ppy build --warm · Migrating a real project

features.ppy, model.ppy, and train.py are hand-written; there is no conversion step.

31_torchrun/features.ppy

"""The per-batch preprocessing of a trainer, as native loops over borrowed buffers.

`standardize` centers and scales every row and appends the pairwise
interactions the model trains on; `bucketize` counts how the standardized
values spread, the kind of running statistic a trainer logs each step. Both
index rather than slice, so their buffers are borrowed and the loops lower
to native code writing into memory the caller owns.
"""

import array
import math

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


def bucketize(
    values: Buffer[float], count: int, counts: Buffer[int], bins: int, low: float, high: float
) -> int:
    """Count the first `count` values into `bins` equal slots over [low, high).

    Returns how many fell outside the range.
    """
    width: float = (high - low) / bins
    outside: int = 0
    for i in range(count):
        offset: float = (values[i] - low) / width
        if offset < 0.0 or offset >= bins:
            outside += 1
        else:
            counts[int(offset)] += 1
    return outside


def _check() -> None:
    """A checksum on a fixed input, so the three paths can be compared."""
    rows: int = 2000
    cols: int = 16
    seed: int = 12345
    raw = array.array("d", [0.0] * (rows * cols))
    for i in range(rows * cols):
        seed = (seed * 1103515245 + 12345) % 2147483648
        raw[i] = seed / 2147483648.0 - 0.5
    out = array.array("d", [0.0] * (rows * cols * 2))
    counts = array.array("q", [0] * 8)
    total: float = standardize(raw, out, rows, cols)
    outside: int = bucketize(out, rows * cols * 2, counts, 8, -3.0, 3.0)
    print(f"checksum={total:.6f} outside={outside} counts={list(counts)}")


if __name__ == "__main__":
    _check()

31_torchrun/model.ppy

"""The model: one function of curated tensor operations, compiled into one ATen region.

Every `at::` call inside the region still goes through the dispatcher, so
autograd sees the graph it always did and `.backward()` is unchanged; what
the region removes is the Python round trip between operators.
"""

import torch


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 _check() -> None:
    """The loss on a fixed input, so the three paths can be compared."""
    torch.manual_seed(0)
    x = torch.randn(8, 4)
    y = torch.randn(8, 1)
    w1 = torch.randn(4, 3)
    w2 = torch.randn(3, 1)
    print("# aten region:", getattr(forward_loss, "__ppy_region__", False))
    print(f"loss={forward_loss(x, y, w1, torch.zeros(3), w2, torch.zeros(1)).item():.6f}")


if __name__ == "__main__":
    _check()

Source: examples/31_torchrun.