traintools

Paper-backed ML training diagnostics for PyTorch. Small tools that answer practical questions while a run is still alive.

Source on GitHub | PyPI

Have a real run or false alarm to share? Use the diagnostic report template.

pip install traintools

Not sure which tool fits the problem?

traintools recommend "my loss became NaN and gradients explode"
traintools integration gradient-health --framework pytorch

Agents can use the JSON CLI, the agent guide, llms.txt, or the optional local MCP server:

pip install "traintools[mcp]"
traintools-mcp

The same local, read-only server is discoverable through the official MCP Registry and can be launched without a persistent install:

uvx --with mcp traintools mcp

For AI coding assistants: start with traintools recommend "<training symptom>" --json, then insert the smallest diagnostic snippet into the user's existing loop. The agent-facing entrypoints are AGENTS.md, llms.txt, JSON schemas, and the optional MCP server.

Tools

Quick Start

from traintools import BatchInspector, GradientHealthMonitor
from traintools.callbacks.pytorch import TraintoolsTracker

tracker = TraintoolsTracker(model, loss_fn)
batch_inspector = BatchInspector(expected_num_classes=10)
grad_health = GradientHealthMonitor(max_grad_norm=1.0)

for step, (x, y) in enumerate(dataloader):
    batch_report = batch_inspector.inspect(x, y, step=step)
    if not batch_report.ok:
        print(batch_report)

    loss = loss_fn(model(x), y)
    loss.backward()

    grad_report = grad_health.inspect(model, step=step, lr=optimizer.param_groups[0]["lr"])
    if not grad_report.ok:
        print(grad_report)

    optimizer.step()
    optimizer.zero_grad()

    decision = tracker.step(step=step, inputs=x, targets=y, val_loss=val_loss)
    if decision and decision.should_stop:
        break

HuggingFace Trainer:

from transformers import Trainer
from traintools.callbacks.huggingface import TraintoolsCallback

trainer = Trainer(model=model, ..., callbacks=[TraintoolsCallback()])

Gradient Noise Scale

GNS is the ratio of per-example gradient variance to gradient signal:

GNS = tr(Sigma) / ||G||^2

It estimates the critical batch size B*: the point where larger batches stop buying much more optimization progress.

• GNS > B: under-batched, larger batches can help
• GNS < B: over-batched, the batch may be larger than needed
• GNS ~= B: near the efficient frontier

traintools uses the unbiased estimators from McCandlish et al. 2018 (Bessel-corrected variance, bias-corrected signal) and tracks GNS as EMA(tr(Sigma)) / EMA(||G||^2).

[step 500] GNS=5010.7 (EMA)  critical_batch=5011  current=64  regime=under-batched
  > Batch size 64 is ~78x below the critical batch (~5011). Larger batches would give cleaner gradients per step.

Free GNS During Gradient Accumulation

If you already use gradient accumulation, the per-micro-batch gradients you compute anyway are exactly the samples GNS needs.

from traintools import GradientAccumulationGNS

gns = GradientAccumulationGNS(model, micro_batch_size=B_micro)

for step in range(num_steps):
    for xm, ym in micro_batches:
        (loss_fn(model(xm), ym) / accum_steps).backward()
        gns.record_microbatch()
    optimizer.step()
    result = gns.compute(step=step)
    optimizer.zero_grad()
    gns.reset_accumulation()

PlasticityProbe

PlasticityProbe measures activations directly:

• dormant unit fraction: units whose activation is near zero for every input
• feature effective rank: normalized effective rank of the activation covariance

Those are combined into a plasticity score in [0, 1].

[step 200] Plasticity Score: 0.706
  All layers healthy.

TrainGuard

TrainGuard fits a power-law or exponential curve to validation loss, bootstraps uncertainty, and only stops when continuing looks unlikely to matter.

