Profiling and tracing

What it is

PyTorch ships several complementary tools for understanding what a job is doing:

• torch.profiler — CPU + CUDA op-level timing, exported to Chrome trace.
• Memory snapshots — torch.cuda.memory._record_memory_history() produces a binary that https://docs.pytorch.org/memory_viz renders into an interactive timeline.
• NCCL flight recorder — ring buffer of recent collectives, dumped on hang.
• Structured tracing (torch._logging.trace_structured) — production-friendly artefact emission, parsed offline by tlparse.
• torch.autograd.profiler.record_function — manual span markers.

For implementation details see Systems / Profiler.

Op-level profiling

import torch
from torch.profiler import profile, ProfilerActivity, schedule

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    schedule=schedule(wait=1, warmup=1, active=3),
    on_trace_ready=torch.profiler.tensorboard_trace_handler("logs/"),
    record_shapes=True,
    with_stack=True,
    profile_memory=True,
) as prof:
    for step in range(10):
        train_step()
        prof.step()

The output JSON opens in chrome://tracing (or the better-suited https://ui.perfetto.dev). Common things to look for:

• CPU-bound gaps — long stretches with no CUDA activity = CPU is the bottleneck.
• Tiny kernels — a sea of <100 µs kernels indicates you should compile.
• cudaStreamSynchronize — explicit host-device sync, often unintended.
• cudaMemcpyAsync H2D/D2H — accidental tensor moves.

Memory snapshots

For OOM debugging:

torch.cuda.memory._record_memory_history(max_entries=100_000)
# ... run code ...
torch.cuda.memory._dump_snapshot("oom_snapshot.pickle")

Open the file at https://docs.pytorch.org/memory_viz. Each allocation is shown with its Python stack trace and lifetime; the visualizer makes leaks and fragmentation obvious.

The OOM trace is also automatically dumped on torch.cuda.OutOfMemoryError if PYTORCH_CUDA_ALLOC_CONF=alloc_oom_observer:1 is set.

NCCL flight recorder

For debugging distributed hangs:

TORCH_NCCL_TRACE_BUFFER_SIZE=2000000 \
TORCH_NCCL_DUMP_ON_TIMEOUT=1 \
TORCH_NCCL_DEBUG_INFO_PIPE_FILE=/tmp/nccl_dump.bin \
torchrun ...

Each NCCL collective writes a row to a per-rank ring buffer; on timeout (or SIGUSR2) the buffer is dumped. torch.distributed.flight_recorder decodes the dumps and pinpoints which rank/collective deviated.

Structured tracing

torch._logging.trace_structured emits typed JSON-lines artefacts to a file controlled by TORCH_TRACE. The compiler stack uses it heavily — every Inductor-generated kernel, every guard failure, every recompilation gets its own entry. Offline, tlparse (a separate Rust tool) renders these into interactive HTML.

from torch._logging import trace_structured

trace_structured(
    "artifact",
    metadata_fn=lambda: {"name": "my_debug", "encoding": "string"},
    payload_fn=lambda: payload,
)

This is the canonical way (per CLAUDE.md) to log debug artefacts in PyTorch internals.

record_function markers

from torch.autograd.profiler import record_function

with record_function("custom_phase"):
    do_thing()

Adds a custom span to whatever profiler is active. Cheap when no profiler is on. Works at any code site.

Compile-time logs

TORCH_LOGS=... enables structured logs from the compile stack. A few useful values:

Value	What you get
dynamo	Frame conversions
recompiles	Recompilation reasons
graph_breaks	Every graph break with traceback
aot_graphs	The fwd / bwd FX graphs
aot_joint_graph	The pre-partition joint graph
inductor	Inductor scheduling
output_code	The generated Triton / C++ source
kernel_code	Per-kernel source
+all	Everything (very verbose)

Combine with TORCH_COMPILE_DEBUG=1 to dump everything into a directory.

Where to look

Path	Contents
torch/profiler/	torch.profiler API
torch/csrc/profiler/	C++ collector + Kineto integration
torch/_logging/	trace_structured
torch/cuda/memory.py	Memory snapshot APIs
torch/distributed/flight_recorder/	NCCL flight recorder decoder

Where to read next

• Systems / Profiler — implementation.
• Features / torch.compile — compile-time logs.
• Features / Distributed training — flight recorder usage.