Skip to content

TraceML

Diagnose slow PyTorch training with zero-code instrumentation. Catch regressions in CI.

TraceML shows where each training step goes—input loading, data transfer, forward, backward, and optimizer work—then identifies the bottleneck and saves evidence you can compare locally or check in CI.

Works automatically with: Hugging Face Trainer · PyTorch Lightning · RF-DETR

pip install traceml-ai
traceml run train.py

For these standard trainers, no TraceML code is required in the training script. Plain PyTorch loops and other stacks use a small explicit integration.

A completed run ends with a diagnosis, its evidence, and the next place to investigate:

Verdict: INPUT-BOUND
Why: Input Wait took 64% of Step Time.
Next: Increase workers, prefetch, or storage throughput.

  • Quickstart

    Run a supported trainer without changing its script, or instrument a custom PyTorch loop.

    Get started

  • Reading output

    Understand the terminal card, the dashboard, and the fields in final_summary.json.

    Read the guide

  • Integrations

    Choose the supported setup for Hugging Face, Lightning, RF-DETR, Accelerate, MONAI, Ray, DeepSpeed, W&B, or MLflow.

    See integrations

  • Compare runs

    See what changed between two runs and optionally fail CI on a Step Time regression.

    Compare runs

  • Public API

    The full traceml_ai reference: init(), trace_step(), and summary().

    API reference

  • Developer guide

    Architecture, pipeline contracts, and how to contribute to TraceML.

    Contribute

TraceML is open source under the Apache 2.0 license. Learn more at traceopt.ai or star the project on GitHub.