analyze-memory-snapshot

Files (8)

• SKILL.md
• checkers.py · evals/
• eval.yaml · evals/
• snapshot_step_005.json · evals/fixtures/leak_diff/
• snapshot_step_050.json · evals/fixtures/leak_diff/
• oom_snapshot.json · evals/fixtures/oom_snapshot/
• sample_snapshot.json · evals/fixtures/summarize_fixture_snapshot/
• analyze_snapshot.py · scripts/

SKILL.md (first lines)

# Analyze Memory Snapshot

Analyze PyTorch CUDA memory snapshots by inspecting allocator state and event
history. Always prefer evidence from the snapshot over generic memory advice.

## Workflow

1. Get the snapshot path if the user did not provide one. Accept `.pickle`,
   `.pkl`, `.json`, or a directory containing rank-local snapshots. If the user
   only has a screenshot from memory_viz, ask for the underlying snapshot file.
2. Load the snapshot locally. PyTorch `_dump_snapshot()` files are pickles, so
   only unpickle files the user provided or that already live in the workspace.
   JSON snapshots can be loaded directly.
3. If available, run this skill's helper for a first-pass summary:

   ```bash
   python scripts/analyze_snapshot.py /path/to/snapshot.pickle
   ```

   If there are multiple ranks, run it per rank and compare peaks, active bytes,
   inactive bytes, OOM events, and top stack traces.

   If the user has (or can capture) two snapshots from different points in the
   same run, diff them — this converts leak hunting from interpretation into a
   lookup of which allocation sites grew:

   ```bash
   python scripts/analyze_snapshot.py step100.pickle --diff step200.pickle
   ```

   For very large snapshots (hundreds of MB of trace history), the helper caps
   trace analysis at `--max-trace-events` (default 1M) and reports how many old
   events it dropped, so it stays usable without loading everything into your
   analysis.
4. Read the raw snapshot when the helper output is insufficient. Inspect:
   - `segments`: allocator segments returned by `cudaMalloc`
   - `blocks`: live, awaiting-free, and inactive pieces inside each segment
   - `device_traces`: chronological allocation events when memory history was
     enabled
   - `frames`: stack traces for allocation sites
5. Report the memory picture in terms of allocator semantics:
   - **Reserved memory**: sum of segment `total_size`
   - **Live tensor memory**: blocks in `active_allocated`
   - **Delayed frees**: blocks in `active_awaiting_free`, often stream-related
   - **Inactive cached memory**: `inactive` blocks that PyTorch can reuse
   - **Requested vs rounded size**: `requested_size` versus block `size`

## What To Diagnose

**Where memory is going**

Group live `active_allocated` blocks by the most useful user frame. Report the
largest stack traces first, with bytes, percent of live memory, block count, and
representative file/function/line. Separate parameters, optimizer state,
activations, temporary buffers, communication buffers, and user containers when
the stack names make that clear.

Always attribute memory to the user's model code, not allocator or framework
internals. Snapshot stacks are dominated by frames inside
`torch/nn/modules/*`, `torch/autograd/*`, and `torch/optim/*`; skip past those
to the first frame in the user's own files and name the module or line that
owns the memory (for example "38% of live memory is activations of
`model.decoder.layers[*].mlp`, allocated at `modeling.py:214`"). Only fall
back to torch-internal frames when no user frame exists in the stack.

**When memory is deallocated**

Use `device_traces` to distinguish `free_requested` from `free_completed`.
Pending or delayed frees can mean tensors are still needed by another stream or
that work has not reached the free point yet. In multi-step traces, compare
repeated allocation sites with their matching free events: allocation sites that
appear every step without corresponding frees are stronger retention evidence
than one-time model or optimizer allocations. If traces include OOM events,
report the failed allocation size, `device_free`, OOM stack frame, and which live
allocation sites were holding memory immediately before the OOM.

**Leak and retention patterns**

Look for active bytes that grow across steps, repeated active allocations from
the same stack that are not freed, or a sawtooth where many tensors accumulate
and are released together. The strongest evidence is a two-snapshot diff
(`--diff` in the helper): an allocation site whose live bytes and block count
both grow between snapshots taken k steps apart is leaking or retaining.

Recognize these named patterns by their snapshot signature and give the
matching fix:

- **Retained autograd graph**: live activation blocks from step N still
  present at step N+k, often alongside a loss or output tensor per step.
  Typical cause is logging or accumulating a tensor that still holds the
  graph — fix with `loss.detach()` (or `.item()`) before appending, or find
  the `retain_graph=True` that keeps graphs alive.
- **Growing Python container**: same user-code allocation site with block
  count increasing linearly in steps. Fix by detaching and moving to CPU
  before appending, or bounding the container.
- **First-step optimizer state spike**: large `torch/optim/*`-attributed
  allocations appearing once at the first `optimizer.step()` and then stable.
  This is expected retention, not a leak — say so explicitly; the fix for OOM
  here is a smaller model/batch, a lower-memory optimizer, or sharding, not
  leak hunting.
- **Activation memory vs. checkpointing**: live memory dominated by forward
  activations that all free during backward (sawtooth in `device_traces`).
  Recommend activation checkpointing (`torch.utils.checkpoint`) and quantify
  the tradeoff from the snapshot: the activation bytes that would be freed
  versus recompute cost.
- **Gradient memory under DDP vs. FSDP**: full-size gradient and parameter
  blocks per rank suggest DDP replication; if the user is memory-bound and
  gradients+params+optimizer state dominate live memory, mention
  FSDP/ZeRO-style sharding as the structural fix.
- **Python reference cycles**: frees that only happen in bursts when GC runs.
  Recommend `from torch.utils.viz._cycles import warn_tensor_cycles` and call
  `warn_tensor_cycles()` during initialization to find object cycles that keep
  CUDA tensors alive.

**Fragmentation and allocator behavior**

Do not equate `nvidia-smi` usage with live tensors. PyTorch uses a caching CUDA
allocator, so reserved memory can be much larger than allocated tensor memory.
Large inactive blocks are reusable by PyTorch; `empty_cache()` may return fully

From meta-pytorch/skills, BSD-3-Clause.

More from meta-pytorch/skills