Nemo Mbridge Memory Snapshot Analysis

Memory Snapshot Analysis

Stable documentation: @docs/training/profiling.md Related skills: @skills/nemo-mbridge-perf-memory-tuning/SKILL.md (how to fix the memory problem this skill helps you find)

What It Is

Bridge writes PyTorch CUDA memory snapshots when record_memory_history=True is set on ProfilingConfig. Each snapshot pickle contains:

The bundled scripts replay device_traces to reconstruct memory over time, group allocations by source location, and diff two snapshots to explain which tensors account for a memory difference.

PyTorch's own memory_viz also renders a timeline from these files and is worth using for a single snapshot. What it cannot do is put two runs on one axis — which is the whole question when you are asking why config B peaks higher than config A. plot_timeline.py overlays them, and compare_snapshots.py then attributes the gap to specific allocation sites. Add scriptable JSON output on top and that is the reason this skill exists alongside memory_viz.

No dependencies — Python stdlib only, and they run on any python3 back to 3.9, so a login node or laptop system interpreter is fine (macOS ships 3.9; from __future__ import annotations keeps the newer type syntax from being evaluated at import). No virtualenv, no uv sync, no GPU. The HTML plot loads Plotly.js from a CDN in the browser, so it adds nothing to the Python environment.

Pick a Script

plot_timeline.py is the only script that writes a file; the rest print to stdout. It also takes two snapshots and overlays them, which is usually the fastest way to see where two runs diverge before asking compare_snapshots.py why.

A good default loop for a memory regression: overview both files, overlay them on a timeline, read off the time where they split, then drill in at that moment.

python3 skills/nemo-mbridge-memory-snapshot-analysis/scripts/plot_timeline.py A_0.pickle B_0.pickle --labels A B -o compare.html
python3 skills/nemo-mbridge-memory-snapshot-analysis/scripts/replay_to_time.py A_0.pickle B_0.pickle --time 83.0

How Bridge Produces These Files

Read this before interpreting any numbers. Bridge's capture path differs from a hand-rolled _record_memory_history() call in ways that change what the numbers mean.

from megatron.bridge.training.config import ProfilingConfig

cfg.profiling = ProfilingConfig(
    record_memory_history=True,
    memory_snapshot_path="/nemo_run/snapshot.pickle",
    profile_step_end=15,
    profile_ranks=[0],
)

Five consequences worth internalizing:

1. external_annotations is populated, but never with ProfilerStep. This is the single most misleading thing about these files. The list is not empty — a real Bridge pretrain snapshot carries thousands of entries, but all of them are allocator-level record_function markers:

  1448  nccl:all_gather_into_tensor_coalesced
   588  nccl:all_reduce
   504  nccl:reduce_scatter_tensor_coalesced
   462  nccl:_all_gather_base
    96  enumerate(DataLoader)#_SingleProcessDataLoaderIter.__next__
    44  Optimizer.step#FusedAdam.step
    24  CustomFSDP.forward

ProfilerStep#N markers come from torch.profiler's step boundaries, and record_memory_history and use_pytorch_profiler are independent flags. So a snapshot can be rich in annotations and still have zero step boundaries. Consequences:

• replay_step.py cannot run — it needs step boundaries, and exits 1.
• compare_snapshots.py automatically falls back to full-trace replay.
• plot_timeline.py and replay_to_time.py are unaffected — they never needed steps.
• The NCCL/optimizer/dataloader markers are still useful for orienting yourself in the timeline, and replay_step.py surfaces them per step when step boundaries do exist.

To get per-step drill-down, set use_pytorch_profiler=True alongside record_memory_history=True. Note finalize() forbids combining it with use_nsys_profiler.

2. Where the trace starts depends on the Bridge version. Since fix: record CUDA memory history before snapshot (9e9e58b9f, 2026-04-23), start_memory_history_recording() runs at setup.py:330, before _build_distributed_model(). Weights and optimizer state are then allocated inside the trace, so they appear as ordinary alloc events attributable to their construction site and baseline_at_start is small.

Snapshots captured before that commit — or by an external callback that enables recording at train start — begin after model construction instead, and their baseline_at_start is the whole model. A verified pre-fix capture shows a 26.29 GiB baseline for exactly this reason. Check the earliest frames in the trace to tell which regime you are in.

3. The trace is capped at 100k entries. trace_alloc_max_entries=100_000 is hardcoded in start_memory_history_recording(). PyTorch keeps the most recent entries once the cap is hit, so a long enough run silently loses its beginning. Symptom: a large unmatched frees count and a baseline_at_start that no longer reconciles with nvidia-smi. If you need the start of the run, dump earlier by lowering profile_step_end. Snapshots captured outside this code path are not subject to the cap and can be much larger — a real 113-second capture contained 184,679 events.

4. One dump per rank, at profile_step_end. The snapshot is written when iteration == profile_step_end, with the rank appended before the extension. memory_snapshot_path="/nemo_run/snapshot.pickle" with profile_ranks=[0] yields /nemo_run/snapshot_0.pickle. Compare like for like — rank 0 against rank 0.

5. OOM produces its own snapshot automatically. An out-of-memory observer is attached at capture time and dumps snapshot_oom_rank-{N}.pickle at the moment of failure. This file is usually the most valuable one you have after a crash: it captures exactly what was live when the allocator gave up. Analyze it with parse_snapshot.py and replay_to_time.py.

Workflows

Scripts live in skills/nemo-mbridge-memory-snapshot-analysis/scripts/.

Load trusted snapshots only. Every script below starts by unpickling the file, and pickle.load executes arbitrary code embedded in it. Analyze snapshots produced by your own training runs or by someone you trust — never a pickle from an untrusted issue attachment, bucket, or download. There is no safe-mode parse: the format requires full deserialization.

Snapshots also carry absolute paths and stack frames from the machine that produced them, so check before sharing one outside your organization.

1. Single snapshot overview

Start here to see what a file actually contains — especially whether it has step annotations and how many events survived the ring buffer.

python3 skills/nemo-mbridge-memory-snapshot-analysis/scripts/parse_snapshot.py snapshot_0.pickle

Shows allocator settings, segments totals and fragmentation, baseline memory, detected steps, event counts by action, and the top 15 allocation sources by throughput. Add --json for machine-readable output.

2. Compare two runs

The primary use case: explaining why config B uses more memory than config A.

python3 skills/nemo-mbridge-memory-snapshot-analysis/scripts/compare_snapshots.py A_0.pickle B_0.pickle
python3 skills/nemo-mbridge-memory-snapshot-analysis/scripts/compare_snapshots.py A_0.pickle B_0.pickle --top 20 --frame-depth 2

Produces an allocator-settings diff, a side-by-side segments baseline, a per-step table (or a full-trace summary when annotations are absent), a drill-down at the largest divergence grouped by source and sorted by delta, and automatic observations.

3. Visualize the timeline

python3 skills/nemo-mbridge-memory-snapshot-analysis/scripts/plot_timeline.py snapshot_0.pickle -o timeline.html
python3 skills/nemo-mbridge-memory-snapshot-analysis/scripts/plot_timeline.py A_0.pickle B_0.pickle --labels "TP2" "TP4" -o compare.html

Standalone HTML with a WebGL Plotly timeline. Two toggles matter:

• X-axis: "Wall Clock" vs "Event Index". Event-index mode matches memory_viz, stretching busy iterations and compressing idle gaps so sawtooth allocation patterns become visible.
• Y-axis: "Absolute" (baseline + delta) vs "Delta from trace start".

Wall-clock mode buckets by time (default 1 ms, --precision) preserving min/max per bucket, so busy traces produce more points.

4. Drill into a specific moment

The natural follow-up to the plot: spot a divergence visually, then find out what is responsible.

python3 skills/nemo-mbridge-memory-snapshot-analysis/scripts/replay_to_time.py snapshot_0.pickle --time 83.0
python3 skills/nemo-mbridge-memory-snapshot-analysis/scripts/replay_to_time.py A_0.pickle B_0.pickle --time 83.0 --frame-depth 2

Replays to the given elapsed second and reports baseline/delta/absolute memory, live allocation counts, and live allocations grouped by source. In comparison mode it sorts by absolute delta and separately totals only-in-A and only-in-B sources.

5. Per-step drill-down (requires annotations)

Only works when the snapshot was captured with use_pytorch_profiler=True.

python3 skills/nemo-mbridge-memory-snapshot-analysis/scripts/replay_step.py snapshot_0.pickle --step 7
python3 skills/nemo-mbridge-memory-snapshot-analysis/scripts/replay_step.py snapshot_0.pickle --all-steps

Per step: alloc/free counts, throughput, peak delta and absolute peak, live allocations at peak grouped by source, and active annotation phases.

Example Output

--- Segments Baseline ---
                            Snapshot A      Snapshot B           Delta
  ────────────────────  ──────────────  ──────────────  ──────────────
  Segments                          24              24               0
  Reserved (total)           53.21 GiB       73.49 GiB  +    20.27 GiB
  Allocated                  45.10 GiB       62.28 GiB  +    17.18 GiB
  Active                     45.10 GiB       62.28 GiB  +    17.18 GiB
  Awaiting free                    0 B             0 B            0 B
  Inactive (reusable)         8.12 GiB       11.21 GiB  +     3.09 GiB
  Baseline at start          39.61 GiB       52.77 GiB  +    13.15 GiB

--- Full-Trace Summary (no ProfilerStep annotations) ---
                                           Snapshot A      Snapshot B      Delta(B-A)
  ───────────────────────────────────  ──────────────  ──────────────  ──────────────
  Peak delta (from trace start)              9.65 GiB       19.15 GiB  +     9.50 GiB
  Absolute peak (baseline+delta)            49.26 GiB       71.92 GiB  +    22.65 GiB
  Alloc count                                     760           1,840  +        1,080

