PyTorch Performance Debugging: CUDA OOM, Slow Training, GPU Utilization and torch.compile

AI, PyTorch

August 08, 2026

Page content

PyTorch: Zero to Hero — Step 09

At this point in the series, the model runs.

That does not mean it runs well.

A training loop can be correct and still waste most of the machine.

A model can fit in memory and still spend half its time waiting on synchronization.

A torch.compile call can make code faster, slower, or simply move the bottleneck somewhere else.

A CUDA out-of-memory error can be caused by the model, the optimizer, activations, fragmentation, a leaked reference, a larger batch, a longer sequence, or an innocent-looking tensor that was kept alive by Python.

Performance work starts when you stop guessing.

The rule for this article is simple:

Measure first. Optimize second. Measure again.

What we are going to debug

This article is built around the questions programmers actually end up asking:

Why is my PyTorch training slow?
Why is GPU utilization low?
Why is GPU utilization sawtoothing?
Why does CUDA run out of memory?
Why does reducing batch size not fix OOM?
Why is torch.compile slower than eager mode?
Why does torch.compile keep recompiling?
Why did mixed precision not help?
Why does a benchmark lie unless I synchronize CUDA?
Why does inference keep allocating memory?
Why does memory not return after del tensor?
Why does one shape work and another shape trigger compilation again?

By the end we will have a repeatable performance-debugging workflow rather than a bag of tuning tricks.

First: build a benchmark that does not lie

CUDA execution is asynchronous.

This means code like this is usually wrong:

import time

start = time.perf_counter()
y = model(x)
elapsed = time.perf_counter() - start

print(elapsed)

The CPU can enqueue GPU work and keep moving.

If you want wall-clock timing around CUDA work, synchronize around the region you are measuring.

import time
import torch


def time_cuda(fn, *, warmup=10, repeats=50):
    for _ in range(warmup):
        fn()

    torch.cuda.synchronize()
    start = time.perf_counter()

    for _ in range(repeats):
        fn()

    torch.cuda.synchronize()
    elapsed = time.perf_counter() - start

    return elapsed / repeats

Use it like this:

model = model.cuda().eval()
x = torch.randn(64, 1024, device="cuda")

with torch.inference_mode():
    seconds = time_cuda(lambda: model(x))

print(f"{seconds * 1000:.3f} ms")

Warmup matters too.

The first iterations may include:

CUDA context initialization
kernel loading
allocator growth
cuDNN/autotuning work
torch.compile compilation

Do not compare cold eager execution against warm compiled execution.

Do not compare one iteration against another configuration that ran fifty.

CUDA events are better for GPU timing

For GPU-only timing, CUDA events are useful:

import torch

start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)

start.record()
y = model(x)
end.record()

end.synchronize()
print(f"{start.elapsed_time(end):.3f} ms")

For a proper benchmark:

def benchmark_cuda(fn, warmup=20, repeats=100):
    for _ in range(warmup):
        fn()

    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)

    start.record()
    for _ in range(repeats):
        fn()
    end.record()

    end.synchronize()
    return start.elapsed_time(end) / repeats

This gives milliseconds per iteration.

Throughput is often more useful than latency

For training, I usually care about examples per second or tokens per second.

elapsed = 1.25
examples = 4096

print(f"examples/sec = {examples / elapsed:.2f}")

For language models:

batch_size = 16
sequence_length = 2048
steps = 50
elapsed = 12.4

tokens = batch_size * sequence_length * steps
print(f"tokens/sec = {tokens / elapsed:,.0f}")

A faster single forward pass does not necessarily mean a faster training system.

A minimal training benchmark

import time
import torch


