CUDA内存使用的理解

为了调试CUDA内存使用情况，PyTorch提供了一种生成内存快照的功能。这些快照记录了在任意时刻分配的CUDA内存状态，并可以选择性地记录产生此快照的所有历史分配事件。

生成的快照可以拖放到 pytorch.org/memory_viz 的交互式查看器中，以便探索这些快照。

生成快照

记录快照的常见做法是：开启内存历史记录，运行待观察的代码，然后保存一个包含序列化快照的文件。

# 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")

使用视觉化工具

打开pytorch.org/memory_viz，并将打包的快照文件拖放到该页面中。这是一个在你的计算机上本地运行的 javascript 应用程序，不会上传任何快照数据。

活跃内存 timeline

活动内存时间线显示了快照中特定 GPU 上随时间变化的所有活张量。在图表上平移或缩放以查看较小的分配情况。将鼠标悬停在已分配的块上，可以看到该块被分配时的堆栈跟踪信息及其地址等详细信息。当数据量很大时，可以调整细节滑块来减少渲染的分配数量并提高性能。

分配器状态记录

分配器状态历史显示了单独的分配器事件在左侧的时间线上。选择时间线上的一个事件，以查看该事件发生时分配器状态的可视化摘要。此摘要展示了从 cudaMalloc 返回的每个单独段以及这些段如何被分割成一个个分配块或空闲空间。将鼠标悬停在段和块上，可以看到内存分配时的堆栈跟踪。将鼠标悬停在事件上，可以看到事件发生时的堆栈跟踪，例如张量释放时的情况。内存不足错误报告为 OOM 事件。查看 OOM 发生时的内存状态可能会提供有关为什么即使仍存在预留内存的情况下分配失败的原因。

堆栈跟踪信息还报告了分配发生的地址。地址 b7f064c000000_0 指的是位于地址 7f064c000000 的 (b) 块，这是该地址被分配的第“_0”次。这个唯一的字符串可以在活动内存时间线中查找，并在活动状态历史记录中搜索，以检查张量被分配或释放时的内存状态。

快照API参考

torch.cuda.memory._record_memory_history(enabled='all', context='all', stacks='all', max_entries=9223372036854775807, device=None)[源代码]

启用与内存分配相关的堆栈跟踪记录，以便你可以确定任何一块内存是在哪里被分配的，特别是在调用torch.cuda.memory._snapshot()时。

除了在每次分配和释放操作中保留堆栈跟踪外，这还将记录所有分配/释放事件的历史。

使用torch.cuda.memory._snapshot() 获取此信息，并利用_memory_viz.py 中的工具来可视化快照。

Python 的追踪收集速度非常快（每条追踪只需 2 微秒），因此如果预期可能会遇到内存问题，可以在生产环境中开启这项功能。

C++ 跟踪收集也非常快（约 50ns/帧），对于许多典型程序而言，每个跟踪大约需要 ~2us，但这会因堆栈深度的不同而有所变化。

参数

• enabled (Literal[None, "state", "all"], optional) – None，禁用内存历史记录。“state”，保留当前分配的内存信息。“all”，除了保留当前分配的内存信息外，还保存所有alloc/free调用的历史记录。默认值为“all”。

• context (Literal[None, "state", "alloc", "all"], 可选) – None：不记录任何回溯。“state”：为当前分配的内存记录回溯信息。“alloc”：除了“state”外，还保留所有分配调用的回溯信息。“all”：在“alloc”的基础上，再保留释放调用的回溯信息。默认值为“all”。

• stacks (Literal["python", "all"], optional) – “python”，在跟踪回溯中包含 Python、TorchScript 和 inductor 帧；“all”，在此基础上还包含 C++ 帧，默认值为 “all”。

• max_entries (int, 可选) – 在记录的历史中保留最多max_entries个分配和释放事件。

torch.cuda.memory._snapshot(device=None)[源代码]

保存调用时的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

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

返回值

快照字典对象

torch.cuda.memory._dump_snapshot(filename='dump_snapshot.pickle')[源代码]

将 torch.memory._snapshot() 字典的 pickle 版本保存到文件中。

此文件可以在 pytorch.org/memory_viz 的交互式快照查看器中打开。

参数

filename (str, optional) – 要创建的文件名称，默认为 “dump_snapshot.pickle”。