Autotrainer

Automatic distributed training and optimization for ML models. Give it a model and data - it detects your hardware (local GPUs or a SLURM cluster), picks the distribution strategy, and can infer the loss function, optimizer, learning rate, and schedule, or search for the best training hyperparameters.

Supports PyTorch (DDP, SLURM multi-node), TensorFlow/Keras (Mirrored / MultiWorker strategies), scikit-learn, XGBoost, and LightGBM - all through one API.

import autotrainer

# Smart defaults: loss, optimizer, LR, and schedule inferred and printed
model, loader, opt, loss_fn, sched = autotrainer.auto(model, loader)

# Or search for the best training recipe for YOUR model
best_model, best_params, study = autotrainer.tune(model, train_loader, val_loader)

# tune() also handles sklearn / XGBoost / LightGBM (curated default spaces)
best_est, params, study = autotrainer.tune(XGBClassifier(), (X, y), (X_val, y_val))

# Or fully hands-free: search the recipe, then train the winner to completion
model, params, study = autotrainer.fit(model, train_loader, val_loader)

python train.py                  # local: prepare() auto-distributes across GPUs
autotrainer run train.py         # equivalent, explicit (same spawn machinery)
srun autotrainer run train.py    # SLURM multi-node (srun starts the tasks)
autotrainer doctor               # diagnose your environment first

Install

From PyPI:

uv pip install autotrainer[torch]            # just PyTorch (recommended)
uv pip install autotrainer[torch,boosting]   # PyTorch + XGBoost/LightGBM

# Or with standard pip:
pip install autotrainer[torch]

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

Install (dev mode)

Using uv (recommended):

uv venv
# On Windows:
.venv\Scripts\activate
# On macOS/Linux:
source .venv/bin/activate

uv pip install -e ".[dev,torch,sklearn,tf,boosting,tune]"

Or using standard pip:

pip install -e ".[dev,torch,sklearn,tf,boosting,tune]"

Use

In your training script, add one line (plus set_epoch at each epoch start, so distributed shuffling gives every epoch a fresh order):

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

for epoch in range(epochs):
    autotrainer.set_epoch(loader, epoch)  # no-op when not distributed
    # ... your normal training loop

Want throughput, not magic? prepare(..., optimize=True) (the default)

prepare() makes your model/loader distribution-ready and — by default on a GPU — also flips on the wins users forget: TF32, cudnn.benchmark for CNNs, sane num_workers / pin_memory / persistent_workers defaults on bare loaders, and AMP — without touching your lr, loss, schedule, or optimizer choice. It's a no-op on CPU (every flag gates on a visible CUDA device), so CPU-only scripts see no change. Pass optimize=False to opt out.

Before: the boilerplate you write today to "use your GPUs well"

None of this touches your recipe (lr, loss, schedule), yet you have to remember all of it every time — and forgetting any one of them silently leaves 2–3× on the table.

import torch

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)   # your hyperparameter
loss_fn = nn.CrossEntropyLoss()                             # your hyperparameter

# Manual GPU optimization — ~10 lines of boilerplate:
torch.backends.cuda.matmul.allow_tf32 = True               # ships off by default
torch.backends.cudnn.allow_tf32 = True
torch.backends.cudnn.benchmark = True                       # free win for CNNs
loader = DataLoader(ds, batch_size=64, shuffle=True,
                    num_workers=8, pin_memory=True,         # avoid GPU starvation
                    persistent_workers=True)
amp_dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
scaler = torch.cuda.amp.GradScaler(enabled=amp_dtype == torch.float16)

for epoch in range(epochs):
    for xb, yb in loader:
        with torch.cuda.amp.autocast(dtype=amp_dtype):
            loss = loss_fn(model(xb), yb)
        scaler.scale(loss).backward()
        scaler.step(optimizer); scaler.update()

After: one line

import autotrainer

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)   # your hyperparameter
loss_fn = nn.CrossEntropyLoss()                             # your hyperparameter

# ONE line detects hardware and sets TF32 / cudnn.benchmark / num_workers /
# pin_memory / persistent_workers / AMP. lr, loss, schedule, optimizer untouched.
# (optimize=True is the default on GPU; shown explicitly here for clarity.)
model, loader, optimizer = autotrainer.prepare(model, loader, optimizer)

