Autotrainer

Run your training loop fast on the hardware you have.

Autotrainer detects your hardware (local GPUs or a SLURM cluster), picks the distribution strategy, applies the throughput wins people forget, and tells you when a run is slow or broken — without touching your lr, loss, schedule, or optimizer.

Works with PyTorch (DDP, SLURM multi-node), TensorFlow/Keras, scikit-learn, XGBoost, and LightGBM through one API.

Install

pip install "autotrainer[torch]"   # PyTorch (recommended)
pip install "autotrainer[all]"     # + TensorFlow, sklearn, XGBoost, LightGBM

Only psutil is required by default; each framework is an opt-in extra (torch, sklearn, tf, boosting, all). Install only what you use.

Use it

Two lines, full training run — infers loss/optimizer/schedule, applies the hardware wins, saves the checkpoint, prints a summary:

import autotrainer

model = autotrainer.train(model, loader, epochs=5, save_path="model.pt")

Or keep your own loop — prepare() returns the same three objects you passed in, now distributed and optimized:

import autotrainer

model, loader, optimizer = autotrainer.prepare(model, loader, optimizer)

for epoch in range(epochs):
    autotrainer.set_epoch(loader, epoch)   # reshuffle each epoch
    ...                                    # your normal training loop

autotrainer.finish()                       # prints the training summary

Then run it. No launcher to configure — on a multi-GPU box prepare() spawns one worker per GPU by itself:

python train.py                  # local: auto-distributes across GPUs
srun autotrainer run train.py    # SLURM multi-node
autotrainer doctor               # check the environment before you burn an allocation
autotrainer ui                   # zero-dependency dashboard for your runs

Why use it

• Your loop stays yours. No base class, no trainer object, no callback system. Deleting the import leaves a working script.
• Your hyperparameters are never touched silently. lr, loss, schedule, and optimizer are yours unless you call auto() or train(). optimize=True only does throughput — TF32, cudnn.benchmark, channels_last, workers, AMP, fused optimizers — and prints every change it makes.
• It's faster. +22% end to end vs torch defaults on a 4-conv net, +123% on optimizer-bound MLPs. Measurements.
• No launcher to configure. python train.py uses every GPU; under SLURM the auto-spawn stands down so srun autotrainer run behaves.
• It tells you why the run is bad. MFU and samples/sec, whether the loader is starving the GPU, and NaN / divergence / fp16-overflow named in plain language with the fix.
• One API across frameworks. PyTorch, TensorFlow/Keras, scikit-learn, XGBoost, LightGBM — most alternatives are PyTorch-only.

Worker counts come from the CPUs you were actually allocated (SLURM_CPUS_PER_TASK, then the affinity mask), so a job granted 8 CPUs of a 128-core node does not spawn 128 workers.

Documentation

Start here

• Entry points — which function do I call for what.
• Performance — what the optimizations measurably buy.
• How it compares — vs Lightning, Accelerate, Ray, torchrun, and when to reach for something else.

Guide

• Getting throughput out of your GPUs — what optimize=True replaces, and train_step().
• Training-loop helpers — set_epoch, zero_grad, eval_mode, accumulate, batch size, find_lr.
• Monitors — bottleneck, throughput, training health, the web UI, trackers, and framework callbacks.
• Scaling up — launching, torch.compile, FSDP, CPU offload, SLURM.
• One-line training — train(), auto(), and the data checks that run before the compute.

Reference

• API reference — generated from the source on every push to main.
• Examples — runnable scripts per framework, SLURM .sbatch templates.
• Environment variables — every knob autotrainer reads.
• Roadmap · CHANGELOG · Contributing · Security · Code of Conduct

HPC tip: on shared network filesystems (NFS/Lustre), route the uv cache to local node storage (export UV_CACHE_DIR=/tmp/$USER-uv-cache) and pass --link-mode=copy, or create .venv in /tmp.