def benchmark_training_step(model, optimizer, x, y, steps=100):
    model.train()

    for _ in range(10):
        optimizer.zero_grad(set_to_none=True)
        logits = model(x)
        loss = torch.nn.functional.cross_entropy(logits, y)
        loss.backward()
        optimizer.step()

    if x.is_cuda:
        torch.cuda.synchronize()

    start = time.perf_counter()

    for _ in range(steps):
        optimizer.zero_grad(set_to_none=True)
        logits = model(x)
        loss = torch.nn.functional.cross_entropy(logits, y)
        loss.backward()
        optimizer.step()

    if x.is_cuda:
        torch.cuda.synchronize()

    elapsed = time.perf_counter() - start

    return {
        "seconds": elapsed,
        "steps_per_second": steps / elapsed,
        "examples_per_second": steps * x.shape[0] / elapsed,
    }

The important part is not this exact helper.

The important part is having one repeatable benchmark before changing anything.

Why is GPU utilization low?

Low GPU utilization usually means one of two things:

GPU has no work

or:

GPU workload is too small to occupy the hardware efficiently

Those are very different problems.

Start by splitting the training step into regions.

for x, y in loader:
    t0 = time.perf_counter()

    x = x.to(device, non_blocking=True)
    y = y.to(device, non_blocking=True)

    t1 = time.perf_counter()

    optimizer.zero_grad(set_to_none=True)

    logits = model(x)
    loss = loss_fn(logits, y)
    loss.backward()
    optimizer.step()

    if device.type == "cuda":
        torch.cuda.synchronize()

    t2 = time.perf_counter()

    print({
        "transfer": t1 - t0,
        "compute": t2 - t1,
    })

For serious work use the profiler, but even crude segmentation can expose a large bottleneck.

Use the PyTorch profiler

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

activities = [ProfilerActivity.CPU]
if torch.cuda.is_available():
    activities.append(ProfilerActivity.CUDA)

with profile(
    activities=activities,
    record_shapes=True,
    profile_memory=True,
    with_stack=True,
) as prof:
    with record_function("train_step"):
        optimizer.zero_grad(set_to_none=True)

        with record_function("forward"):
            logits = model(x)
            loss = loss_fn(logits, y)

        with record_function("backward"):
            loss.backward()

        with record_function("optimizer"):
            optimizer.step()

print(
    prof.key_averages().table(
        sort_by="self_cuda_time_total" if torch.cuda.is_available()
        else "self_cpu_time_total",
        row_limit=30,
    )
)

The profiler answers a much better question than “why is PyTorch slow?”

It lets you ask:

Where is the time actually going?

Add profiler ranges around your own code

from torch.profiler import record_function

with record_function("decode_batch"):
    x, y = decode(batch)

with record_function("host_to_device"):
    x = x.to(device, non_blocking=True)
    y = y.to(device, non_blocking=True)

with record_function("model_forward"):
    logits = model(x)

Do not settle for a profile containing hundreds of kernels with no application-level context.

Mark your own boundaries.

The first CUDA OOM question: how much memory is actually used?

PyTorch gives us useful counters.

import torch


def mib(n):
    return n / 1024**2


def print_cuda_memory():
    if not torch.cuda.is_available():
        print("CUDA unavailable")
        return

    print({
        "allocated_MiB": mib(torch.cuda.memory_allocated()),
        "reserved_MiB": mib(torch.cuda.memory_reserved()),
        "max_allocated_MiB": mib(torch.cuda.max_memory_allocated()),
        "max_reserved_MiB": mib(torch.cuda.max_memory_reserved()),
    })

Call:

torch.cuda.reset_peak_memory_stats()

loss = training_step(...)

print_cuda_memory()

allocated and reserved are not the same thing.

The caching allocator can reserve memory so future allocations do not repeatedly go back to CUDA.

Measure peak memory per configuration

def measure_peak_memory(fn):
    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()

    fn()

    torch.cuda.synchronize()

    return {
        "peak_allocated_MiB": torch.cuda.max_memory_allocated() / 1024**2,
        "peak_reserved_MiB": torch.cuda.max_memory_reserved() / 1024**2,
    }

Then compare batch sizes:

for batch_size in [8, 16, 32, 64]:
    x = torch.randn(batch_size, 1024, device="cuda")

    stats = measure_peak_memory(lambda: model(x))
    print(batch_size, stats)

Now batch-size tuning is evidence-driven.

Memory lives in more places than model parameters

Training memory roughly includes:

parameters
+ gradients
+ optimizer state
+ activations saved for backward
+ temporary workspaces
+ allocator overhead

For Adam-like optimizers, optimizer state can be substantial.

A parameter count alone is not a training-memory estimate.

We can visualize the breakdown:

    pie title Training VRAM breakdown (typical)
    "parameters" : 25
    "gradients" : 25
    "optimizer state" : 30
    "activations" : 15
    "other" : 5

Count parameter bytes

def parameter_bytes(model):
    return sum(
        p.numel() * p.element_size()
        for p in model.parameters()
    )

print(parameter_bytes(model) / 1024**2, "MiB")

Gradients can add roughly another parameter-sized allocation when present.

Optimizer state can add more again.

Activations frequently dominate as batch size or sequence length grows.

Sequence length can hurt much more than batch size

For attention, the attention matrix scales with sequence length squared.

A conceptual shape is:

(B, H, T, T)

Double T and the number of elements in that matrix becomes roughly four times larger.

That is why this:

sequence_length = 4096

can be dramatically more expensive than:

sequence_length = 2048

Even when batch size is unchanged.

Find leaked tensors

One extremely common memory bug is keeping references to tensors that still own computation graphs.

Bad:

losses = []

for batch in loader:
    loss = training_step(batch)
    losses.append(loss)

That may retain graph-connected tensors.

Better:

losses.append(loss.item())

or:

losses.append(loss.detach().cpu())

depending on what you actually need.

Likewise:

outputs.append(logits)

can be disastrous inside a long training loop.

Ask whether the tensor really needs to stay alive.

Track memory across iterations

for step, batch in enumerate(loader):
    loss = training_step(batch)

    if step % 50 == 0:
        print(
            step,
            torch.cuda.memory_allocated() / 1024**2,
            torch.cuda.memory_reserved() / 1024**2,
        )

If allocated memory climbs every iteration, investigate retained references.

A healthy caching allocator may keep reserved memory high, but continually increasing allocated memory is suspicious.

You can visualise this growth:

import matplotlib.pyplot as plt

allocated = []
for step in range(200):
    y = model(x)
    loss = y.square().mean()
    loss.backward()
    allocated.append(torch.cuda.memory_allocated() / 1024**2)

plt.plot(allocated)
plt.xlabel('Step')
plt.ylabel('Allocated (MiB)')
plt.title('GPU memory over iterations (watch for leaks)')
plt.grid(True)
plt.show()

A flat line is healthy; a steady upward slope is a red flag.

`torch.cuda.empty_cache()` is not a memory leak fix

It can release unused cached blocks back to CUDA so other processes can use them.

It does not free live tensors.

This:

torch.cuda.empty_cache()

cannot save you if Python still references the tensor responsible for the allocation.

Start with ownership.

Use `memory_summary()` when OOM gets confusing

print(torch.cuda.memory_summary())

That gives a detailed allocator report.

For deeper allocator debugging, PyTorch can also record and dump CUDA memory snapshots.

torch.cuda.memory._record_memory_history()

# run workload

snapshot = torch.cuda.memory.memory_snapshot()

Or dump a snapshot for external inspection when using the supported snapshot workflow.

This is where you go when aggregate numbers stop being enough.

Activation checkpointing trades compute for memory

If activations dominate memory, checkpointing can recompute some activations during backward instead of retaining all of them.

Conceptually:

less memory
↕
more recomputation

Example:

from torch.utils.checkpoint import checkpoint


def forward(self, x):
    x = checkpoint(self.block1, x, use_reentrant=False)
    x = checkpoint(self.block2, x, use_reentrant=False)
    return self.head(x)

Do not add checkpointing before proving activations are your problem.

Gradient accumulation is not free memory

If batch size 64 does not fit, you can approximate its gradient batch with smaller micro-batches.

accumulation_steps = 4
optimizer.zero_grad(set_to_none=True)

for step, (x, y) in enumerate(loader):
    x = x.to(device)
    y = y.to(device)

    logits = model(x)
    loss = loss_fn(logits, y) / accumulation_steps
    loss.backward()

    if (step + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad(set_to_none=True)

This reduces activation memory per micro-batch.

But it does not make optimizer state or model parameters disappear.

Mixed precision: measure speed and memory separately

Modern CUDA workloads can often benefit from autocast.

with torch.autocast(device_type="cuda", dtype=torch.float16):
    logits = model(x)
    loss = loss_fn(logits, y)

For training with fp16, use scaling where appropriate:

scaler = torch.amp.GradScaler("cuda")

optimizer.zero_grad(set_to_none=True)

with torch.autocast(device_type="cuda", dtype=torch.float16):
    logits = model(x)
    loss = loss_fn(logits, y)

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

Then benchmark it.

Do not assume AMP helped because lower precision sounds faster.

Small workloads can be dominated by overhead.

Some operations do not benefit equally.

BF16 is often worth testing

On supported hardware:

with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
    logits = model(x)

The right dtype depends on hardware, model and numerical behavior.

Test rather than generalizing from another machine.

Inference should usually not build autograd graphs

Bad:

model.eval()
outputs = model(x)

eval() changes module behavior such as dropout and batch norm.

It does not disable gradient tracking.

Use:

model.eval()

with torch.inference_mode():
    outputs = model(x)

or torch.no_grad() where appropriate.

This can reduce memory and overhead.

`optimizer.zero_grad(set_to_none=True)`

For many workloads this is preferable to filling gradient buffers with zeros.

optimizer.zero_grad(set_to_none=True)

It can reduce memory operations and gives a useful semantic difference:

p.grad is None

means no gradient was produced since the reset.

That is also helpful for debugging.

Avoid accidental synchronization

Operations that force a CPU-visible scalar can synchronize CUDA work.

For example:

value = loss.item()

Doing this every iteration can introduce synchronization overhead.

Likewise excessive logging can serialize a pipeline.

Instead of:

for step in range(10000):
    loss = training_step()
    print(loss.item())

consider:

running_loss = 0.0

for step in range(10000):
    loss = training_step()
    running_loss += loss.detach()

    if step % 100 == 0:
        print((running_loss / 100).item())
        running_loss.zero_()

Be thoughtful about where CPU/GPU synchronization occurs.

Tiny kernels can make a powerful GPU look idle

A GPU wants enough parallel work.

If the model is tiny and the batch is tiny, kernel launch overhead and Python overhead can dominate.

You may see low utilization even though nothing is technically broken.

Try measuring throughput across batch sizes:

for batch_size in [1, 8, 32, 128, 512]:
    x = torch.randn(batch_size, features, device="cuda")

    ms = benchmark_cuda(lambda: model(x))
    throughput = batch_size / (ms / 1000)

    print({
        "batch_size": batch_size,
        "ms": ms,
        "examples_per_second": throughput,
    })

Latency may increase while throughput improves dramatically.

Now `torch.compile`

The wrong way to use it is:

model = torch.compile(model)

followed immediately by:

It is slower.

Compilation has upfront cost.

Compare steady-state execution.

eager_model = model
compiled_model = torch.compile(model)

Warm both appropriately.

Then benchmark the same workload.

A compile benchmark

import copy

base = MyModel().cuda().eval()

eager = copy.deepcopy(base)
compiled = torch.compile(copy.deepcopy(base))

x = torch.randn(128, 1024, device="cuda")

with torch.inference_mode():
    eager_ms = benchmark_cuda(lambda: eager(x), warmup=50, repeats=200)
    compiled_ms = benchmark_cuda(lambda: compiled(x), warmup=50, repeats=200)

print({
    "eager_ms": eager_ms,
    "compiled_ms": compiled_ms,
    "speedup": eager_ms / compiled_ms,
})

The speedup number is the interesting part.

Not the fact that compilation succeeded.

`torch.compile` works by capturing graphs

At a high level:

    flowchart LR
    A[Python frame] --> B[TorchDynamo capture]
    B --> C[FX graph]
    C --> D[AOT Autograd]
    D --> E[TorchInductor]
    E --> F[Optimized kernels/code]

The exact internals are richer than this, but this mental model is enough to debug most performance surprises.

Graph breaks matter

Suppose compiled code contains data-dependent Python control flow:

@torch.compile
def fn(x):
    if x.sum() > 0:
        return x * 2
    return x / 2

This kind of data-dependent operation can cause graph breaks.

Another common offender:

value = x.sum().item()

inside compiled code.

A graph break means the compiler captured a region, returned to Python for unsupported behavior, then may resume compiling later.

Too many breaks can destroy the expected speedup.

Use `fullgraph=True` as a debugging tool

compiled = torch.compile(model, fullgraph=True)

Now graph breaks are not silently tolerated as separate regions.

If the function cannot be captured as a full graph, you get a failure.

That makes fullgraph=True useful when trying to identify hidden breaks.

You do not necessarily need it for production execution.

Disable compilation to separate compiler bugs from model bugs

A simple debugging principle:

works eager?
works compiled?

If eager is broken, fix the model.

If eager works and compiled fails, you have narrowed the problem dramatically.

Current PyTorch also exposes compiler stance controls that can force eager execution for debugging without permanently rewriting every compile site.

Recompilation can erase your speedup

Compiled graphs are guarded by assumptions.

If inputs violate those assumptions, PyTorch may compile another graph.

One common reason is changing tensor shapes.

Example:

compiled = torch.compile(model)

compiled(torch.randn(32, 128, device="cuda"))
compiled(torch.randn(32, 256, device="cuda"))
compiled(torch.randn(32, 512, device="cuda"))

Depending on the graph and shape behavior, this may trigger recompilation.

Dynamic shapes exist specifically to reduce unnecessary specialization in suitable workloads.

compiled = torch.compile(model, dynamic=True)

Do not enable it blindly.

Use it when shape variability is actually part of the workload and recompilation is measurable.

Find recompilations with logging

PyTorch exposes compiler logging through environment settings.

A useful debugging pattern is to run the program with relevant TORCH_LOGS categories enabled, such as guards/recompilation-related logging.

The goal is to answer:

Did this input reuse the existing compiled graph?

or:

Did we compile again?

That distinction matters enormously for short-lived jobs and variable-shape workloads.

Dynamic shapes are a trade-off

Static specialization can produce highly optimized code.

Dynamic kernels can avoid repeated compilation.

Neither is universally better.

Benchmark your real shape distribution.

If production uses exactly one shape, dynamic flexibility may buy you nothing.

If production receives dozens of sequence lengths, repeated specialization may be expensive.

`mode="reduce-overhead"`

For some workloads with Python/kernel-launch overhead, try:

compiled = torch.compile(model, mode="reduce-overhead")

Also test the default.

And if you are exploring aggressive autotuning:

compiled = torch.compile(model, mode="max-autotune")

Again:

a mode name is not a benchmark result.

Measure.

Compile only the hot path

You do not have to compile an entire application.

@torch.compile
def model_step(model, x):
    return model(x)

Keep logging, file I/O, control-plane logic and debugging code outside the compiled region when that makes the graph cleaner.

A narrow hot path is often easier to reason about.

Compilation and debugging instrumentation can fight each other

This is an important practical issue.

Debugging code like:

print(x.shape)
print(x.mean().item())

inside a compiled region can change graph behavior or introduce breaks/synchronization.

Separate modes:

if debug:
    eager_model(...)
else:
    compiled_model(...)

When investigating correctness, clarity usually matters more than compiler performance.

Build a small benchmark matrix

from dataclasses import dataclass


@dataclass
class Result:
    name: str
    ms: float
    throughput: float
    peak_memory_mib: float


def evaluate(name, fn, batch_size):
    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()

    ms = benchmark_cuda(fn, warmup=30, repeats=100)

    peak = torch.cuda.max_memory_allocated() / 1024**2

    return Result(
        name=name,
        ms=ms,
        throughput=batch_size / (ms / 1000),
        peak_memory_mib=peak,
    )

Then compare:

results = []

results.append(evaluate("eager-fp32", eager_fn, batch_size))
results.append(evaluate("eager-amp", eager_amp_fn, batch_size))
results.append(evaluate("compile-fp32", compiled_fn, batch_size))
results.append(evaluate("compile-amp", compiled_amp_fn, batch_size))

for r in results:
    print(r)

This is much better than optimizing one feature at a time without preserving a baseline.

Memory is a performance dimension too

A faster configuration that uses 40% more VRAM may reduce the maximum batch size.

A slightly slower configuration that doubles batch size may produce higher overall throughput.

Record both:

latency
throughput
peak allocated memory
peak reserved memory

Performance is multi-dimensional.

A practical CUDA OOM decision tree

When you hit:

CUDA out of memory

work in this order.

    flowchart TD
    OOM[Out of Memory] --> REPRO[Can same batch trigger it again?]
    REPRO -->|Yes| PEAK[Measure peak allocation]
    REPRO -->|No| INTER[Check if intermittent]
    PEAK --> CLIMB[Memory climbing over iterations?]
    CLIMB -->|Yes| LEAK[Check retained references, .item misuse]
    CLIMB -->|No| PRESSURE[Reduce activation pressure]
    PRESSURE --> REDUCE[Smaller batch / shorter sequence / lower resolution]
    REDUCE --> CKPT[Activation checkpointing?]
    CKPT --> AMP[Test mixed precision]
    AMP --> FOOT[Inspect optimizer/model persistent footprint]
    FOOT --> SNAP[Use memory summary / snapshots]

1. Reproduce it reliably

Can the same batch trigger it again?

2. Measure peak allocation

torch.cuda.reset_peak_memory_stats()

3. Check whether memory climbs between iterations

If yes, suspect retained references.

4. Reduce activation pressure

Try:

smaller batch
shorter sequence
smaller image resolution
activation checkpointing
mixed precision

5. Inspect optimizer/model footprint

Large models may be dominated by persistent state.

6. Inspect allocator state

Use memory summaries/snapshots if fragmentation or allocation history is unclear.

Do not start by repeatedly calling empty_cache() and hoping.

A practical slow-training decision tree

    flowchart TD
    SLOW[Training is slow] --> WAIT[Is GPU waiting for data?]
    WAIT -->|Yes| DLFIX[Check DataLoader, num_workers]
    WAIT -->|No| TRANSFER[Expensive host-to-device transfer?]
    TRANSFER -->|Yes| PIN[Test pin_memory + non_blocking]
    TRANSFER -->|No| TINY[GPU workload tiny?]
    TINY -->|Yes| BS[Increase batch size]
    TINY -->|No| SYNC[Accidental synchronisation?]
    SYNC -->|Yes| REMOVE[Reduce .item .cpu calls]
    SYNC -->|No| DOMINATE[One operation dominates?]
    DOMINATE -->|Yes| PROF[Use profiler]
    DOMINATE -->|No| AMP_TEST[Test mixed precision]
    AMP_TEST --> COMP[Test torch.compile]
    COMP --> RECOMP[Recompiling constantly?]
    RECOMP -->|Yes| DYNAMIC[Add dynamic=True or fix shapes]
    RECOMP -->|No| DONE[Benchmark complete]

Is the GPU waiting for data?

Measure DataLoader wait time.

We covered this in Step 05.

Is host-to-device transfer expensive?

Test pinned memory and non-blocking transfers.

Is the GPU workload tiny?

Sweep batch size.

Are you synchronizing every iteration?

Look for:

.item()
.cpu()
print(...)

inside hot loops.

Is one operation dominating?

Use the profiler.

Does AMP help?

Benchmark it.

Does `torch.compile` help?

Benchmark steady state.

Is compilation recompiling constantly?

Inspect graph breaks, guards and shape variability.

A reusable environment report

When asking another programmer—or an LLM—to investigate performance, give it evidence.

import platform
import torch


def environment_report():
    report = {
        "python": platform.python_version(),
        "torch": torch.__version__,
        "cuda_available": torch.cuda.is_available(),
    }

    if torch.cuda.is_available():
        report.update({
            "cuda_runtime": torch.version.cuda,
            "device": torch.cuda.get_device_name(0),
            "capability": torch.cuda.get_device_capability(0),
            "total_memory_GiB": (
                torch.cuda.get_device_properties(0).total_memory / 1024**3
            ),
        })

    return report

print(environment_report())

Then provide:

model
input shapes
batch size
sequence length
precision
optimizer
peak memory
profiler summary
throughput eager
throughput compiled
compile logs if relevant

That is a dramatically better debugging prompt than:

PyTorch is slow. Fix it.

Performance regression tests

Performance can regress just like correctness.

You can keep lightweight benchmark assertions outside noisy CI environments.

For example:

baseline_ms = 4.2
current_ms = benchmark_cuda(lambda: model(x))

regression = current_ms / baseline_ms - 1

print(f"regression={regression:.1%}")

For stable dedicated hardware, you can define thresholds.

Do not use extremely tight thresholds on shared CI machines.

Profile memory by phase

def phase_memory(label):
    print(
        label,
        {
            "allocated": torch.cuda.memory_allocated() / 1024**2,
            "reserved": torch.cuda.memory_reserved() / 1024**2,
            "peak": torch.cuda.max_memory_allocated() / 1024**2,
        },
    )

optimizer.zero_grad(set_to_none=True)
phase_memory("start")

logits = model(x)
phase_memory("after_forward")

loss = loss_fn(logits, y)
loss.backward()
phase_memory("after_backward")

optimizer.step()
phase_memory("after_step")

This often reveals whether the peak comes from forward activations, backward or optimizer state initialization.

The first optimizer step can be special

Some optimizers allocate state lazily on the first update.

Therefore:

memory before optimizer.step()

may not represent steady-state training memory.

Always profile more than one step when optimizer state matters.

Avoid measuring debug mode as production performance

This series has emphasized instrumentation heavily.

Instrumentation itself costs time.

Hooks, anomaly detection, memory history, stack traces and profiler recording are debugging tools.

They are not necessarily settings you leave enabled for maximum throughput.

Use two modes:

forensic mode
performance mode

Forensic mode maximizes evidence.

Performance mode removes unnecessary instrumentation after the problem is understood.

A compact performance harness

Here is a useful starting point for real projects:

import time
import torch


class PerfHarness:
    def __init__(self, device="cuda"):
        self.device = torch.device(device)

    def reset_memory(self):
        if self.device.type == "cuda":
            torch.cuda.empty_cache()
            torch.cuda.reset_peak_memory_stats()

    def sync(self):
        if self.device.type == "cuda":
            torch.cuda.synchronize()

    def benchmark(self, fn, warmup=20, repeats=100):
        for _ in range(warmup):
            fn()

        self.sync()
        start = time.perf_counter()

        for _ in range(repeats):
            fn()

        self.sync()
        elapsed = time.perf_counter() - start

        result = {
            "ms": elapsed * 1000 / repeats,
        }

        if self.device.type == "cuda":
            result.update({
                "peak_allocated_MiB": (
                    torch.cuda.max_memory_allocated() / 1024**2
                ),
                "peak_reserved_MiB": (
                    torch.cuda.max_memory_reserved() / 1024**2
                ),
            })

        return result

Use it before and after every optimization.

The performance debugging ladder

When a PyTorch workload is slow, work downward through evidence.

    flowchart TD
    M1[Define the metric] --> M2[Create reproducible benchmark]
    M2 --> M3[Separate input wait / transfer / compute]
    M3 --> M4[Profile CPU and CUDA work]
    M4 --> M5[Inspect synchronization points]
    M5 --> M6[Measure peak memory]
    M6 --> M7[Test batch/shape scaling]
    M7 --> M8[Test mixed precision]
    M8 --> M9[Test torch.compile]
    M9 --> M10[Inspect graph breaks/recompiles]
    M10 --> M11[Repeat benchmark]
    style M11 fill:#c8e6c9

1. define the metric
2. create a reproducible benchmark
3. separate input wait / transfer / compute
4. profile CPU and CUDA work
5. inspect synchronization points
6. measure peak memory
7. test batch/shape scaling
8. test mixed precision
9. test torch.compile
10. inspect graph breaks/recompiles
11. repeat benchmark

Notice what is missing:

randomly change flags until GPU utilization looks higher

That is not engineering.

Three performance mistakes worth memorizing

Mistake 1: timing CUDA without synchronization

start = time.perf_counter()
y = model(x)
print(time.perf_counter() - start)

This can measure enqueue time instead of execution time.

Mistake 2: treating reserved memory as a leak

torch.cuda.memory_reserved()

High reserved memory is not by itself evidence of live tensors.

Look at allocated memory and ownership.

Mistake 3: assuming `torch.compile` means faster

Compilation is an optimization opportunity, not a performance guarantee.

Benchmark steady state and investigate graph breaks/recompilations when results are disappointing.

What `torch.compile` changes in the programmer’s job

This connects back to the philosophy of the previous article.

An LLM can generate:

model = torch.compile(model)

in less than a second.

That is not the hard part.

The hard part is answering:

Did it help?
Why?
On which shapes?
At what memory cost?
Did it recompile?
Where did it graph break?
Does it remain numerically correct?

The more code generation becomes automatic, the more valuable these questions become.

The programmer increasingly owns the measurement boundary.

Challenge 1: find the synchronization bug

Benchmark this:

for _ in range(1000):
    logits = model(x)
    losses.append(logits.mean().item())

Then compare against a version that avoids a CPU-visible scalar every iteration.

Measure the difference.

Challenge 2: create an artificial memory leak

Run:

saved = []

for _ in range(100):
    y = model(x)
    loss = y.square().mean()
    saved.append(loss)

    print(torch.cuda.memory_allocated())

Then replace:

saved.append(loss)

with:

saved.append(loss.item())

Observe the allocator behavior.

Challenge 3: force compile shape variation

Call a compiled function with a sequence of different input shapes.

Measure:

first-call latency
steady-state latency
subsequent new-shape latency

Then inspect compiler logs.

The point is to see compilation cost directly rather than treating torch.compile as a black box.

Challenge 4: build a performance table

For one real model, compare:

eager fp32
eager AMP
compiled fp32
compiled AMP

Record:

ms / step
examples or tokens / second
peak allocated MiB
peak reserved MiB

Do not pick a winner until you have all four columns.

Where we are now

We have moved through almost the entire stack:

✓ tensors and shapes
✓ autograd
✓ manual neural networks
✓ nn.Module and state
✓ DataLoader performance
✓ CNN geometry
✓ attention and masks
✓ model-not-learning debugging
✓ CUDA / profiling / torch.compile

There is one stage left.

In Step 10 we will stop examining the pieces separately and assemble them into the thing this series has been quietly preparing us to build:

a small language model from scratch in PyTorch

No pretrained model.

No AutoModelForCausalLM.from_pretrained(...).

We will build:

text
  ↓
tokenization
  ↓
embeddings
  ↓
positional information
  ↓
causal multi-head self-attention
  ↓
feed-forward network
  ↓
transformer blocks
  ↓
logits
  ↓
cross entropy
  ↓
next-token generation

And because of the previous nine stages, every major tensor in that model should be something we already know how to inspect when it goes wrong.