scaler = autotrainer.GradScaler()   # no-op when bf16 is available
for epoch in range(epochs):
    autotrainer.set_epoch(loader, epoch)                    # reshuffles in DDP
    for xb, yb in loader:
        with autotrainer.autocast_context():                # bf16 if supported, else fp16
            loss = loss_fn(model(xb), yb)
        scaler.scale(loss).backward()
        scaler.step(optimizer); scaler.update()

What changed: one prepare(...) call (optimize is on by default) + two no-op-on-CPU helpers. What didn't: your lr, your loss, your schedule, your optimizer. Same script runs unchanged on a laptop (everything degrades to no-ops) and on an A100.

Even simpler: train_step() runs the whole step

prepare() can't wrap your loop, so the AMP block above is the one thing you still hand-write. train_step() does it for you - forward under autocast, loss, scale/backward/step/update, then zero the grads - in one call, returning the (detached) loss for logging:

model, loader, optimizer = autotrainer.prepare(model, loader, optimizer)
scaler = autotrainer.GradScaler()          # no-op on CPU / bf16; omit to skip

for epoch in range(epochs):
    autotrainer.set_epoch(loader, epoch)
    model.train()
    for xb, yb in loader:
        loss = autotrainer.train_step(model, loss_fn, xb, yb, optimizer, scaler=scaler)

Same contract: lr / loss / schedule / optimizer are yours; train_step only runs the forgettable, order-sensitive bookkeeping (backward stays outside autocast). Pass autocast=False to keep full precision, or omit scaler on CPU / bf16 GPUs.

What it prints when it runs

Nothing is silent — every speedup is named, the user is explicitly told their hyperparameters weren't touched, and when AMP is on prepare() points at train_step() (and still shows the manual form) so the loop is one call either way (the helpers are no-ops on CPU, so the snippet is safe verbatim):

[autotrainer] mode=local_multi_gpu nodes=1 procs/node=4 world_size=4
[autotrainer] DistributedSampler installed (shuffle=True) - call autotrainer.set_epoch(loader, epoch) ...
[autotrainer] optimize: TF32, cudnn.benchmark, num_workers=8, pin_memory, persistent_workers, AMP (hyperparameters untouched)
[autotrainer] optimize: AMP is on. Simplest - one call per step:
    scaler = autotrainer.GradScaler()   # once, before the loop
    loss = autotrainer.train_step(model, loss_fn, xb, yb, opt, scaler=scaler)
  or wrap the step yourself:
    with autotrainer.autocast_context():
        out = model(x); loss = loss_fn(out, y)
    scaler.scale(loss).backward(); scaler.step(opt); scaler.update()

Training-loop helpers (zero_grad, eval_mode, accumulate)

The small things users forget inside the loop. None touch lr / loss / schedule / optimizer choice.

import autotrainer

for epoch in range(epochs):
    autotrainer.set_epoch(loader, epoch)
    model.train()
    for xb, yb in loader:
        with autotrainer.autocast_context():
            loss = loss_fn(model(xb), yb)
        scaler.scale(loss).backward()
        scaler.step(opt)
        scaler.update()
        autotrainer.zero_grad(opt)            # set_to_none=True, saves memory

    # eval_mode restores the prior train/eval state - kills the classic
    # "forgot to flip back to train() after validation" silent bug.
    with autotrainer.eval_mode(model):
        val_loss = evaluate(model, val_loader)

Gradient accumulation when the effective batch is larger than the physical one — scales the step count, not the lr:

# Effective batch = 4 micro-batches; opt steps once per 4 backwards.
with autotrainer.accumulate(opt, steps=4, scaler=scaler) as acc:
    for micro_xb, micro_yb in micro_batches:
        with autotrainer.autocast_context():
            loss = loss_fn(model(micro_xb), micro_yb) / 4
        acc.backward(loss)

Auto batch size

# Grow batch size until OOM, back off one step. Pass loss_fn for an
# accurate fwd+bwd measurement; without it the sweep is forward-only
# (conservative). lr and schedule are NOT changed - pair with accumulate()
# to scale the step to the new effective batch.
model, loader, opt = autotrainer.prepare(
    model, loader, opt, optimize=True, auto_bs=True, loss_fn=loss_fn
)

