TrainTools MCP server

traintools

<!-- mcp-name: io.github.AparajeetS/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 "" --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

Shortened. The full README is on GitHub.