--- Drill-Down: Full Trace (peak diff: +9.50 GiB) ---
  Source                                            Peak A        Peak B    Delta(B-A)
  ────────────────────────────────────────────  ──────────  ──────────  ────────────
  router_forward@router.py:145                    1.81 GiB    4.51 GiB  +   2.70 GiB
  mlp_forward@mlp.py:301                          1.59 GiB    4.28 GiB  +   2.69 GiB
  attention_forward@attention.py:88               1.77 GiB    3.99 GiB  +   2.21 GiB

Interpreting Results

All sizes are binary — KiB/MiB/GiB, 1024-based — so they line up directly with nvidia-smi and torch.cuda.memory_allocated. A decimal "GB" reads about 7.4% larger for the same bytes, which is enough to make a correct cross-check look like a discrepancy.

Segments vs device traces

segments is the allocator's state at dump time — use it for total memory accounting. device_traces is the event log — use it to understand what changes over time. The two are cross-referenced to derive baseline_at_start, which is what turns relative deltas into real GPU memory figures. A timeline alone cannot tell you how much memory the model was actually using.

Reserved, allocated, active, inactive

The four segment numbers are distinct and easy to conflate:

Fragmentation is reported as inactive / reserved. It is deliberately computed from active, not allocated: awaiting-free blocks are not reusable yet and are not fragmentation, so counting them as such over-reports the problem. When Awaiting free is non-zero, expect the two definitions to disagree.

Absolute peak

absolute_peak = baseline_at_start + cumulative_delta_at_step_start + step_peak_delta. Compare this against nvidia-smi or the memory/ metrics Bridge logs to TensorBoard to sanity-check that the trace covers what you think it covers.

Per-step replays seed their live set with whatever was already live when the step opened, so the source table accounts for memory carried into the step (weights, optimizer state, graph pools, activations held across the boundary) and not just what the step itself allocated. On a real annotated capture this recovered about 3.4 GiB per step that the table previously omitted.

The table will not sum to absolute_peak, and the gap is not a single clean quantity. The table shows allocations the trace can see and attribute to a stack frame; absolute_peak also includes memory that predates the trace and has no frames, and the two are related through step_start_delta, which nets frees of pre-trace memory against later allocations. Treat the table as "what is attributable and how it ranks", not as a decomposition of the peak. Use the segments numbers for total accounting.

Unmatched frees

free_completed events with no matching alloc are allocations that predate the visible trace. A small count is normal and harmless. A large count means the trace does not reach back to the allocations it is freeing — either the 100k cap dropped the beginning of the run, or recording started after model construction (see point 2 above). Either way, baseline_at_start is doing more guessing than measuring, so cross-check it against nvidia-smi before trusting absolute numbers.

Source grouping

Allocations are grouped by their first non-internal stack frame, rendered as function@file.py:line. PyTorch, CUDA, Hydra, and stdlib frames are skipped. When results look over-aggregated (everything attributed to one wrapper), raise --frame-depth 2 to split by two levels of user code.

CUDA graphs

A snapshot with dramatically fewer allocs indicates CUDA graph replay. Graphs pre-allocate at capture time and reuse on replay, so those allocations stop appearing in device_traces. The memory did not disappear — it moved into the graph's pinned pool, which is why enabling graphs can raise peak memory even as the trace gets quieter.

A matched pair of Bridge captures — same model and config, graphs off vs on — shows both halves of this:

The trace got 4.5x quieter and peak memory went up 34 GiB. Running compare_snapshots.py on the pair attributes most of that increase to the graph memory pool rather than to any model tensor — which is the signal you want, and the reason to diff rather than eyeball a single file. See @skills/nemo-mbridge-perf-cuda-graphs/SKILL.md.

Options

Diagnostics go to stderr via logging; report bodies go to stdout, so --json output pipes cleanly.

Troubleshooting

"no ProfilerStep annotations found" — Expected for a Bridge run with record_memory_history=True and no PyTorch profiler, and not a sign of a broken snapshot: the file still has NCCL and optimizer annotations, just no step boundaries. Use compare_snapshots.py (auto-falls back), plot_timeline.py, or replay_to_time.py. Add use_pytorch_profiler=True on the next run if you need per-step analysis.

"no device traces found" — The pickle has only segments. Either record_memory_history was False, or the snapshot came from somewhere other than Bridge's capture path. parse_snapshot.py still summarizes segments.

Numbers don't match nvidia-smi — Check the unmatched-free count first; a wrapped ring buffer makes baseline_at_start meaningless. Also confirm you are looking at the rank that actually peaked — with pipeline parallelism the first PP stage usually holds the most memory, and profile_ranks defaults to [0].

Steps don't line up between two snapshots — Matching is by step number, so different warmup counts pair the wrong steps. Force alignment with --step N.

Large files — The whole pickle is deserialized into memory; budget roughly 10-20x the file size in RAM. Loading progress is logged to stderr.