理解 CUDA 内存使用情况#
创建日期: 2023 年 8 月 23 日 | 最后更新日期: 2025 年 9 月 2 日
为了调试 CUDA 内存使用情况，PyTorch 提供了一种生成内存快照的方法，该方法可以记录任何时间点的已分配 CUDA 内存状态，并可选择记录导致该快照的分配事件历史。
生成的快照可以拖放到托管在 pytorch.org/memory_viz 上的交互式查看器中，用于探索快照。
注意
本文档所述的内存分析器和可视化工具只能看到通过 PyTorch 分配器分配和管理的 CUDA 内存。直接从 CUDA API 分配的任何内存都无法在 PyTorch 内存分析器中看到。
NCCL（用于 CUDA 设备上的分布式通信）是一个常见的库，它会分配一些 PyTorch 内存分析器不可见的 GPU 内存。有关更多信息，请参阅 识别非 PyTorch 分配。
生成快照#
记录快照的常见模式是启用内存历史记录，运行要观察的代码，然后将包含序列化快照的文件保存下来。
# enable memory history, which will
# add tracebacks and event history to snapshots
torch.cuda.memory._record_memory_history()
run_your_code()
torch.cuda.memory._dump_snapshot("my_snapshot.pickle")
使用可视化工具#
打开 https://pytorch.ac.cn/memory_viz 并将序列化快照文件拖放到可视化工具中。可视化工具是一个在您的计算机本地运行的 JavaScript 应用程序。它不会上传任何快照数据。
活动内存时间线#
活动内存时间线显示了快照中特定 GPU 上随时间推移的所有活动张量。平移/缩放图表以查看更小的分配。鼠标悬停在已分配的块上，可以看到该块分配时的时间堆栈跟踪，以及地址等详细信息。可以通过调整详细程度滑块来渲染更少的分配，并在数据量很大时提高性能。
分配器状态历史#
分配器状态历史显示左侧时间线上的单个分配器事件。选择时间线中的一个事件，即可看到该事件下分配器状态的可视化摘要。此摘要显示了从 `cudaMalloc` 返回的每个单独的段，以及它如何被分割成单独的分配块或空闲空间。鼠标悬停在段和块上，可以看到内存分配时的堆栈跟踪。鼠标悬停在事件上，可以看到事件发生时（例如张量被释放时）的堆栈跟踪。内存不足错误会报告为 OOM 事件。查看 OOM 期间的内存状态可能会提供有关为什么尽管仍有保留内存但分配失败的见解。
堆栈跟踪信息还报告了分配发生的地址。地址 `b7f064c000000_0` 指的是地址 `7f064c000000` 处的 (b)lock，这是该地址第 “_0” 次分配。可以在活动内存时间线中查找此唯一字符串，并在活动状态历史中搜索，以检查分配或释放张量时的内存状态。
识别非 PyTorch 分配#
如果您怀疑 CUDA 内存正在 PyTorch 之外分配，您可以使用 `pynvml` 包收集原始 CUDA 分配信息，并将其与 pytorch 报告的分配进行比较。
要收集 PyTorch 之外的原始内存使用情况，请使用 device_memory_used()
import torch
device_idx = ...
print(torch.cuda.device_memory_used(device_idx))
快照 API 参考#
- torch.cuda.memory._record_memory_history(enabled='all', context='all', stacks='all', max_entries=9223372036854775807, device=None, clear_history=False, compile_context=False, global_record_annotations=False)[source]#
启用内存分配相关的堆栈跟踪记录，以便在
torch.cuda.memory._snapshot()中确定任何内存块的分配者。除了为当前每个分配和释放保留堆栈跟踪外，这还将启用所有分配/释放事件历史记录的记录。
使用
torch.cuda.memory._snapshot()检索此信息，并使用 _memory_viz.py 中的工具可视化快照。缓冲区行为#
启用时，此操作最多存储 `max_entries` 个 `TraceEntry` 实例。Python 堆栈跟踪集合的默认值是 `sys.maxsize`，这意味着长期运行或无限期运行的作业应设置合理的限制，以避免过度使用内存。每个条目预计会占用几 KB。
运行时间更长的工作流或 `max_entries` 值较小的工作流将只存储最后累积的 `max_entries` 个条目，这意味着新条目会覆盖旧条目。
C++ 实现，用于参考环形缓冲区实现
if (record_history) { if (alloc_trace->size() < alloc_trace_max_entries_) { alloc_trace->emplace_back(te); } else { (*alloc_trace)[alloc_trace_next++] = te; if (alloc_trace_next == alloc_trace_max_entries_) { alloc_trace_next = 0; } } }
延迟影响#
Python 堆栈跟踪集合速度很快（每个堆栈跟踪 2 微秒），因此如果您预计需要调试内存问题，可以考虑在生产作业中启用此功能。
C++ 堆栈跟踪集合速度也很快（约 50ns/帧），对于许多典型程序来说，这大约是每个堆栈跟踪 2 微秒，但可能会因堆栈深度而异。
- 参数 `enabled`：
“`None`”，禁用内存历史记录。 “`state`”，保留当前已分配内存的信息。 “`all`”，另外保留所有分配/释放调用的历史记录。默认为“`all`”。
- 类型 `enabled`：
Literal[None, “state”, “all”], optional
- 参数 `context`：
“`None`”，不记录任何堆栈跟踪。 “`state`”，记录当前已分配内存的堆栈跟踪。 “`alloc`”，另外保留分配调用的堆栈跟踪。 “`all`”，另外保留释放调用的堆栈跟踪。默认为“`all`”。
- 类型 `context`：
Literal[None, “state”, “alloc”, “all”], optional
- 参数 `stacks`：
“`python`”，包含 Python、TorchScript 和 Inductor 帧的堆栈跟踪。“`all`”，另外包含 C++ 帧。默认为“`all`”。
- 类型 `stacks`：
Literal[“python”, “all”], optional
- 参数 `max_entries`：
在记录的历史记录中保留最多 `max_entries` 个分配/释放事件。
- 类型 `max_entries`：
int, optional
- torch.cuda.memory._snapshot(device=None, augment_with_fx_traces=False)[source]#
保存调用时 CUDA 内存状态的快照。
状态表示为一个具有以下结构的字典。
class Snapshot(TypedDict): segments: List[Segment] device_traces: List[List[TraceEntry]] class Segment(TypedDict): # Segments are memory returned from a cudaMalloc call. # The size of reserved memory is the sum of all Segments. # Segments are cached and reused for future allocations. # If the reuse is smaller than the segment, the segment # is split into more then one Block. # empty_cache() frees Segments that are entirely inactive. address: int total_size: int # cudaMalloc'd size of segment stream: int segment_type: Literal["small", "large"] # 'large' (>1MB) allocated_size: int # size of memory in use active_size: int # size of memory in use or in active_awaiting_free state blocks: List[Block] class Block(TypedDict): # A piece of memory returned from the allocator, or # current cached but inactive. size: int requested_size: int # size requested during malloc, may be smaller than # size due to rounding address: int state: Literal[ "active_allocated", # used by a tensor "active_awaiting_free", # waiting for another stream to finish using # this, then it will become free "inactive", ] # free for reuse frames: List[Frame] # stack trace from where the allocation occurred class Frame(TypedDict): filename: str line: int name: str # Optional FX debug fields (present when augment_with_fx_traces=True # and the frame corresponds to FX-generated code) fx_node_op: str # FX node operation type (e.g., 'call_function', 'output') fx_node_name: str # FX node name (e.g., 'linear', 'relu_1') fx_original_trace: str # Original model source code stack trace class TraceEntry(TypedDict): # When `torch.cuda.memory._record_memory_history()` is enabled, # the snapshot will contain TraceEntry objects that record each # action the allocator took. action: Literal[ "alloc" # memory allocated "free_requested", # the allocated received a call to free memory "free_completed", # the memory that was requested to be freed is now # able to be used in future allocation calls "segment_alloc", # the caching allocator ask cudaMalloc for more memory # and added it as a segment in its cache "segment_free", # the caching allocator called cudaFree to return memory # to cuda possibly trying free up memory to # allocate more segments or because empty_caches was called "oom", # the allocator threw an OOM exception. 'size' is # the requested number of bytes that did not succeed "snapshot", # the allocator generated a memory snapshot # useful to coorelate a previously taken # snapshot with this trace ] addr: int # not present for OOM frames: List[Frame] size: int stream: int device_free: int # only present for OOM, the amount of # memory cuda still reports to be free
- 参数:
device (Device) – 要捕获快照的设备。如果为 None，则捕获当前设备。
augment_with_fx_traces – 如果为 True，则使用 FX 调试信息增强堆栈跟踪帧，将生成的 FX 代码映射回原始模型源代码。这会将 `fx_node_op`、`fx_node_name`、`fx_original_trace` 和 `fx_node_info` 字段添加到 Frame 对象。默认值为 False。
- 返回:
快照字典对象
