Use with AI · Chapter Application

Audit PyTorch CUDA Memory and Training Performance

Diagnose CUDA OOM, low GPU utilization and slow PyTorch training by measuring memory, synchronization, input wait and compute before applying optimizations.

Chapter Application PyTorch: Zero to Hero PyTorch CUDA training and inference code Advanced
Read the source chapter → View the learning path →

How to use this

  1. Open a repository-aware AI assistant.
  2. Give it access to the repository or files you want reviewed.
  3. Copy the prompt below and run it unchanged first.
  4. Use the evidence it finds to decide what to inspect or change next.
PromptCopy and run against your own project
You are auditing a PyTorch repository for CUDA memory problems or poor training performance.

Do not start by reducing batch size, enabling torch.compile, adding mixed precision, or changing the model. First determine which resource is actually limiting the workload.

Inspect the training/inference loop, model construction, optimizer, data transfer, precision policy, logging/metrics and any profiling code.

Work in this order:

1. Establish workload context:
   - device(s)
   - model/parameter size where inferable
   - optimizer
   - batch and sequence/image sizes
   - dtype / autocast policy
   - gradient accumulation
   - checkpointing
   - torch.compile usage

2. Reconstruct GPU memory ownership:
   - parameters
   - gradients
   - optimizer state
   - activations
   - temporary tensors
   - cached allocator memory
   - tensors accidentally retained across iterations

3. Inspect CUDA OOM risks:
   - references preserving computation graphs
   - loss/output tensors appended without detach
   - unnecessary activation retention
   - oversized temporary tensors
   - sequence/batch growth
   - fragmentation versus true capacity pressure
   - optimizer-state cost
   - validation/inference accidentally building graphs

4. Inspect performance risks:
   - CPU/data starvation
   - frequent cuda synchronize behavior
   - .item() or CPU transfers in the hot loop
   - excessive tiny kernels
   - repeated device allocations/copies
   - non_blocking used without the conditions required for overlap
   - logging/profiling that forces synchronization
   - poor batch sizing
   - compilation graph breaks or recompilation
   - mixed precision that is absent, misapplied or numerically unsafe

5. Separate the diagnosis into:
   - memory capacity bound
   - allocator/retention defect
   - data/input bound
   - CPU/synchronization bound
   - GPU compute bound
   - kernel/launch overhead bound
   - compile/runtime graph issue
   - insufficient evidence

6. For every proposed optimization, state:
   - evidence supporting the bottleneck
   - expected mechanism of improvement
   - likely trade-off
   - measurement that would confirm success

Output:

## Resource map
Summarize workload, precision and major memory consumers.

## Findings
Rank confirmed defects separately from optimization opportunities.

## Measurement plan
Specify appropriate PyTorch profiler/timing/memory observations needed to discriminate competing hypotheses.

## Minimal interventions
Recommend the smallest experiment first. Examples may include fixing retained graphs, changing transfer behavior, precision, batch size, checkpointing or compilation — but only when justified.

## Verification
Compare before/after peak allocated memory, peak reserved memory, step time, throughput and GPU utilization where available.

Never claim torch.compile, AMP, larger batches or more workers will improve performance without measuring the result in this workload.