TraceML
Find out why your PyTorch training is slow—before it wastes GPU hours.
Quickstart • Try in Colab • Integrations • Documentation • Discord
TraceML is an open-source tool that explains why PyTorch training is slow. It analyzes the full training run and gives you:
- A diagnosis: data loading, GPU compute, waiting, memory growth, or a slow distributed worker.
- The evidence: timing, CPU/GPU usage, memory, and per-worker measurements.
- The next step: what part of your training setup to investigate first.
Here is an example where TraceML finds that a slow DataLoader is leaving the GPU waiting:
+----------------------------------------------------------------------------+
| TraceML Run Summary | duration 52.4s |
+----------------------------------------------------------------------------+
| |
| TraceML Verdict: INPUT-BOUND / CRITICAL |
| Why: Input Wait is 64.0% of the typical GPU Step Time. |
| Next: Increase workers, prefetch, or storage throughput. |
| |
| Section Status |
| Section Status Severity |
| ------------------------------------------------ |
| Step Time INPUT-BOUND CRITICAL |
| System LOW GPU UTIL INFO |
| Process NORMAL INFO |
| Step Memory BALANCED INFO |
| |
| System Evidence |
| Metric Average |
| ---------------------------------- |
| CPU Util 18.4% |
| GPU Util 24.0% |
| GPU Memory 3.33GB |
| GPU Temp 42C |
| |
| Step Time Evidence |
| Phase Average Share |
| ------------------------------------------------ |
| Step Time 200.4ms 100.0% |
| Input Wait 128.0ms 64.0% |
| Traced Step Time 72.0ms supplemental |
| Compute 68.0ms 34.0% |
| Residual 3.6ms 1.8% |
| H2D 0.4ms 0.2% |
+----------------------------------------------------------------------------+
Running distributed training? See a rank-straggler diagnosis
+----------------------------------------------------------------------------+
| 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 |
| --------------------------------------------------------------------------|
| Step Time 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 |
+----------------------------------------------------------------------------+
Quickstart
1. Install
pip install traceml-ai
Or, in a project managed by uv:
uv add traceml-ai
2. Instrument the training step
+ 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:
logs/<run_name>/final_summary.json
logs/<run_name>/final_summary.txt
See the
full quickstart for source-only examples,
Docker (both require a repository checkout), Colab, direct python/torchrun
launches, HTML reports, and advanced options.
What TraceML Diagnoses
| Diagnosis | Where to investigate |
|---|---|
| Input-bound | DataLoader workers, transforms, tokenization, collation, or storage |
| 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 |
| Run regression | Code, data, environment, hardware, or infrastructure changes |
Read How to Read TraceML Output for the diagnosis rules, evidence fields, and recommended next actions.
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.
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, Ray Train, and DeepSpeed
- Trackers: W&B and MLflow
Summary mode supports single-node and multi-node runs. Live terminal and dashboard modes are explicit single-node options. See the FAQ for current support and limitations.
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
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Filter files by name, interpreter, ABI, and platform.
If you're not sure about the file name format, learn more about wheel file names.
Copy a direct link to the current filters
File details
Details for the file traceml_ai-0.3.6.tar.gz.
File metadata
- Download URL: traceml_ai-0.3.6.tar.gz
- Upload date:
- Size: 3.7 MB
- Tags: Source
- Uploaded using Trusted Publishing? Yes
- Uploaded via: twine/7.0.0 CPython/3.13.14
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
dc2e2e734910170ad2291f0a2ce72d5004ae12fd6b1962803951a4e62a465ccb
|
|
| MD5 |
0040d4ec98b6f5f94e2a1347bf55a356
|
|
| BLAKE2b-256 |
cb2ae0c5830d2a87e9fb64a9d9a8d808f76cec218ac387814f624d8a5455a11d
|
Provenance
The following attestation bundles were made for traceml_ai-0.3.6.tar.gz:
Publisher:
release.yml on traceopt-ai/traceml
-
Statement:
-
Statement type:
https://in-toto.io/Statement/v1 -
Predicate type:
https://docs.pypi.org/attestations/publish/v1 -
Subject name:
traceml_ai-0.3.6.tar.gz -
Subject digest:
dc2e2e734910170ad2291f0a2ce72d5004ae12fd6b1962803951a4e62a465ccb - Sigstore transparency entry: 2344978438
- Sigstore integration time:
-
Permalink:
traceopt-ai/traceml@d5cec51180c4c828f195656890ef5bc91c3cc4d4 -
Branch / Tag:
refs/heads/main - Owner: https://github.com/traceopt-ai
-
Access:
public
-
Token Issuer:
https://token.actions.githubusercontent.com -
Runner Environment:
github-hosted -
Publication workflow:
release.yml@d5cec51180c4c828f195656890ef5bc91c3cc4d4 -
Trigger Event:
workflow_dispatch
-
Statement type:
File details
Details for the file traceml_ai-0.3.6-py3-none-any.whl.
File metadata
- Download URL: traceml_ai-0.3.6-py3-none-any.whl
- Upload date:
- Size: 559.6 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? Yes
- Uploaded via: twine/7.0.0 CPython/3.13.14
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
5fba1341f4cfcbc895c726ba07aa6b61ffd06215371afbe6cad5ef9ea8b089d9
|
|
| MD5 |
b9226dfb006c250524c78c2ebc2f504a
|
|
| BLAKE2b-256 |
ebbf2c7f1a811b3b61ad1dad65bbf85bcb8d6b70c6f02e71aa429900f993529d
|
Provenance
The following attestation bundles were made for traceml_ai-0.3.6-py3-none-any.whl:
Publisher:
release.yml on traceopt-ai/traceml
-
Statement:
-
Statement type:
https://in-toto.io/Statement/v1 -
Predicate type:
https://docs.pypi.org/attestations/publish/v1 -
Subject name:
traceml_ai-0.3.6-py3-none-any.whl -
Subject digest:
5fba1341f4cfcbc895c726ba07aa6b61ffd06215371afbe6cad5ef9ea8b089d9 - Sigstore transparency entry: 2344978473
- Sigstore integration time:
-
Permalink:
traceopt-ai/traceml@d5cec51180c4c828f195656890ef5bc91c3cc4d4 -
Branch / Tag:
refs/heads/main - Owner: https://github.com/traceopt-ai
-
Access:
public
-
Token Issuer:
https://token.actions.githubusercontent.com -
Runner Environment:
github-hosted -
Publication workflow:
release.yml@d5cec51180c4c828f195656890ef5bc91c3cc4d4 -
Trigger Event:
workflow_dispatch
-
Statement type:
