LLM Trace Analysis, in depth: reading PyTorch profiler traces to find idle GPUs, exposed communication and launch overhead

What is in a trace

A PyTorch profiler trace is a JSON file with a traceEvents array. Most entries are complete events ("ph": "X") with a name, a category, a start timestamp ts and a duration dur, both in microseconds. The categories you will use most are CPU operators (cpu_op, such as aten::mm), CUDA runtime calls (cuda_runtime, such as cudaLaunchKernel), GPU kernels (kernel), memory copies and your own record_function annotations. GPU events carry the device and stream in args, and a correlation ID links each kernel back to the runtime call that launched it, which is how the viewer draws arrows from CPU to GPU.

The key mental model is that the CPU and GPU run asynchronously. The CPU thread enqueues kernels and moves on; the GPU drains the queue. While the CPU stays ahead, the GPU never waits. When the CPU falls behind, because of a slow dataloader, Python overhead, or a synchronising call like .item(), the GPU queue empties and the trace shows gaps. Nearly every finding in a trace is a statement about that race between the two timelines.

Capturing a trace that represents steady state

The first steps of a run include CUDA context creation, memory allocator warm-up, autotuning and compilation. A trace of step 0 measures those, not your training loop. Use the profiler's schedule to skip them, record a few steady-state steps and write one file per rank:

import torch, torch.distributed as dist
from torch.profiler import (profile, schedule, record_function,
                            ProfilerActivity, tensorboard_trace_handler)

rank = dist.get_rank()
prof = profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    schedule=schedule(wait=5, warmup=2, active=3, repeat=1),   # record steps 7-9
    on_trace_ready=tensorboard_trace_handler("traces/run42", worker_name=f"rank{rank}"),
    record_shapes=False, with_stack=False, profile_memory=False,  # keep overhead low
)
with prof:
    for step, batch in enumerate(loader):
        with record_function("forward"):
            loss = model(**batch).loss
        with record_function("backward"):
            loss.backward()
        with record_function("optimizer"):
            optimizer.step()
            optimizer.zero_grad(set_to_none=True)
        prof.step()              # advances the schedule; required
        if step == 12:
            break

Turn on record_shapes, with_stack or profile_memory only for a second, targeted capture: they add CPU overhead that changes the very gaps you are trying to measure. Two or three active steps are enough; traces of large models grow to hundreds of megabytes per rank per step and become hard to open. Capture every rank at the same steps, because many problems only appear when you compare ranks.

The anatomy of a training step

With FSDP, as described in the FSDP article, a step has a recognisable shape. Before each layer's forward, an all-gather collects its sharded parameters. During backward, another all-gather fetches parameters again and a reduce-scatter distributes gradients. Collectives run as NCCL kernels on their own stream, so they can overlap with compute on the default stream. The optimizer step at the end is a burst of small elementwise kernels unless it uses fused implementations.

Four numbers to compute first

Before reading individual kernels, reduce the trace to four numbers per rank and step. GPU busy percentage is the union of all kernel intervals divided by the step's wall time. It must be a union, because kernels on different streams overlap; summing durations can exceed 100%. Exposed communication is NCCL kernel time with no compute kernel running at the same moment: that is time collectives add to the step. Launch gaps are the many short idle intervals between kernels, typically a few to tens of microseconds each, which add up when kernels are tiny. Large gaps are idle stretches of a millisecond or more, nearly always the GPU waiting for the host.

These are simple interval arithmetic. The script below depends only on kernel events and their ts and dur fields, so it works across PyTorch versions. It was checked against a synthetic trace with hand-computed answers.

import json

def union(intervals):
    """Merge [start, end) intervals into a sorted, non-overlapping list."""
    out = []
    for s, e in sorted(intervals):
        if out and s <= out[-1][1]:
            out[-1][1] = max(out[-1][1], e)
        else:
            out.append([s, e])
    return out

def length(merged):
    return sum(e - s for s, e in merged)

def intersect(a, b):
    """Total overlap between two merged interval lists."""
    i = j = total = 0
    while i < len(a) and j < len(b):
        lo, hi = max(a[i][0], b[j][0]), min(a[i][1], b[j][1])
        total += max(0, hi - lo)
        if a[i][1] < b[j][1]:
            i += 1
        else:
            j += 1
    return total

def analyse(path, device=0):
    events = json.load(open(path))["traceEvents"]
    kernels = [e for e in events
               if e.get("cat") == "kernel" and e.get("args", {}).get("device", device) == device]
    if not kernels:
        raise SystemExit("no GPU kernels in trace: was CUDA activity enabled?")
    comm, comp = [], []
    for k in kernels:
        span = (k["ts"], k["ts"] + k["dur"])              # microseconds
        (comm if "nccl" in k["name"].lower() else comp).append(span)
    t0 = min(s for s, _ in comm + comp)
    t1 = max(e for _, e in comm + comp)
    busy, comm_u, comp_u = union(comm + comp), union(comm), union(comp)
    gaps = [b[0] - a[1] for a, b in zip(busy, busy[1:])]
    return {
        "busy_pct": 100 * length(busy) / (t1 - t0),
        "exposed_comm_ms": (length(comm_u) - intersect(comm_u, comp_u)) / 1e3,
        "launch_gap_ms": sum(g for g in gaps if g < 50) / 1e3,
        "large_gaps_ms": sorted((g / 1e3 for g in gaps if g >= 1000), reverse=True)[:5],
    }

Run it per rank and per step, not on the whole file, by first slicing events to the window of each ProfilerStep#N annotation. The 50 microsecond and 1 millisecond thresholds are starting points; adjust them to your kernel sizes.

Holistic Trace Analysis

Meta's open-source Holistic Trace Analysis (HTA) library computes these and more across a directory of per-rank traces. Its documented entry point is TraceAnalysis:

from hta.trace_analysis import TraceAnalysis

analyzer = TraceAnalysis(trace_dir="traces/run42")
temporal = analyzer.get_temporal_breakdown()      # idle / compute / non-compute per rank
idle     = analyzer.get_idle_time_breakdown()     # why the GPU was idle
kernels  = analyzer.get_gpu_kernel_breakdown()    # time by kernel type and name
overlap  = analyzer.get_comm_comp_overlap()       # % of comm overlapped by compute
launches = analyzer.get_cuda_kernel_launch_stats()  # CPU launch vs GPU run times

It also offers frequent kernel sequence mining, memory bandwidth and queue length summaries, and a TraceDiff module for comparing two runs. Return types differ between functions and releases, so check the documentation for your installed version. HTA reads distributed metadata that the profiler records in each trace to tell ranks apart, so feed it unmodified per-rank files from the same steps. The custom script above stays useful for CI, where a few lines of interval arithmetic are easier to pin than a library's output schema.

Patterns and what they mean

The straggler row deserves emphasis. A collective cannot finish until every rank arrives, so on fast ranks it appears as a long NCCL kernel, which looks like slow communication. The slow rank is the one whose collective is short, because everyone else was already waiting. Techniques for hiding collectives behind compute are covered in collective communication overlap.

Inference traces: prefill and decode

Inference traces look different. Prefill processes the whole prompt in large matrix multiplies and is usually compute-bound, with dense, long kernels. Decode generates one token per sequence per step. Each step runs every layer with tiny inputs, so kernels last microseconds and the trace becomes a comb of short kernels separated by launch gaps. A decode step at batch size 1 can show the GPU busy well under half the time, purely because the CPU cannot launch kernels fast enough.

The standard fix is CUDA graphs: record the kernel sequence of one decode step once and replay it with a single launch. Serving engines such as vLLM capture graphs for decode by default, so if your serving trace still shows a comb, check whether graph capture was disabled or the batch size fell outside the captured sizes. The other lever is batching more sequences per step, which makes each kernel longer and the launch overhead relatively smaller. In serving, measure per-step time and tokens per second alongside the trace, since a trace of a single request hides queueing and scheduler effects.

Worked example: an eight-GPU FSDP step

Consider an illustrative case: a 7B model trained with FSDP on eight GPUs, step time 1.24 s, with nvidia-smi showing near-full utilisation. The numbers below show the method, not a benchmark. Running the script on rank 0 for steps 7 to 9 gives about 71% GPU busy, 210 ms of exposed communication, 40 ms of launch gaps and one large gap of about 150 ms at the start of each step.

The large gap lines up with the CPU row sitting in the dataloader's __next__. The workers were tokenising on the fly. Moving to pre-tokenised shards and adding prefetching removes it. The exposed communication sits at the end of backward: the reduce-scatters for the first layers have no backward compute left to hide behind, and in this run backward prefetch was off, so all-gathers also waited. Enabling backward prefetch reduces the exposed time to the unavoidable tail. A second capture after both changes is the only way to confirm the effect; record the new four numbers next to the old ones. A result in the region of 85-90% busy and a step under one second would be a plausible outcome for fixes of this kind, but only the second trace tells you.

Pitfalls in the analysis itself

• Profiler overhead. Stack and shape recording slows the CPU and creates gaps that are not there in production. Diagnose gaps from a lean capture.
• Clock alignment across ranks. Each rank's timestamps come from its own host clock. Compare durations across ranks freely; compare absolute times only after aligning on a shared event, such as the end of a collective.
• Attributing GPU time to CPU operators. Because execution is asynchronous, the CPU operator that is running when the GPU is slow is usually not the cause. Follow correlation IDs.
• Warm-up and compilation. Steps with torch.compile recompilation or autotuning are outliers; skip them or capture them separately.
• Kernel names. Classifying communication by "nccl" in the name is a heuristic; check how your collectives library names its kernels.

What to do next

1. Add record_function annotations for forward, backward, optimizer and data loading, then capture three steady-state steps on every rank with the schedule above.
2. Compute GPU busy, exposed communication, launch gaps and large gaps per rank and step, and write them down as the baseline.
3. Fix the largest bucket first: host waits, then exposed communication, then launch overhead.
4. Compare ranks; if one collective is shorter on one rank, investigate that rank's host, data or GPU.
5. Run HTA on the trace directory for the kernel breakdown and overlap percentage.
6. Re-capture after every change and keep the four numbers in the pull request.
7. For kernel-level questions that remain, move to Nsight Compute, and for multi-node communication read NCCL collectives.