It helps you answer:

• Is the GPU computing or waiting for the input pipeline?
• Which phase is making each training step slower?
• Is one distributed rank holding back the others?
• Is memory usage silently growing?
• Did a code, data, or infrastructure change make the run slower?

TraceML produces actionable diagnostics with under 1% overhead in current benchmarks.

Quickstart

1. Install TraceML

For the live browser dashboard:

pip install traceml-ai

Or install with uv:

uv pip install traceml-ai

For a uv-managed project, add TraceML to the project's dependencies instead:

uv add traceml-ai

Extras use the same syntax with pip and uv. For example:

pip install "traceml-ai[dashboard]"
uv add "traceml-ai[dashboard]"
uv add "traceml-ai[torch]"

To try TraceML immediately in an environment managed by uv, install the torch extra required by the example:

uv pip install "traceml-ai[torch]"
traceml run examples/quickstart.py

Using Hugging Face Trainer, PyTorch Lightning, Ray Train, W&B, or MLflow? Start with the native integration path in Use With Your Stack.

2. Instrument the training step

Add TraceML around the core training step. You do not need to change your model, optimizer, loss function, or dataloader.

import traceml_ai as traceml

traceml.init(mode="auto")

for batch in dataloader:
    with traceml.trace_step(model):
        optimizer.zero_grad(set_to_none=True)
        outputs = model(batch["x"])
        loss = criterion(outputs, batch["y"])
        loss.backward()
        optimizer.step()

3. Run your training

Start your training with the live browser dashboard:

traceml run train.py

TraceML prints the dashboard URL, usually http://127.0.0.1:8765. Open it to see live bottleneck diagnostics while the job runs.

Or try the self-contained example first:

traceml run examples/quickstart.py

Try a real GPU demo in Colab

• Any PyTorch loop: data-loading bottleneck before/after
• Hugging Face Trainer: data-loading bottleneck before/after

traceml run train.py

ssh -L 8765:127.0.0.1:8765 user@remote-host

If you want a live view without a browser or SSH tunnel, use terminal mode:

traceml run train.py --mode=cli

Use summary mode when no live display is needed, such as headless jobs, CI, DDP, FSDP, Slurm, or multi-node runs:

traceml run train.py --mode=summary

For DDP, FSDP, Slurm, and multi-node runs, see Distributed Training.

Or launch your script directly

Prefer to launch training yourself with python or torchrun? Start the TraceML aggregator once with traceml serve, then run your script directly. traceml.init(...) connects to the aggregator over TCP:

# terminal 1: start the TraceML aggregator
traceml serve --aggregator-host 127.0.0.1 --aggregator-port 29765

# terminal 2: run your script directly
python train.py

For torchrun and multi-node, bind the aggregator so workers on other nodes can reach it:

# on the aggregator node (--nnodes x --nproc-per-node = total workers, so the
# aggregator waits for every rank before finalizing)
traceml serve --aggregator-bind-host 0.0.0.0 --aggregator-host <node0-ip> \
  --aggregator-port 29765 --nnodes <N> --nproc-per-node <M>

# on each training node: point workers at node 0's aggregator, then launch
TRACEML_AGGREGATOR_HOST=<node0-ip> TRACEML_AGGREGATOR_PORT=29765 \
  torchrun ... train.py

traceml.init(...) takes runtime settings as arguments (for example traceml.init(mode="auto", logs_dir="logs", aggregator_port=29765)), and falls back to TRACEML_* environment variables and traceml.yaml. The aggregator endpoint is configured with traceml serve flags. See Public API.

Try it in Docker

Build the image from the repository checkout, then run the default CPU demo:

docker build -t traceml-demo .
docker run --rm traceml-demo

The image installs TraceML from the checkout with PyTorch support and runs the slow DataLoader diagnosis demo in summary mode. The final report should identify the run as input-bound; no external dataset or GPU is required.

If the host has the NVIDIA Container Toolkit configured and the installed PyTorch build supports CUDA, the same demo can optionally access the GPU:

docker run --rm --gpus all traceml-demo

Example Diagnosis

Instead of showing only utilization charts, TraceML explains what is slowing the job, presents the supporting evidence, and tells you where to investigate.

INPUT STRAGGLER / CRITICAL

Rank 0 spent 254.5 ms waiting for input versus 3.8 ms on the median rank.

Next: inspect the dataloader, preprocessing, collate_fn, storage, and worker configuration on rank 0.

