Memory snapshot: python-only stacks for fast dumps + configurable max_entries#3628
Open
SherlockNoMad wants to merge 1 commit into
Open
Memory snapshot: python-only stacks for fast dumps + configurable max_entries#3628SherlockNoMad wants to merge 1 commit into
SherlockNoMad wants to merge 1 commit into
Conversation
…le max_entries `MemoryProfiler` calls `torch.cuda.memory._record_memory_history` with torch's default `stacks="all"`, which captures and symbolizes C++ frames. Symbolization happens at dump time and is extremely slow for this workload: on DeepSeek-v3 16B (8xH100) a single `_snapshot()` dump took **~461 s**, which dominates the run and makes `--profiler.enable_memory_snapshot` impractical. Switch to `stacks="python"` (Python frames only): the dump drops to **~4 s** (~100x faster) while still giving actionable per-allocation Python stacks in the PyTorch memory visualizer. Also make the history cap configurable: replace the hardcoded `MEMORY_SNAPSHOT_MAX_ENTRIES = 100000` with a `Profiler.Config` field `memory_snapshot_max_entries` (default 1_000_000). 100k only covered the last ~2-3 steps; python-only stacks make a larger buffer cheap, and exposing it lets runs bound it explicitly (torch's own default is effectively unbounded). Test plan: - deepseek_v3 16B, 8xH100, `--profiler.enable_memory_snapshot`: snapshot dump ~461 s (stacks="all") -> ~4 s (stacks="python"); the .pickle opens in the PyTorch memory visualizer with Python stacks.
SherlockNoMad
requested review from
fegin,
tianyu-l,
wconstab and
wwwjn
as code owners
meta-cla
Bot
added
the
CLA Signed
tianyu-l
approved these changes
Jun 11, 2026
yushangdi
approved these changes
Jun 11, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
MemoryProfilercallstorch.cuda.memory._record_memory_historywith torch'sdefault
stacks="all", which captures and symbolizes C++ frames.Symbolization runs at dump time and is pathologically slow for this workload.
This PR:
stacks="python") for memory snapshots.MEMORY_SNAPSHOT_MAX_ENTRIES = 100000with aProfiler.Configfieldmemory_snapshot_max_entries(default1_000_000),overridable via
--profiler.memory_snapshot_max_entries.Why / evidence
On DeepSeek-v3 16B (8×H100,
--profiler.enable_memory_snapshot), a single_snapshot()dump:stacks"all"(default)"python"(this PR)~100× faster. With
stacks="all"the dump dominated the run and made--profiler.enable_memory_snapshotimpractical;stacks="python"still producesactionable per-allocation Python stacks in the PyTorch memory visualizer.
On
max_entries: the old 100k ring buffer only retained the last ~2–3 steps ofhistory. With the much cheaper python-only stacks a larger buffer is affordable,
so the default is raised to 1M and exposed as config. (Torch's own default is
effectively unbounded —
sys.maxsize— which is unsafe for long runs, so we keepan explicit, configurable cap.)
Test plan
deepseek_v3 16B, 8×H100,
--profiler.enable_memory_snapshot:stacks="all") → ~4 s (stacks="python").pickleopens in the PyTorch memory visualizer with Python stacksmax_entries