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.

For the live browser dashboard:

pip install traceml-ai

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

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()

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

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.

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.

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

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.

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 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.

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.

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

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

FSDP support means timing collection and rank-skew surfacing; explicit FSDP collective attribution is on the roadmap.

On the roadmap:

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

• 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

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.

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.

Apache 2.0. See LICENSE.

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