Dataloader bottleneck monitor

The cheapest piece of the training-triage family. Sample per-step data-load vs compute time and get a plain-language warning when the loader is starving the GPU:

mon = autotrainer.BottleneckMonitor(warmup=10)
for xb, yb in loader:
    with mon.data_time():
        pass  # the wait for the next batch
    with mon.step_time():
        loss = loss_fn(model(xb), yb); loss.backward(); opt.step()
    mon.tick()
    if mon.should_report():
        mon.report()   # -> "[autotrainer] bottleneck: dataloader is 78% of
                       #     step time ... - raise num_workers / pin_memory / prefetch"

Opt-in; zero overhead when not constructed.

Training health monitor (triage)

BottleneckMonitor and ThroughputMonitor answer "is it fast?"; TrainingMonitor answers "is it healthy?" — the silent numerical failures. Call step() once per optimizer step (after backward(), before zero_grad()) and it flags each problem once, in plain language, with the fix:

mon = autotrainer.TrainingMonitor()
for xb, yb in loader:
    loss = loss_fn(model(xb), yb)
    loss.backward()
    mon.step(loss, model=model, optimizer=opt, scaler=scaler)
    opt.step(); opt.zero_grad()
mon.report()   # one-line all-clear, or a recap of what fired

It catches NaN/Inf loss (lr too high, or bad inputs), loss divergence, fp16 gradient overflow (→ "switch to bf16"), and non-finite / spiking / vanishing gradients — e.g.:

[autotrainer] triage: gradient norm spiked to 4.1e+03 (128x the recent median)
    - add gradient clipping (torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0))

Everything but loss is optional (pass model for gradient checks, optimizer for a concrete lr hint, scaler for fp16 overflow). Like the other monitors it only observes — it never touches your lr, loss, schedule, or optimizer. Opt-in; zero overhead when not constructed.

Then launch. prepare() auto-distributes — on a multi-GPU box it spawns one worker per GPU the first time it's called, so a bare python train.py uses all visible GPUs with no launcher:

python train.py                   # local: prepare() auto-spawns per-GPU workers
autotrainer run train.py          # equivalent, explicit (same spawn machinery)
autotrainer info                  # show what was detected