+----------------------------------------------------------------------------+
|  TraceML Run Summary | duration 40.1s                                      |
+----------------------------------------------------------------------------+
|                                                                            |
|  TraceML Verdict: INPUT STRAGGLER / CRITICAL                               |
|  Why: Rank r0 input wait was 254.5ms vs median rank r1 at 3.8ms.           |
|  Next: Inspect dataloader, collate_fn, preprocessing, and storage on the   |
|  slow rank.                                                                |
|                                                                            |
|  Section Status                                                            |
|  Section       Status                  Severity                            |
|  ------------------------------------------------                          |
|  Step Time     INPUT STRAGGLER         CRITICAL                            |
|  System        LOW GPU UTIL            INFO                                |
|  Process       NORMAL                  INFO                                |
|  Step Memory   BALANCED                INFO                                |
|                                                                            |
|  System Evidence                                                           |
|  Metric          Median        Worst         Skew        Scope             |
|  --------------------------------------------------------------------------|
|  CPU Util        18.4%         71.2%         52.8pp      node=n1           |
|  GPU Util        14.0%         0.0%          14.0pp      node=n0           |
|  GPU Memory      6.20GB        8.90GB        43.5%       node=n1           |
|  GPU Temp        42C           58C           16C         node=n1           |
|                                                                            |
|  Step Time Evidence                                                        |
|  Phase           Median        Worst         Skew        Scope             |
|  --------------------------------------------------------------------------|
|  Total           303.7ms       304.1ms       0.1%        rank=r0 node=n0   |
|  Input Wait      3.8ms         254.5ms       6597.4%     rank=r0 node=n0   |
|  Compute         259.5ms       261.0ms       0.6%        rank=r2 node=n1   |
+----------------------------------------------------------------------------+

In this example, rank 0 is the slow input rank and can hold back the aligned distributed step.

Want to reproduce a specific bottleneck? See examples/ for self-contained demos covering dataloader bottlenecks, H2D timing, DDP rank stragglers, Lightning, Hugging Face, Ray, and tracker-friendly summary logging.

What TraceML Helps You Triage

Use TraceML as the first check before opening a heavier profiler. It surfaces the likely bottleneck category so you know where to look next.

Display Modes

Choose the interface that fits the environment without changing the saved end-of-run artifacts.

Headless, CI, or capturing stdout? Use --mode=summary. TraceML still writes the same .json and .txt artifacts at the end of the run.

Saved Run Artifacts

TraceML writes two end-of-run artifacts:

logs/<run_name>/final_summary.json
logs/<run_name>/final_summary.txt

Reprint a saved summary without rerunning training:

traceml view logs/<run_name>/final_summary.json

Create a self-contained HTML report during the run:

traceml run train.py --html-report

Or render one later from a saved summary:

traceml view logs/<run_name>/final_summary.json --html

For experiment trackers, call traceml.summary() near the end of your script to get a flat dictionary of diagnosis statuses and average metrics. Keep final_summary.json when you want the complete run artifact or an input for traceml compare.

Compare Runs and Catch Regressions

Compare a slow run against a known-good baseline:

traceml compare input_slow/final_summary.json input_fixed/final_summary.json

+--------------------------------------------------------------------------------------+
|  TraceML Compare                                                                     |
+--------------------------------------------------------------------------------------+
|  Verdict: IMPROVEMENT                                                                |
|  Why: Step time decreased by 95.6%.                                                  |
|                                                                                      |
|  Metric                         A                B                Delta              |
|  Total step                     294.0 ms         13.0 ms          -280.9 ms (-95.6%) |
|  Input                          66.4 ms          2.7 ms           -63.7 ms (-95.9%)  |
+--------------------------------------------------------------------------------------+

See Compare Runs for the full report format.

Use With Your Stack

TraceML supports:

• Custom PyTorch training loops
• Hugging Face Trainer
• PyTorch Lightning
• Ray Train
• W&B and MLflow summary logging
• DDP and FSDP
• Slurm and multi-node summary reports

See Use With Your Stack for integration examples.

Where TraceML Fits

Start with TraceML to identify the bottleneck category. Open a deeper profiler when you need operator- or kernel-level detail.

Current Support

Works today:

• Single-GPU training
• Single-node multi-GPU DDP and FSDP
• Multi-node DDP summary reports
• Multi-node runs on Slurm
• Run-to-run comparison from final_summary.json
• Custom PyTorch loops, Hugging Face, PyTorch Lightning, and Ray Train

Validated on single-process and DDP. Runs on FSDP; under-reports due to collective masking.

On the roadmap:

• Multi-node live CLI and browser dashboard
• Explicit collective and NCCL timing

Troubleshooting Guides

• Find why PyTorch training is slow
• Find input pipeline bottlenecks
• Debug low GPU utilization
• Debug DDP rank stragglers
• Find PyTorch memory creep
• Distributed training
• Running on Slurm
• Use With Your Stack
• Compare Runs
• How to Read Output
• FAQ

Feedback

For bugs, unexpected results, or feature requests, open a GitHub issue using the matching issue template.

The templates ask for the information needed to reproduce training-environment problems, including hardware, topology, launch command, TraceML version, PyTorch and CUDA versions, and redacted summary output.

• Open a GitHub issue
• Security policy
• support@traceopt.ai

If TraceML helped you find a real bottleneck, use the I found a bottleneck issue template. These reports help other training teams recognize similar problems.

Contributing

Contributions are welcome, especially:

• Real slowdown examples and reproductions
• Distributed training edge cases
• Documentation improvements
• Framework integrations

See CONTRIBUTING.md for development setup and contribution guidelines.

License

Apache 2.0. See LICENSE.

TraceOpt is a trademark of OptAI UG (haftungsbeschränkt).