TraceML: Stop Flying Blind Inside Your Training Loop

You have always been able to watch the loss curve. TraceML lets you see training efficiency inside the loop.

Job running. Loss going down. GPUs allocated. No crashes.

But inside each training step, you still do not know how time is actually split across data loading, forward, backward, and optimizer.

That is not a failure. It is just a blind spot that has always existed. A training loop can be correct, the loss curve can look fine, and the run can still be quietly inefficient. Maybe the dataloader is eating half the step. Maybe one rank is lagging in a distributed job. Maybe memory is slowly climbing during the run. Maybe everything is actually fine.

Until now, there was no easy way to know. That is the gap TraceML is built to close.

The Missing Layer

When engineers investigate a training run, the usual toolkit looks like this:

• nvidia-smi and cluster dashboards for GPU utilization
• W&B, MLflow, or TensorBoard for loss curves and run history
• PyTorch Profiler or Nsight Systems when something needs deep inspection

All useful. But there is a layer missing between the job is running and open a heavyweight profiler: lightweight step-level visibility during normal training, with minimal setup.

System monitors see the machine. Experiment trackers see the outcome. Deep profilers see kernels and timelines, but only once you already suspect something specific, and only if you’re willing to pay the overhead.

What about just knowing, on any run, how time is splitting across the training step?

That is what TraceML gives you.

One Boundary, Live Diagnostics

TraceML is open source and integrates with a small code change. Mark the training step once:

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

Then run:

traceml run train.py

That is the core idea:

Mark the training step once. TraceML turns that boundary into structured diagnostics.

TraceML uses the step boundary to organize timing, memory, rank, process, and system signals into a diagnosis of where the run is wasting time.

What You Get

While the job runs, TraceML opens a live terminal view alongside your logs.

TraceML live dashboard on a PyTorch training run. The run is classified as compute-bound, with backward dominating step time, while the memory panel flags growing reserved memory across the observed window.

At a glance it answers:

• Is dataloader fetch taking too much of the step?
• Is compute dominating?
• Are ranks behaving differently in supported distributed runs?
• Is memory stable or showing pressure?
• Is the run balanced enough that deeper profiling is not the next move?

At the end of the run, TraceML writes final_summary.json,structured JSON, not a giant trace. Small enough to store, easy to log, simple to compare across runs. A simplified version of the summary looks like:

{
  "step_time": {
    "diagnosis": "INPUT_BOUND",
    "dataloader_pct": 47.0,
    "forward_pct": 31.0,
    "backward_pct": 18.0,
    "optimizer_pct": 4.0
  },
  "step_memory": {
    "diagnosis": "BALANCED"
  }
}

Each diagnosis maps directly to where you look next:

A lot of engineering time disappears into investigating runs that do not need it. Knowing the run is balanced means you can move on.

Tool Zero, Not a Profiler Replacement

TraceML is not trying to replace PyTorch Profiler or Nsight Systems. When you need kernel-level traces, CUDA timelines, or NCCL behavior, those are the right tools.

But you should not need a heavyweight profiler just to answer the first question.

Run TraceML first. Classify the run. Then decide whether deeper profiling is worth the cost and where to point it. That is why it makes sense to think of

TraceML as tool zero for PyTorch training diagnostics.

It does not replace the microscope. It tells you where to aim it.

For ML Platform Teams

The most useful part is not just the live view, it is the final_summary.json.

That file is small, structured, and easy to collect across every job. It can be logged into W&B or MLflow, stored as a CI artifact, or diffed against a previous run. TraceML also ships a compare workflow:

traceml compare run_a.json run_b.json

When a training regression appears, the question should not only be did throughput drop? It should be where did the time move? That is the kind of signal final_summary.json is designed to make collectable across runs.

Try It

pip install traceml-ai
traceml run train.py

Repo: https://github.com/traceopt-ai/traceml

TraceML is open source and Apache-2.0 licensed.

Works today with single GPU and single-node DDP/FSDP. Multi-node support is coming soon.

If you work on PyTorch training, distributed jobs, or ML infrastructure: feedback, issues, and contributions are very welcome.

Before opening a profiler, know what kind of slow you are looking at.