Train faster on the CPU

install() also speeds up the parts of a training loop that run outside the model's forward pass. The optimizer step is the largest: torch runs about eight small kernels per parameter tensor from a Python loop, and Kwker runs the whole step in one pass. On 4 cores, an AdamW step over 1.3 million parameters took 0.63 ms instead of 3.3 ms (5.2x), and an SGD step over 11 million parameters 2.0 ms instead of 6.3 ms (3.2x). The parameters come out bit for bit the same.

Prerequisites

• Linux on x86-64 and the Kwker wheel built for your torch release (Installation).
• A training loop on the CPU that uses torch.optim.Adam, AdamW or SGD.

Step 1: a training loop

A small model, a fixed seed and twenty steps:

import torch
import kwker.torch_ops as sto

def train(seed=0, steps=20):
    torch.manual_seed(seed)
    model = torch.nn.Sequential(torch.nn.Linear(64, 128), torch.nn.ReLU(), torch.nn.Linear(128, 10))
    opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
    x, y = torch.randn(256, 64), torch.randint(0, 10, (256,))
    for _ in range(steps):
        opt.zero_grad()
        torch.nn.functional.cross_entropy(model(x), y).backward()
        opt.step()
    return model

reference = train()                    # torch's own optimizer step
print(type(reference[0]).__name__)

Linear

Step 2: the same loop with Kwker

sto.install()
fast = train()                         # the same code, Kwker's fused optimizer step
sto.uninstall()
print(all(torch.equal(a, b) for a, b in zip(reference.parameters(), fast.parameters())))

True

Every parameter is equal to the bit. The fused step follows torch's own rounding, including its square root.

What else gets faster

• Optimizer steps. Adam, AdamW and SGD (with weight decay, momentum, dampening, Nesterov and maximize) on float32 and float64 CPU tensors. AMSGrad, complex parameters, tensor learning rates, gradient scaling and sparse gradients run on torch's code.
• Data loading. In torchvision transforms, the v2 color adjustments, ToDtype + Normalize (fused when they sit next to each other in a Compose) and PIL's Image.resize, each bit for bit the original.
• Embeddings and sparse tensors. The embedding weight gradients, EmbeddingBag in max mode and sparse coalescing.
• Compiled training. torch.compile(model, backend="kwker") also tunes training graphs, at their first call.

Step 3: measure your own training script

python -m kwker audit -- train.py --epochs 1

It runs the script several times with and without Kwker, alternating, and reports the wall time, the CPU time and whether the outputs agree.

Next steps

• Speed up inference with one line: the same install() for inference.
• Compile with torch.compile: the backend for training and inference graphs.
• Runtime controls: switch single features off.