Set prepare(..., auto_launch=False) to opt out (e.g. you're managing processes yourself, or running under your own launcher). The auto-spawn only fires when all three hold: no RANK env var set (fresh parent, not a worker), not under SLURM, and ≥2 GPUs on one node — so it never double-spawns under srun or loops on already-launched workers.

On SLURM (multi-node), the front door stays srun autotrainer run — SLURM itself starts one task per GPU across nodes, so auto-spawn is correctly skipped there:

On SLURM, inside your sbatch script:

#SBATCH --nodes=2
#SBATCH --ntasks-per-node=4
#SBATCH --gres=gpu:4
srun autotrainer run train.py

torch.compile and FSDP

Two more opt-ins on prepare(), both preserving the "don't touch hyperparameters" contract:

# torch.compile: 1.5-2x on modern torch+GPU for many models. Compiled
# BEFORE the DDP wrap (the documented-supported order; the reverse causes
# graph breaks). On compile failure (dynamic shapes, missing Triton on
# Windows) falls back to the uncompiled model with a warning.
model, loader, opt = autotrainer.prepare(
    model, loader, opt, optimize=True, compile=True
)

# FSDP instead of DDP when the model is too large to replicate on every GPU.
# Shards params/grads/optimizer state across ranks. use_orig_params=True so
# your existing optimizer keeps working unchanged.
model, loader, opt = autotrainer.prepare(
    model, loader, opt, optimize=True, fsdp=True
)

What it prints:

[autotrainer] optimize: TF32, torch.compile(mode=default), FSDP (hyperparameters untouched)

compile_mode: default | reduce-overhead (CUDA graphs, fastest for small models with static shapes) | max-autotune (kernel search, slow first compile). None of these flags touch lr / loss / schedule / optimizer.

CPU offload and SLURM scratch

CPU offload pairs with FSDP for the case where the model OOMs even when sharded across ranks — move params to CPU, bring them to GPU only for the forward/backward:

model, loader, opt = autotrainer.prepare(
    model, loader, opt, optimize=True, fsdp=True, cpu_offload=True
)

Trades throughput for memory headroom. Ignored with a warning on the DDP path or single-process (no FSDP = nothing to offload).

SLURM node-local scratch — the classic HPC footgun is every rank writing to $HOME (NFS, slow, shared) instead of $TMPDIR (node-local, fast, auto-cleaned). One call at the top of your script wires the obvious env vars and warns when the scratch looks networked:

import autotrainer
autotrainer.configure_scratch()  # sets TORCHINDUCTOR_CACHE_DIR to $TMPDIR
# ... your training script ...

Under SLURM this uses $TMPDIR/autotrainer-<jobid>; outside SLURM, the system temp dir. node_scratch() returns the path if you want to write your own checkpoints there too.

One-call training: fit()

fit() is the whole pipeline in one call - give it a model and data, get back the best model it can produce on your hardware:

model, params, study = autotrainer.fit(model, train_loader, val_loader, trials=30)

1. Tune: an ASHA (successive-halving) search over the training recipe - lr, weight decay, optimizer, batch size, LR schedule + warmup, gradient clipping, and (for classification) label smoothing - on short trials, with the default space chosen from your model and its inferred loss. Most candidates get a small budget and only survivors are promoted, so the wide space stays cheap. Launched distributed, the trials are split across all ranks via a shared journal-file study - one trial per process, every GPU busy during the search.
2. Train: the winning recipe is retrained from your model's original init through prepare() - which auto-distributes it across every GPU/node (a bare python train.py spawns one worker per GPU locally; under SLURM, srun autotrainer run does the same across nodes) - with the winning schedule, mixed precision, and early stopping on the val loss. The best epoch's weights are returned.

Pass test_loader= to also get an honest held-out score the search never saw (printed and stored on study.user_attrs["test_loss"]) - the number to trust once a wide search has been optimizing against your val set.

Pass checkpoint="fit.ckpt" to make it preemption-safe: the full training state is saved every epoch, and rerunning the same script resumes where it died (skipping the search) - ideal for requeued SLURM jobs.

For long searches under multi-process launches, raise the collective timeout with AUTOTRAINER_TIMEOUT (seconds) - see .env.example.

Optional: auto batch size

best = autotrainer.find_batch_size(model, my_one_step_fn)

Documentation

• CHANGELOG - version history.
• API reference - published from CI on every push to main; build locally with pdoc -o docs/build src/autotrainer.
• Public API & deprecation policy - what autotrainer.__all__ exports is stable; the rest is internal.
• Examples - runnable scripts for each framework and SLURM .sbatch templates.
• Environment variables - every knob autotrainer reads.
• Contributing | Security policy | Code of Conduct.

Roadmap

Toward 1.0:

• Stabilization: the public API is frozen as of 0.10; 1.0 now only blocks on real multi-node SLURM validation (the deprecated train_loader=/val_loader= aliases have been removed).

Understanding your training run (the theme after 1.0 - autotrainer should explain runs, not just launch them):

• Preflight estimation (doctor --profile): dry-run a few batches, then report projected training time, memory headroom, and cost per GPU count - answer "how many GPUs do I actually need?" before burning an allocation.
• Data sanity checks in auto(): the same one-batch peek that infers the loss can flag class imbalance (suggest weighted loss), unnormalized inputs, and train/val overlap.
• Training cards: every fit() emits a reproducibility card (recipe, seeds, environment, val curve) and replay reruns it.

Deeper SLURM ergonomics:

• autotrainer sbatch train.py --nodes 2 --time 4h: generate and submit a correct sbatch script (no more --ntasks-per-node != GPUs footguns).
• Preemption handling: catch SLURM's requeue signal, checkpoint via fit()'s resume support, and continue after requeue automatically.

More breadth:

• Multi-node boosting (xgboost.dask / lightgbm.dask across a SLURM allocation) — currently single-node threads only.
• Augmentation and architecture-aware search: the recipe search now covers schedule/warmup/grad-clip/label-smoothing; data-augmentation policies and width/depth are the natural next frontier (a bigger, opt-in commitment that would step beyond "the model is yours").

See CHANGELOG.md for the full version history, and open or upvote issues to prioritize these.