{
  "markdown": "# traintools\n\n<!-- mcp-name: io.github.AparajeetS/traintools -->\n\n[![CI](https://github.com/AparajeetS/Traintools/actions/workflows/ci.yml/badge.svg)](https://github.com/AparajeetS/Traintools/actions/workflows/ci.yml)\n[![PyPI](https://img.shields.io/pypi/v/traintools.svg)](https://pypi.org/project/traintools/)\n[![Python](https://img.shields.io/pypi/pyversions/traintools.svg)](https://pypi.org/project/traintools/)\n[![License: MIT](https://img.shields.io/badge/License-MIT-blue.svg)](LICENSE)\n\nPaper-backed ML training diagnostics for PyTorch. Small tools that answer\npractical questions while a run is still alive.\n\n[Source on GitHub](https://github.com/AparajeetS/Traintools) | [PyPI](https://pypi.org/project/traintools/)\n\nHave a real run or false alarm to share? Use the\n[diagnostic report template](https://github.com/AparajeetS/Traintools/issues/new?template=diagnostic-report.yml).\n\n```bash\npip install traintools\n```\n\nNot sure which tool fits the problem?\n\n```bash\ntraintools recommend \"my loss became NaN and gradients explode\"\ntraintools integration gradient-health --framework pytorch\n```\n\nAgents can use the JSON CLI, [the agent guide](https://github.com/AparajeetS/Traintools/blob/main/AGENTS.md), [llms.txt](https://github.com/AparajeetS/Traintools/blob/main/llms.txt),\nor the optional local MCP server:\n\n```bash\npip install \"traintools[mcp]\"\ntraintools-mcp\n```\n\nThe same local, read-only server is discoverable through the official MCP\nRegistry and can be launched without a persistent install:\n\n```bash\nuvx --with mcp traintools mcp\n```\n\nFor AI coding assistants: start with `traintools recommend \"<training symptom>\"\n--json`, then insert the smallest diagnostic snippet into the user's existing\nloop. The agent-facing entrypoints are [AGENTS.md](https://github.com/AparajeetS/Traintools/blob/main/AGENTS.md),\n[llms.txt](https://github.com/AparajeetS/Traintools/blob/main/llms.txt), JSON\nschemas, and the optional MCP server.\n\n## Tools\n\n| Tool | Question it answers |\n|---|---|\n| **Gradient Noise Scale (GNS)** | Is my batch size wasting compute? |\n| **GradientAccumulationGNS** | Can I get GNS for free during gradient accumulation? |\n| **PlasticityProbe** | Is my network losing the ability to learn? |\n| **TrainGuard** | Should I stop training yet? |\n| **BatchInspector** | Is this batch broken, imbalanced, or out of scale? |\n| **GradientHealthMonitor** | Are gradients finite, clipped, vanished, exploded, or too large for the weights? |\n| **ExampleDynamicsTracker** | Which examples are forgotten, hard, ambiguous, or likely mislabeled? |\n| **GradientConfusionMonitor** | Are micro-batch gradients fighting each other and slowing SGD? |\n| **AUMTracker** | Which examples look mislabeled by margin dynamics? |\n| **EL2NTracker** | Which examples are important or pruneable early in training? |\n| **NeuralCollapseMonitor** | Has the classifier entered neural-collapse geometry? |\n\n## Quick Start\n\n```python\nfrom traintools import BatchInspector, GradientHealthMonitor\nfrom traintools.callbacks.pytorch import TraintoolsTracker\n\ntracker = TraintoolsTracker(model, loss_fn)\nbatch_inspector = BatchInspector(expected_num_classes=10)\ngrad_health = GradientHealthMonitor(max_grad_norm=1.0)\n\nfor step, (x, y) in enumerate(dataloader):\n    batch_report = batch_inspector.inspect(x, y, step=step)\n    if not batch_report.ok:\n        print(batch_report)\n\n    loss = loss_fn(model(x), y)\n    loss.backward()\n\n    grad_report = grad_health.inspect(model, step=step, lr=optimizer.param_groups[0][\"lr\"])\n    if not grad_report.ok:\n        print(grad_report)\n\n    optimizer.step()\n    optimizer.zero_grad()\n\n    decision = tracker.step(step=step, inputs=x, targets=y, val_loss=val_loss)\n    if decision and decision.should_stop:\n        break\n```\n\nHuggingFace Trainer:\n\n```python\nfrom transformers import Trainer\nfrom traintools.callbacks.huggingface import TraintoolsCallback\n\ntrainer = Trainer(model=model, ..., callbacks=[TraintoolsCallback()])\n```\n\n## Gradient Noise Scale\n\nGNS is the ratio of per-example gradient variance to gradient signal:\n\n```text\nGNS = tr(Sigma) / ||G||^2\n```\n\nIt estimates the critical batch size B*: the point where larger batches stop\nbuying much more optimization progress.\n\n- `GNS > B`: under-batched, larger batches can help\n- `GNS < B`: over-batched, the batch may be larger than needed\n- `GNS ~= B`: near the efficient frontier\n\n`traintools` uses the unbiased estimators from McCandlish et al. 2018\n(Bessel-corrected variance, bias-corrected signal) and tracks GNS as\n`EMA(tr(Sigma)) / EMA(||G||^2)`.\n\n```text\n[step 500] GNS=5010.7 (EMA)  critical_batch=5011  current=64  regime=under-batched\n  > Batch size 64 is ~78x below the critical batch (~5011). Larger batches would give cleaner gradients per step.\n```\n\n### Free GNS During Gradient Accumulation\n\nIf you already use gradient accumulation, the per-micro-batch gradients you\ncompute anyway are exactly the samples GNS needs.\n\n```python\nfrom traintools import GradientAccumulationGNS\n\ngns = GradientAccumulationGNS(model, micro_batch_size=B_micro)\n\nfor step in range(num_steps):\n    for xm, ym in micro_batches:\n        (loss_fn(model(xm), ym) / accum_steps).backward()\n        gns.record_microbatch()\n    optimizer.step()\n    result = gns.compute(step=step)\n    optimizer.zero_grad()\n    gns.reset_accumulation()\n```\n\n## PlasticityProbe\n\nPlasticityProbe measures activations directly:\n\n- dormant unit fraction: units whose activation is near zero for every input\n- feature effective rank: normalized effective rank of the activation covariance\n\nThose are combined into a plasticity score in `[0, 1]`.\n\n```text\n[step 200] Plasticity Score: 0.706\n  All layers healthy.\n```\n\n## TrainGuard\n\nTrainGuard fits a power-law or exponential curve to validation loss, bootstraps\nuncertainty, and only stops when continuing looks unlikely to matter.\n\n```text\n[step 400] STOP\n  current loss: 0.6536\n  predicted final: 0.6119\n  expected improvement: 0.0417 (90% CI: [0.0012, 0.0821])\n  estimated plateau at step: 3200\n  reason: No improvement in 300 steps (best=0.6350 at step 93).\n```\n\n## BatchInspector\n\nBatchInspector catches bad tensors and labels before they quietly poison a run.\n\n```python\nfrom traintools import BatchInspector\n\ninspector = BatchInspector(expected_num_classes=10, max_abs_value=1e4)\nreport = inspector.inspect(inputs=x, targets=y, step=step)\nif not report.ok:\n    print(report)\n```\n\nIt checks for empty tensors, NaNs/infs, extreme scales, constant tensors, labels\noutside the expected class range, and severe batch imbalance.\n\n## GradientHealthMonitor\n\nGradientHealthMonitor is called after `backward()` and before `optimizer.step()`.\n\n```python\nfrom traintools import GradientHealthMonitor\n\nmonitor = GradientHealthMonitor(max_grad_norm=1.0)\nloss.backward()\nreport = monitor.inspect(model, step=step, lr=optimizer.param_groups[0][\"lr\"])\nif not report.ok:\n    print(report)\n```\n\nIt reports global and per-layer gradient norms, non-finite gradients, likely\nvanishing/exploding gradients, clipping coefficient, and update-to-weight ratio.\n\n## ExampleDynamicsTracker\n\nExampleDynamicsTracker implements two underused training-dynamics probes:\n\n- example forgetting events from Toneva et al. 2019\n- dataset-cartography-style confidence and variability from Swayamdipta et al. 2020\n\nUse stable dataset ids, logits, and labels during a normal classification run.\n\n```python\nfrom traintools import ExampleDynamicsTracker\n\ndynamics = ExampleDynamicsTracker()\n\nfor step, (ids, x, y) in enumerate(dataloader):\n    logits = model(x)\n    loss = loss_fn(logits, y)\n    dynamics.update(ids, logits, y, step=step)\n    loss.backward()\n    optimizer.step()\n    optimizer.zero_grad()\n\nprint(dynamics.summary())\nprint(\"likely noisy or brittle:\", [ex.example_id for ex in dynamics.most_forgotten(20)])\nprint(\"ambiguous:\", [ex.example_id for ex in dynamics.cartography_region(\"ambiguous\")])\n```\n\nA forgetting event is a transition from correct classification to incorrect\nclassification for the same example. Repeatedly forgotten examples are often\nambiguous, mislabeled, or distribution-edge cases. Unforgettable examples can be\nuseful candidates for pruning or curriculum experiments.\n\n## GradientConfusionMonitor\n\nGradientConfusionMonitor estimates whether micro-batch gradients are aligned or\nfighting each other, following the gradient-confusion idea from Sankararaman et\nal. 2019.\n\n```python\nfrom traintools import GradientConfusionMonitor\n\nconfusion = GradientConfusionMonitor(n_splits=4)\nreport = confusion.estimate(model, loss_fn, x, y, step=step)\nif not report.ok:\n    print(report)\n```\n\nIt reports mean/min/max pairwise gradient cosine, the fraction of negative\ngradient pairs, and a compact conflict score. High conflict can point to noisy\nlabels, incompatible samples, depth/initialization issues, or a need for a\ndifferent batching/curriculum strategy.\n\n## AUMTracker\n\nAUMTracker implements the Area Under the Margin statistic from Pleiss et al.\n2020. For each example, it averages:\n\n```text\ntrue_class_logit - max(other_class_logits)\n```\n\nLow-AUM examples are candidates for label audit or ambiguity review.\n\n```python\nfrom traintools import AUMTracker\n\naum = AUMTracker(low_aum_threshold=0.0)\nfor step, (ids, x, y) in enumerate(dataloader):\n    logits = model(x)\n    aum.update(ids, logits, y, step=step)\n\nprint([ex.example_id for ex in aum.lowest_aum(20)])\n```\n\n## EL2NTracker\n\nEL2NTracker implements the cheap example-importance score from Paul et al.\n2021:\n\n```text\n||softmax(logits) - one_hot(label)||_2\n```\n\nHigh EL2N examples tend to be important, hard, noisy, or distribution-edge\nexamples. Low EL2N examples can be candidates for data-pruning experiments.\n\n```python\nfrom traintools import EL2NTracker\n\nel2n = EL2NTracker()\nel2n.update(ids, logits, y, step=step)\nimportant = el2n.highest(100)\nprune_candidates = el2n.lowest(100)\n```\n\n## NeuralCollapseMonitor\n\nNeuralCollapseMonitor measures late-stage classifier geometry from Papyan, Han,\nand Donoho 2020:\n\n- NC1: within-class feature variability relative to between-class variability\n- NC2: class means approaching simplex ETF geometry\n- NC3: classifier weights aligning with class means, when weights are provided\n- NCC accuracy: nearest-class-center accuracy\n\n```python\nfrom traintools import NeuralCollapseMonitor\n\ncollapse = NeuralCollapseMonitor()\nreport = collapse.measure(features, labels, classifier_weight=model.fc.weight)\nprint(report)\n```\n\n## Installation\n\n```bash\n# Core: PyTorch only\npip install traintools\n\n# Curve fitting and plotting helpers\npip install traintools[full]\n\n# HuggingFace Trainer integration\npip install traintools[hf]\n\n# Local MCP server for compatible AI clients\npip install traintools[mcp]\n\n# Development\npip install -e \".[dev]\"\n```\n\n## Project Status\n\n`traintools` is alpha software. The diagnostics are intentionally small and\nwell-tested, but thresholds are heuristics and should be interpreted as training\nsignals, not automatic truth. Bug reports, benchmark traces, and real-world\nfailure cases are especially welcome.\n\nProblem-oriented guides live in [the documentation](https://github.com/AparajeetS/Traintools/tree/main/docs/problems). Diagnostic\nobjects can be written as versioned JSON with `write_json_report`.\n\n## References\n\n- McCandlish, Kaplan, Amodei et al. 2018. *An Empirical Model of Large-Batch Training*.\n- Dohare, Sutton et al. 2024. *Loss of plasticity in deep continual learning*.\n- Lyle et al. 2023. *Understanding Plasticity in Neural Networks*.\n- Domhan et al. 2015. *Speeding up automatic hyperparameter optimization of DNNs by extrapolation of learning curves*.\n- Toneva et al. 2019. *An Empirical Study of Example Forgetting during Deep Neural Network Learning*.\n- Swayamdipta et al. 2020. *Dataset Cartography: Mapping and Diagnosing Datasets with Training Dynamics*.\n- Sankararaman et al. 2019. *The Impact of Neural Network Overparameterization on Gradient Confusion and Stochastic Gradient Descent*.\n- Pleiss et al. 2020. *Identifying Mislabeled Data using the Area Under the Margin Ranking*.\n- Paul, Ganguli, and Dziugaite 2021. *Deep Learning on a Data Diet: Finding Important Examples Early in Training*.\n- Papyan, Han, and Donoho 2020. *Prevalence of Neural Collapse during the Terminal Phase of Deep Learning Training*.\n\n## License\n\nMIT\n",
  "bytes": 12270,
  "sha": "26c5c536d37b6b1fe23ca9dcf203ce554af1b3d80ddadce0c9b0963fdc12e3ef",
  "repo_slug": "aparajeets/traintools",
  "fonte": "repo",
  "truncated": false,
  "api": "https://agentalog.com/api/listings/mcp_io_github_aparajeets_traintools_818bfe60/readme"
}