TraceML

Find out why your PyTorch training is slow—before it wastes GPU hours.

PyPI version CI Python 3.10+ License GitHub stars

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

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.

Key dates

PyPI data
Data sourced directly from PyPI's database.
  • Released:
Latest release

1 maintainer

PyPI data
Data sourced directly from PyPI's database.
Avatar for abhinavsriva from gravatar.com abhinavsriva

Credits

Author/Maintainer: OptAI UG (haftungsbeschränkt)

License

Apache 2.0

Requires

Python >=3.10

Provides Extra

dashboard torch dev hf lightning ray deepspeed docs

Tags

pytorch profiler gpu training bottleneck cuda dataloader observability machine-learning deep-learning

Classifiers

Report project as malware

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

traceml_ai-0.3.6.tar.gz (3.7 MB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

traceml_ai-0.3.6-py3-none-any.whl (559.6 kB view details)

Uploaded Python 3

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

Hashes for traceml_ai-0.3.6.tar.gz
Algorithm Hash digest
SHA256 dc2e2e734910170ad2291f0a2ce72d5004ae12fd6b1962803951a4e62a465ccb
MD5 0040d4ec98b6f5f94e2a1347bf55a356
BLAKE2b-256 cb2ae0c5830d2a87e9fb64a9d9a8d808f76cec218ac387814f624d8a5455a11d

See more details on using hashes here.

Provenance

The following attestation bundles were made for traceml_ai-0.3.6.tar.gz:

Publisher: release.yml on traceopt-ai/traceml

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

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

Hashes for traceml_ai-0.3.6-py3-none-any.whl
Algorithm Hash digest
SHA256 5fba1341f4cfcbc895c726ba07aa6b61ffd06215371afbe6cad5ef9ea8b089d9
MD5 b9226dfb006c250524c78c2ebc2f504a
BLAKE2b-256 ebbf2c7f1a811b3b61ad1dad65bbf85bcb8d6b70c6f02e71aa429900f993529d

See more details on using hashes here.

Provenance

The following attestation bundles were made for traceml_ai-0.3.6-py3-none-any.whl:

Publisher: release.yml on traceopt-ai/traceml

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

Release history Release notifications | RSS feed

This release

0.3.6 This release

2 files

0.3.5.1

2 files

0.3.5

2 files

0.3.4

2 files

0.3.3

2 files

0.3.2

2 files

0.3.1

2 files

0.3.0

2 files

0.2.15

2 files

0.2.14

2 files

0.2.13

2 files

0.2.12

2 files

0.2.11

2 files

0.2.10

2 files

0.2.9

2 files

0.2.8

2 files

0.2.7

2 files

0.2.6

2 files

0.2.5

2 files

0.2.4

2 files

0.2.3

2 files

0.2.2

2 files

0.2.1

2 files

0.2.0

2 files

Pre-release

0.2.0a0

2 files

0.1.9

2 files

0.1.8

2 files

Yanked

0.1.6

2 files

Yanked reason: Declares compatibility with numpy>=2, but predates its modern feature set; as a result, the resolver silently selects it in modern environments and causes runtime failures. Please install traceml-ai>=0.3.5.1 instead.

Yanked

0.1.5

2 files

Yanked reason: Declares compatibility with numpy>=2, but predates its modern feature set; as a result, the resolver silently selects it in modern environments and causes runtime failures. Install traceml-ai>=0.3.5.1 instead.

Yanked

0.1.3

2 files

Yanked reason: Declares compatibility with numpy>=2, but predates its modern feature set; as a result, the resolver silently selects it in modern environments and causes runtime failures.

Yanked

0.1.1

2 files

Yanked reason: Declares compatibility with numpy>=2, but predates its modern feature set; as a result, the resolver silently selects it in modern environments and causes runtime failures.

Yanked

0.1.0

2 files

Yanked reason: Declares compatibility with numpy>=2, but predates its modern feature set; as a result, the resolver silently selects it in modern environments and causes runtime failures.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Sentry Error logging StatusPage Status page