bench_parallel: add opt-in --memory-snapshot (CUDA memory snapshot) for full-graph models#76
Conversation
| os.makedirs(os.path.dirname(memory_snapshot) or ".", exist_ok=True) | ||
| with gpu_setup_lock(): | ||
| with torch.no_grad(): | ||
| for _ in range(3): |
There was a problem hiding this comment.
_record_memory_history starts AFTER 3 warmup iterations, so the snapshot captures steady-state only — compile-time + first-run allocations (often the actual peak!) are invisible. If the tool's purpose is debugging OOMs, peak usually happens during compile/first iteration. Suggest either a flag to record from process start, or explicitly documenting 'steady-state snapshot' in the help so nobody reads it as peak-memory forensics. Also: max_entries=100000 silently drops older events once exceeded — big models can blow through that; worth surfacing (or making it a flag).
| memory_snapshot = None | ||
| if MEMORY_SNAPSHOT: | ||
| try: | ||
| _rel = os.path.relpath(repro_path, {args_dict.get("models_root", "repros/models")!r}) |
There was a problem hiding this comment.
Asymmetry with #75: this derives _rel from args_dict['models_root'] (new plumbing) while #75 derives it from root + 'repros/models' inline. Whichever lands second should adopt the other's mechanism so the two artifact trees mirror identically (profiles//... and memory_snapshots//... diverge today if --models-root is non-default). models_root plumbing (this PR's approach) looks like the better one.
|
|
||
| if args.peak_memory: | ||
| # --peak-memory: 1 worker per GPU. | ||
| if args.peak_memory or args.memory_snapshot: |
There was a problem hiding this comment.
Consolidation with #75: this PR and #75 instantiate the same scaffold twice — this 1-worker/GPU gating (same line; textual conflict for whichever merges second), TAG plumbing, the relpath+fallback artifact-path derivation, the warmup→lock→instrumented-run→export→try/except block shape, and the near-identical injection test. Only the ~15-line core (profiler trace vs memory snapshot) differs. Suggest either stacking this on #75 (or merging both into one 'full-graph debug artifacts' PR), or having whichever lands second extract a shared capture-run helper (warmup/lock/path/makedirs/except, with the models_root-based path derivation from this PR as the single mechanism) so --profile, --memory-snapshot, and the existing --peak-memory become thin instantiations. Also note: if both flags are set, each runs its own extra warmup+instrumented pass — they should compose.
…or full-graph models Add --memory-snapshot (+ --memory-snapshot-dir) to the --full-graphs path. For each model, dump a CUDA memory snapshot (torch.cuda.memory._record_memory_ history + _dump_snapshot) to <dir>/<tag>/<suite>/<mode>/<model>/<graph>.pickle for pytorch.org/memory_viz, recorded in a separate run from the do_bench timing. Output is written under a per-<tag> subdir (tag defaults to "latest") so tagged A/B runs don't clobber each other. Path base honors --models-root. Parallel across GPUs but pinned to 1 worker/GPU (reuses the --peak-memory guard). The snapshot run is exception-guarded so an OOM/IO failure can't discard the already-measured timing, and recording is always disabled again in a finally so it can't leak into the next model. Records the snapshot path in the result JSON. Opt-in and additive: off by default the result schema is unchanged. Full-graph path only. Adds a worker-script test.
c87296c to
5afa67a
Compare
|
Thanks for the review! pushed fixes to updated the --memory-snapshot help text to note it's a steady-state snapshot |
2 participants
Opt-in
--memory-snapshotflag on the--full-graphspath dumps a CUDA memory snapshot per model (torch.cuda.memory._record_memory_history+_dump_snapshot) to<dir>/<tag>/<suite>/<mode>/<model>/<graph>.pickle, loadable in pytorch.org/memory_viz. This is the "Level 3" memory profiler (full allocation timeline + stacks + fragmentation), complementing--peak-memory(the number).do_benchtiming loop.--tagsubdir (defaultlatest) so tagged A/B runs don't clobber; path base honors--models-root.--peak-memoryguard; per-process/per-GPU, verified with a 2-GPU sweep → distinct collision-free pickles).finallyso it can't leak into the next model.