PyTorch MemoryViz

PyTorch MemoryViz (VS Code Extension)

Explore PyTorch CUDA memory traces directly in VS Code.

This extension embeds the upstream PyTorch MemoryViz UI inside a webview and bundles all dependencies. You can open .pickle/.pkl or profiler .json traces.

See https://pytorch.org/blog/understanding-gpu-memory-1/ for more information how to record the traces.

Features

• Offline MemoryViz UI
• Open traces from the Explorer (context menu) or “Open With → PyTorch MemoryViz”
• Remote-friendly: works over SSH/WSL; files are read on the remote, not the local machine

How to use

1. Install “PyTorch MemoryViz” from the Marketplace.
2. Open your workspace that contains trace files (.pickle, .pkl, .json).
3. Open a trace in any of these ways:
   • Right-click a file in Explorer → “Open with PyTorch MemoryViz”
   • Right-click → “Open With…” → select “PyTorch MemoryViz”
   • From the MemoryViz tab’s toolbar, choose “Choose files (remote)” to pick traces from your remote workspace, or “Choose files (local)” to pick from your machine
4. Interact with the chart: zoom, pan, and click a block to inspect the stack.

Tip: To generate traces, see PyTorch documentation for torch.cuda.memory._snapshot() or torch.profiler.profile(..., profile_memory=True).

Requirements

• VS Code 1.80+
• Trace files generated by PyTorch (CUDA memory snapshot .pickle/.pkl, or profiler .json)

Privacy & Security

• All trace files are processed locally within your VS Code session

Feedback / Issues

• Open issues and feature requests at the GitHub repository.

For development/build instructions, see CONTRIBUTING.md.