[step 400] STOP
  current loss: 0.6536
  predicted final: 0.6119
  expected improvement: 0.0417 (90% CI: [0.0012, 0.0821])
  estimated plateau at step: 3200
  reason: No improvement in 300 steps (best=0.6350 at step 93).

BatchInspector

BatchInspector catches bad tensors and labels before they quietly poison a run.

from traintools import BatchInspector

inspector = BatchInspector(expected_num_classes=10, max_abs_value=1e4)
report = inspector.inspect(inputs=x, targets=y, step=step)
if not report.ok:
    print(report)

It checks for empty tensors, NaNs/infs, extreme scales, constant tensors, labels outside the expected class range, and severe batch imbalance.

GradientHealthMonitor

GradientHealthMonitor is called after backward() and before optimizer.step().

from traintools import GradientHealthMonitor

monitor = GradientHealthMonitor(max_grad_norm=1.0)
loss.backward()
report = monitor.inspect(model, step=step, lr=optimizer.param_groups[0]["lr"])
if not report.ok:
    print(report)

It reports global and per-layer gradient norms, non-finite gradients, likely vanishing/exploding gradients, clipping coefficient, and update-to-weight ratio.

ExampleDynamicsTracker

ExampleDynamicsTracker implements two underused training-dynamics probes:

• example forgetting events from Toneva et al. 2019
• dataset-cartography-style confidence and variability from Swayamdipta et al. 2020

Use stable dataset ids, logits, and labels during a normal classification run.

from traintools import ExampleDynamicsTracker

dynamics = ExampleDynamicsTracker()

for step, (ids, x, y) in enumerate(dataloader):
    logits = model(x)
    loss = loss_fn(logits, y)
    dynamics.update(ids, logits, y, step=step)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

print(dynamics.summary())
print("likely noisy or brittle:", [ex.example_id for ex in dynamics.most_forgotten(20)])
print("ambiguous:", [ex.example_id for ex in dynamics.cartography_region("ambiguous")])

A forgetting event is a transition from correct classification to incorrect classification for the same example. Repeatedly forgotten examples are often ambiguous, mislabeled, or distribution-edge cases. Unforgettable examples can be useful candidates for pruning or curriculum experiments.

GradientConfusionMonitor

GradientConfusionMonitor estimates whether micro-batch gradients are aligned or fighting each other, following the gradient-confusion idea from Sankararaman et al. 2019.

from traintools import GradientConfusionMonitor

confusion = GradientConfusionMonitor(n_splits=4)
report = confusion.estimate(model, loss_fn, x, y, step=step)
if not report.ok:
    print(report)

It reports mean/min/max pairwise gradient cosine, the fraction of negative gradient pairs, and a compact conflict score. High conflict can point to noisy labels, incompatible samples, depth/initialization issues, or a need for a different batching/curriculum strategy.

AUMTracker

AUMTracker implements the Area Under the Margin statistic from Pleiss et al. 2020. For each example, it averages:

true_class_logit - max(other_class_logits)

Low-AUM examples are candidates for label audit or ambiguity review.

from traintools import AUMTracker

aum = AUMTracker(low_aum_threshold=0.0)
for step, (ids, x, y) in enumerate(dataloader):
    logits = model(x)
    aum.update(ids, logits, y, step=step)

print([ex.example_id for ex in aum.lowest_aum(20)])

EL2NTracker

EL2NTracker implements the cheap example-importance score from Paul et al. 2021:

||softmax(logits) - one_hot(label)||_2

High EL2N examples tend to be important, hard, noisy, or distribution-edge examples. Low EL2N examples can be candidates for data-pruning experiments.

from traintools import EL2NTracker

el2n = EL2NTracker()
el2n.update(ids, logits, y, step=step)
important = el2n.highest(100)
prune_candidates = el2n.lowest(100)

NeuralCollapseMonitor

NeuralCollapseMonitor measures late-stage classifier geometry from Papyan, Han, and Donoho 2020: