TraceML
Diagnose slow PyTorch training. Catch regressions in CI.
Works with: PyTorch · Hugging Face Trainer · PyTorch Lightning · RF-DETR · Ray Train · Slurm (support details)
Quickstart • Performance checks • Try in Colab • Integrations • Documentation • GitHub Issues
TraceML is an open-source tool designed for lightweight, always-on diagnostics during PyTorch training. See how much time is spent waiting for the next batch, where step time goes, and whether a slow worker is holding up a distributed run. Compare completed runs to see what changed, and apply a Step Time threshold to the same comparison in CI.
At the end of a run, it gives you:
- A diagnosis: waiting for input, costly data transfers, compute-heavy steps, or a likely slow distributed worker.
- The evidence: step timings, CPU/GPU usage, memory trends, and per-worker comparisons.
- The next step: what part of your training setup to investigate first.
Example diagnosis
+----------------------------------------------------------------------------------------------------------------------------------------------------------+
| TraceML Run Summary |
| bert_finetune · 1 rank · 1 GPU observed · 256 common steps · 52.4s |
+----------------------------------------------------------------------------------------------------------------------------------------------------------+
| |
| Verdict: INPUT-BOUND (CRITICAL) |
| Why: Input Wait took 64% of Step Time. |
| Next: Increase workers, prefetch, or storage throughput. |
| |
| STEP TIMING (Window Average), GPU Clock || STEP MEMORY: BALANCED |
| Step Time 200.4 ms 100% || |
| ├─ Input Wait 128.0 ms 64% ◀ cause || |
| ├─ Compute 68.0 ms 34% || avg per-step peak avg |
| │ ├─ Forward 24.0 ms 12% || Allocated 2.9 GB |
| │ ├─ Backward 38.0 ms 19% || Reserved 3.2 GB |
| │ └─ Optimizer 6.0 ms 3% || |
| ├─ H2D 0.4 ms <1% || |
| └─ Residual 3.6 ms 2% || |
| DataLoader fetch: 120.0 ms (CPU, supplemental) || |
| |
| SYSTEM METRICS: LOW GPU UTIL || PROCESS METRICS: NORMAL |
| Evidence: GPU utilization averaged 24%. || |
| || |
| avg || avg |
| CPU 18% || CPU capacity 14% |
| RAM used 6.2 GB (19%) || RSS used 3.1 GB (10%) |
| GPU util 24% || CUDA allocated 2.9 GB |
| GPU memory/device 3.3 GB (21%) || CUDA reserved 3.2 GB (20%) |
| GPU temperature 42C || |
| GPU power 58W || |
| |
| |
| Full evidence: logs/bert_finetune/final_summary.json (--html-report) |
+----------------------------------------------------------------------------------------------------------------------------------------------------------+
TraceML produces this diagnosis at the end of the instrumented training run.
Want the complete evidence? Jump to the single-run and distributed example reports.
Quickstart
1. Install
TraceML expects an existing PyTorch project. Install the TraceML package with:
pip install traceml-ai
Using uv instead? Run uv add traceml-ai.
2. Instrument the training step
Add TraceML around the core step in your existing PyTorch training script:
+ 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
traceml run train.py
Summary mode is the default. TraceML prints the final diagnosis and writes
final_summary.json and final_summary.txt under logs/<run_name>/.
No training script ready? Try the Colab example.
Example Reports
See the complete single-run report
+----------------------------------------------------------------------------------------------------------------------------------------------------------+
| TraceML Run Summary |
| bert_finetune · 1 rank · 1 GPU observed · 256 common steps · 52.4s |
+----------------------------------------------------------------------------------------------------------------------------------------------------------+
| |
| Verdict: INPUT-BOUND (CRITICAL) |
| Why: Input Wait took 64% of Step Time. |
| Next: Increase workers, prefetch, or storage throughput. |
| |
| STEP TIMING (Window Average), GPU Clock || STEP MEMORY: BALANCED |
| Step Time 200.4 ms 100% || |
| ├─ Input Wait 128.0 ms 64% ◀ cause || |
| ├─ Compute 68.0 ms 34% || avg per-step peak avg |
| │ ├─ Forward 24.0 ms 12% || Allocated 2.9 GB |
| │ ├─ Backward 38.0 ms 19% || Reserved 3.2 GB |
| │ └─ Optimizer 6.0 ms 3% || |
| ├─ H2D 0.4 ms <1% || |
| └─ Residual 3.6 ms 2% || |
| DataLoader fetch: 120.0 ms (CPU, supplemental) || |
| |
| SYSTEM METRICS: LOW GPU UTIL || PROCESS METRICS: NORMAL |
| Evidence: GPU utilization averaged 24%. || |
| || |
| avg || avg |
| CPU 18% || CPU capacity 14% |
| RAM used 6.2 GB (19%) || RSS used 3.1 GB (10%) |
| GPU util 24% || CUDA allocated 2.9 GB |
| GPU memory/device 3.3 GB (21%) || CUDA reserved 3.2 GB (20%) |
| GPU temperature 42C || |
| GPU power 58W || |
| |
| |
| Full evidence: logs/bert_finetune/final_summary.json (--html-report) |
+----------------------------------------------------------------------------------------------------------------------------------------------------------+
Running distributed training? See a rank-straggler diagnosis
+----------------------------------------------------------------------------------------------------------------------------------------------------------+
| TraceML Run Summary |
| ddp_pretrain · 4/4 ranks · 4 GPUs observed · 2/2 nodes · 250 common steps · 40.1s |
+----------------------------------------------------------------------------------------------------------------------------------------------------------+
| |
| Verdict: INPUT STRAGGLER (CRITICAL) |
| Why: R0/N0 waited 254.5 ms for input; R1/N0 waited 3.8 ms for input. |
| Next: Inspect input wait on the slow rank. |
| Scope: N = node · R = global rank · G = GPU index |
| |
| STEP TIMING (Median R1/N0), GPU Clock || STEP MEMORY: BALANCED · 4/4 ranks |
| Step Time 303.7 ms 100% || |
| ├─ Input Wait 3.8 ms 1% || |
| ├─ Compute 259.5 ms 85% || avg per-step peak median rank avg worst rank avg |
| │ ├─ Forward 80.0 ms 26% || Allocated 8.5 GB 9.4 GB, R2/N1 |
| │ ├─ Backward 169.5 ms 56% || Reserved 8.9 GB 9.8 GB, R2/N1 |
| │ └─ Optimizer 10.0 ms 3% || |
| ├─ H2D 1.1 ms <1% || |
| └─ Residual 39.3 ms 13% || |
| DataLoader fetch: 3.7 ms (CPU, supplemental) || |
| |
| SYSTEM METRICS: LOW GPU UTIL · 2/2 nodes || PROCESS METRICS: NORMAL · 4/4 ranks |
| Evidence: GPU utilization averaged 14%. || |
| || |
| median node avg worst node avg || median rank avg worst rank avg |
| CPU 18% 26%, N1 || CPU capacity 12% 81%, R2/N1 |
| RAM used 16.0 GB (27%) 20.8 GB (35%), N1 || RSS used 3.1 GB (10%) 5.4 GB (17%), R1/N0 |
| GPU util 9% 9%, N1 || CUDA allocated 2.9 GB 4.6 GB, R3/N1 |
| GPU memory/device 5.0 GB (31%) 7.0 GB (44%), N1 || CUDA reserved 3.2 GB (20%) 6.8 GB (43%), R3/N1 |
| GPU temperature 58C 70C, N1 || |
| GPU power 220W 280W, N1 || |
| |
| |
| Full evidence: logs/ddp_pretrain/final_summary.json (--html-report) |
+----------------------------------------------------------------------------------------------------------------------------------------------------------+
Read How to Read TraceML Output for the complete field definitions, diagnosis rules, evidence, and recommended actions.
What TraceML Diagnoses
| Diagnosis | Where to investigate |
|---|---|
| Input-bound | DataLoader workers, transforms, tokenization, collation, or storage |
| H2D-bound | Pinned memory, non-blocking copies, batch size, or transfer overlap |
| Compute-bound | Model compute, mixed precision, batch size, or deeper profiling |
| Residual-heavy | Work outside traced phases, CPU stalls, logging, checkpointing, validation, or unobserved transfers |
| Rank straggler | Rank-local input, data imbalance, node variance, or networking |
| Memory creep | Retained tensors, logging references, or cached activations |
Compare Runs
After fixing a bottleneck, compare two summaries to see whether training improved and what changed:
traceml compare before/final_summary.json after/final_summary.json
For example, reducing the DataLoader bottleneck shown above changes the diagnosis and cuts step time:
+--------------------------------------------------------------------------------------+
| TraceML Compare |
+--------------------------------------------------------------------------------------+
| |
| A: before_dataloader_fix |
| B: after_dataloader_fix |
| Delta: B - A |
| Primary diagnosis: INPUT-BOUND -> COMPUTE-BOUND (changed) |
| |
| Verdict: IMPROVEMENT |
| Why: GPU Step Time decreased by 59.9%. |
+--------------------------------------------------------------------------------------+
See the full comparison
+--------------------------------------------------------------------------------------+
| TraceML Compare |
+--------------------------------------------------------------------------------------+
| |
| A: before_dataloader_fix |
| B: after_dataloader_fix |
| Delta: B - A |
| Primary diagnosis: INPUT-BOUND -> COMPUTE-BOUND (changed) |
| |
| Verdict: IMPROVEMENT |
| Why: GPU Step Time decreased by 59.9%. |
| |
| Step Time (GPU comparison clock) |
| Metric A B Delta |
| Step time diagnosis INPUT-BOUND COMPUTE-BOUND changed |
| GPU Step Time 200.4 ms 80.4 ms -120.0 ms (-59.9%) |
| Input 128.0 ms 8.0 ms -120.0 ms (-93.8%) |
| H2D 0.4 ms 0.4 ms +0.0 ms (+0.0%) |
| Compute 68.0 ms 68.0 ms +0.0 ms (+0.0%) |
| Residual 3.6 ms 3.6 ms +0.0 ms (+0.0%) |
| |
| Step Memory |
| Metric A B Delta |
| Step memory diagnosis BALANCED BALANCED same |
| Peak reserved 3.1 GB 3.1 GB 0 B (+0.0%) |
| Memory skew 0.0% 0.0% +0.0 pp |
| |
| Process |
| Metric A B Delta |
| Process diagnosis NORMAL NORMAL same |
| Process CPU avg 95.0% 110.0% +15.0 pp |
| Process RSS avg 1.4 GB 1.6 GB +0.2 GB (+14.3%) |
| |
| System |
| Metric A B Delta |
| System diagnosis LOW GPU UTIL NORMAL changed |
| System CPU avg 18.4% 32.0% +13.6 pp |
| System RAM avg 12.0 GB 13.5 GB +1.5 GB (+12.5%) |
| GPU util avg 24.0% 88.0% +64.0 pp |
| GPU memory avg 18.0% 18.0% +0.0 pp |
+--------------------------------------------------------------------------------------+
See Compare Runs for the complete workflow and artifact format.
Performance Regression Checks
The regression guard is an experimental pilot. Compatible guarded runs can use the same comparison command as a local CI gate. Pass the reference first and the candidate second, then set the maximum allowed Step Time increase:
traceml compare \
logs/reference/final_summary.json \
logs/candidate/final_summary.json \
--max-step-time-regression-pct 5 \
--output compare/reference-vs-candidate
TraceML writes JSON and text evidence before returning the CI exit code. See the Regression Guard for run configuration and supported environments, and Compare Runs for result meanings and exit codes.
Save the Result
Send the compact diagnosis to an existing W&B run:
import traceml_ai as traceml
import wandb
...
summary = traceml.summary(print_text=True)
if summary is not None:
wandb.log(summary)
The same result can be stored in MLflow. See W&B and MLflow for complete examples.
Want live diagnostics during training?
Use the live terminal view locally or over SSH:
traceml run train.py --mode=cli
Use the browser dashboard on a single node:
traceml run train.py --mode=dashboard
For remote browser access and SSH tunneling, see the full quickstart.
Distributed Training and Integrations
- Distributed: DDP, FSDP, and multi-node or Slurm
- Frameworks: Hugging Face, PyTorch Lightning, RF-DETR, Ray Train, DeepSpeed, and MONAI
- Trackers: W&B and MLflow
Summary mode is the documented path for single-node and multi-node runs. Live terminal and dashboard modes are explicit single-node options. See the FAQ for current support and limitations.
Distributed GPU analysis currently assumes homogeneous GPU hardware across ranks. Heterogeneous GPU configurations may produce inaccurate cross-rank analysis and diagnoses.
Learn More
Community
If TraceML helps you find a bottleneck, consider starring the repository. Contributions and real-world slowdown reports are welcome:
License
Apache 2.0. See LICENSE.
1 maintainer
Metadata
1 maintainer
Release files for traceml-ai 0.4.2
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Source distribution (sdist)
| File | Size | Uploaded | Provenance | Actions |
|---|---|---|---|---|
| traceml_ai-0.4.2.tar.gz | 4.2 MB | With provenance | Details for traceml_ai-0.4.2.tar.gz |
Built distribution (wheel)
| File | Interpreter | ABI | Platform | Actions Reset |
|---|---|---|---|---|
| traceml_ai-0.4.2-py3-none-any.whl | Python 3 | none | any | Details for traceml_ai-0.4.2-py3-none-any.whl |
Total release size: 5.0 MB
Release files / traceml_ai-0.4.2.tar.gz
| Download URL | traceml_ai-0.4.2.tar.gz |
|---|---|
| Size | 4.2 MB |
| Tags | Source |
|
SHA-256 checksum How to use checksums |
4a8a8382c22939df707c25298843427e26ae43046fd379a53557aa5ba6c606e3
|
|
BLAKE2b-256 checksum How to use checksums |
28b9eb5307ac65ec93fdf07e01126d977bd8b428e85ac17076f5ab0f6cab32d3
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/7.0.0 CPython/3.13.14
|
Provenance
Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.
PyPI Publish Attestation
PyPI verified that this artifact, at this checksum, originated from the publisher listed below.
Signed by GitHub Actions, verified by PyPI on Oct 2, 2026.
Transparency logRelease files / traceml_ai-0.4.2-py3-none-any.whl
| Download URL | traceml_ai-0.4.2-py3-none-any.whl |
|---|---|
| Size | 727.0 kB |
| Tags | Python 3 |
|
SHA-256 checksum How to use checksums |
3824c6dbc46fcbad356fc9cd003160f0f3e21b834aa57b5478eb9232b45c9265
|
|
BLAKE2b-256 checksum How to use checksums |
63de2dc600172a4bef786a5cf4a22254e3fd02521d6dd237c5e81e45c531331a
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/7.0.0 CPython/3.13.14
|
Provenance
Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.
PyPI Publish Attestation
PyPI verified that this artifact, at this checksum, originated from the publisher listed below.
Signed by GitHub Actions, verified by PyPI on Oct 2, 2026.
Transparency logSecurity for traceml-ai 0.4.2
Provenance
Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.
Full provenance
All 2 files in this release have been uploaded with provenance metadata.
Attestation for traceml_ai-0.4.2.tar.gz
PyPI Publish Attestation
PyPI verified that this artifact, at this checksum, originated from the publisher listed below.
Signed by GitHub Actions, verified by PyPI on Oct 2, 2026.
Transparency logA note on package security
PyPI is a repository for distributing packages, not a service for auditing them. While third-party security researchers and our staff do their best to detect and remove malicious packages, PyPI cannot guarantee the quality, security or suitability of any given package.
To learn more about protecting yourself from malicious actors, visit our package security documentation.
