mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 00:21:07 +01:00
1610ea8ef8
Summary: Adds comprehensive memory instrumentation to the CUDA caching memory allocator. # Counters Added comprehensive instrumentation for the following stats: - Allocation requests (`allocation`) - Allocated memory (`allocated_bytes`) - Reserved segments from cudaMalloc (`segment`) - Reserved memory (`reserved_bytes`) - Active memory blocks (`active`) - Active memory (`active_bytes`) - Inactive, non-releasable blocks (`inactive_split`) - Inactive, non-releasable memory (`inactive_split_bytes`) - Number of failed cudaMalloc calls that result in a cache flush and retry (`cuda_malloc_retries`) - Number of OOMs (`num_ooms`) Except for the last two, these stats are segmented between all memory, large blocks, and small blocks. Along with the current value of each stat, historical counts of allocs/frees as well as peak usage are tracked by the allocator. # Snapshots Added the capability to get a "memory snapshot" – that is, to generate a complete dump of the allocator block/segment state. # Implementation: major changes - Added `torch.cuda.memory_stats()` (and associated C++ changes) which returns all instrumented stats as a dictionary. - Added `torch.cuda.snapshot()` (and associated C++ changes) which returns a complete dump of the allocator block/segment state as a list of segments. - Added memory summary generator in `torch.cuda.memory_summary()` for ease of client access to the instrumentation stats. Potentially useful to dump when catching OOMs. Sample output here: https://pastebin.com/uKZjtupq # Implementation: minor changes - Add error-checking helper functions for Python dicts and lists in `torch/csrc/utils/`. - Existing memory management functions in `torch.cuda` moved from `__init__.py` to `memory.py` and star-imported to the main CUDA module. - Add various helper functions to `torch.cuda` to return individual items from `torch.cuda.memory_stats()`. - `torch.cuda.reset_max_memory_cached()` and `torch.cuda.reset_max_memory_allocated()` are deprecated in favor of `reset_peak_stats`. It's a bit difficult to think of a case where only one of those stats should be reset, and IMO this makes the peak stats collectively more consistent. - `torch.cuda.memory_cached()` and `torch.cuda.max_memory_cached()` are deprecated in favor of `*memory_reserved()`. - Style (add access modifiers in the allocator class, random nit fixes, etc.) # Testing - Added consistency check for stats in `test_cuda.py`. This verifies that the data from `memory_stats()` is faithful to the data from `snapshot()`. - Ran on various basic workflows (toy example, CIFAR) # Performance Running the following speed benchmark: https://pastebin.com/UNndQg50 - Before this PR: 45.98 microseconds per tensor creation - After this PR: 46.65 microseconds per tensor creation Pull Request resolved: https://github.com/pytorch/pytorch/pull/27361 Differential Revision: D17758747 Pulled By: jma127 fbshipit-source-id: 5a84e82d696c40c505646b9a1b4e0c3bba38aeb6
66 lines
1.5 KiB
ReStructuredText
66 lines
1.5 KiB
ReStructuredText
torch.cuda
|
|
===================================
|
|
|
|
.. currentmodule:: torch.cuda
|
|
|
|
.. automodule:: torch.cuda
|
|
:members:
|
|
|
|
Random Number Generator
|
|
-------------------------
|
|
.. autofunction:: get_rng_state
|
|
.. autofunction:: get_rng_state_all
|
|
.. autofunction:: set_rng_state
|
|
.. autofunction:: set_rng_state_all
|
|
.. autofunction:: manual_seed
|
|
.. autofunction:: manual_seed_all
|
|
.. autofunction:: seed
|
|
.. autofunction:: seed_all
|
|
.. autofunction:: initial_seed
|
|
|
|
|
|
Communication collectives
|
|
-------------------------
|
|
|
|
.. autofunction:: torch.cuda.comm.broadcast
|
|
|
|
.. autofunction:: torch.cuda.comm.broadcast_coalesced
|
|
|
|
.. autofunction:: torch.cuda.comm.reduce_add
|
|
|
|
.. autofunction:: torch.cuda.comm.scatter
|
|
|
|
.. autofunction:: torch.cuda.comm.gather
|
|
|
|
Streams and events
|
|
------------------
|
|
|
|
.. autoclass:: Stream
|
|
:members:
|
|
|
|
.. autoclass:: Event
|
|
:members:
|
|
|
|
Memory management
|
|
-----------------
|
|
.. autofunction:: empty_cache
|
|
.. autofunction:: memory_stats
|
|
.. autofunction:: memory_summary
|
|
.. autofunction:: memory_snapshot
|
|
.. autofunction:: memory_allocated
|
|
.. autofunction:: max_memory_allocated
|
|
.. autofunction:: reset_max_memory_allocated
|
|
.. autofunction:: memory_reserved
|
|
.. autofunction:: max_memory_reserved
|
|
.. autofunction:: reset_max_memory_reserved
|
|
.. autofunction:: memory_cached
|
|
.. autofunction:: max_memory_cached
|
|
.. autofunction:: reset_max_memory_cached
|
|
|
|
NVIDIA Tools Extension (NVTX)
|
|
-----------------------------
|
|
|
|
.. autofunction:: torch.cuda.nvtx.mark
|
|
.. autofunction:: torch.cuda.nvtx.range_push
|
|
.. autofunction:: torch.cuda.nvtx.range_pop
|