skills.rest
NVIDIA-NeMo
NVIDIA-NeMoOFFICIAL
Explore all skills by this author →
WebsiteGitHub

nemo-mbridge-memory-snapshot-analysis

Analyze and compare PyTorch CUDA memory snapshots from Megatron Bridge training runs.

896|481|Updated May 21, 2025
Search Tags:#memory profiling#CUDA#pyTorch#megatron bridge#OOM debugging#GPU memory#snapshot analysis
One-click install
npx skills add https://github.com/NVIDIA-NeMo/Megatron-Bridge --skill nemo-mbridge-memory-snapshot-analysis
Or copy as Structured Prompt for Agent▼
Please help me install this Agent Skill.
Skill: nemo-mbridge-memory-snapshot-analysis
Source: https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/skills/nemo-mbridge-memory-snapshot-analysis
Command: npx skills add https://github.com/NVIDIA-NeMo/Megatron-Bridge --skill nemo-mbridge-memory-snapshot-analysis

SYSTEM DOCUMENTATION & REQUIREMENTS

💡 This Skill includes scripts (resource) components.

What problem does it solve?

Debugging GPU out-of-memory failures and peak-memory regressions in Megatron Bridge training requires interpreting raw PyTorch CUDA memory snapshot pickles, which are opaque event logs that cannot be compared across runs without tooling.

Core Features & Use Cases

  • Snapshot Overview: Parse a snapshot pickle to inspect allocator settings, segment totals, fragmentation, baseline memory, and top allocation sources by throughput.
  • Two-Run Comparison: Diff two snapshots side-by-side and attribute the peak memory gap to specific source code locations, with automatic fallback to full-trace replay when step annotations are absent.
  • Interactive Timeline: Generate a standalone HTML Plotly timeline that overlays two runs on one axis, with wall-clock and event-index x-axis modes.
  • Use Case: You ran the same recipe with two parallelism configs and run B peaks 20 GiB higher; run compare_snapshots.py on both rank-0 pickles to see exactly which allocations (e.g., router_forward, mlp_forward) account for the difference.

Quick Start

Ask the assistant to analyze your Megatron Bridge memory snapshot pickle, for example to compare snapshot_0.pickle files from two runs and explain which allocations cause the peak memory difference.

Frequently Asked Questions about nemo-mbridge-memory-snapshot-analysis

High-intent search queries and answers about installing and using this skill.

FAQPage Schema
How do I compare two PyTorch CUDA memory snapshots?▼

Run compare_snapshots.py with the two snapshot pickle files to get an allocator-settings diff, segments baseline comparison, and a drill-down attributing the peak difference to source locations. Use --frame-depth 2 for finer source grouping and compare the same rank from each run.

How do I plot GPU memory usage over time from a memory snapshot?▼

Run plot_timeline.py with one or two snapshot pickles to generate a standalone HTML file with an interactive Plotly timeline. It supports wall-clock and event-index x-axis modes and can overlay two runs on one axis to show where they diverge.

Why does my memory snapshot have no ProfilerStep annotations?▼

ProfilerStep markers come from torch.profiler step boundaries, which require use_pytorch_profiler=True in addition to record_memory_history=True. The snapshot is still valid: compare_snapshots.py falls back to full-trace replay, and plot_timeline.py and replay_to_time.py work without step markers.

What dependencies are needed to analyze CUDA memory snapshots?▼

The analysis scripts use only the Python standard library and run on any python3 back to 3.9, with no GPU or virtualenv required. The generated HTML timeline loads Plotly.js from a CDN in the browser, adding nothing to the Python environment.

Is it safe to load memory snapshot pickle files from others?▼

No, unpickling executes arbitrary code embedded in the file, so only analyze snapshots from training runs you trust. Snapshots also embed absolute paths and stack frames from the source machine, so review them before sharing externally.

Why do snapshot numbers not match nvidia-smi?▼

A high unmatched-free count means the 100k-entry trace cap dropped the run's beginning, making baseline_at_start unreliable. Also verify you are inspecting the rank that actually peaked, since profile_ranks defaults to rank 0 and pipeline stages differ in memory use.

Related Skills

Explore More Software Engineering Skills→
NVIDIANVIDIA
2.8k

moe-optimization-workflow

Diagnose and optimize MoE training bottlenecks across memory, compute, and communication.

Official
Advanced
yo-stevenyo-steven

perf-memory-tuning

Configure PyTorch CUDA allocator expandable segments to reduce GPU memory fragmentation in Megatron Bridge training.

Community
Intermediate
Clay-HHKClay-HHK
2

tensorboard

Visualize training metrics and debug models with TensorBoard event logs.

Community
Intermediate
Tsin418Tsin418

gate-info-coinanalysis

Generate a five-section markdown market analysis report for one crypto ticker.

Community
Advanced