PyTorch内存管理：Caching Allocator（二）一文读懂Block

Sun, 08 Feb 2026 17:19:40 +0800

前言

PyTorch内存管理：Caching Allocator 的第一部分是对PyTorch内存管理的基本情况的概括。这篇文章讲从代码入手，尝试理解PyTorch Caching Allocator 中最核心的数据结构：Block，以及它在实际内存分配中是如何被管理的。

Block

先通过代码看看它的定义：

struct Block {
  c10::DeviceIndex device; // gpu
  cudaStream_t stream; // allocation stream
  stream_set stream_uses; // streams on which the block was used
  size_t size; // block size in bytes
  size_t requested_size; // memory originally requested
  BlockPool* pool{nullptr}; // owning memory pool
  void* ptr{nullptr}; // memory address
  bool allocated{false}; // in-use flag
  bool mapped{true}; // is the virtual address range this Block references
                     // backed by physical pages. Always true when
                     // expandable_segment_ is null. When false
                     // This Block will be aligned to the segment size
                     // of its expandable_segment_.
  Block* prev{nullptr}; // prev block if split from a larger allocation
  Block* next{nullptr}; // next block if split from a larger allocation
  int event_count{0}; // number of outstanding CUDA events
  int64_t gc_count_base{0}; // get_free_blocks_call_count when Block is inserted
  std::shared_ptr<GatheredContext> context_when_allocated;
  // only set for the first block in the segment (when prev == null)
  // this records the frame information when cudaMalloc was called
  // whereas context_when_allocated records the last time we handed this
  // memory out from our cache.
  std::shared_ptr<GatheredContext> context_when_segment_allocated;

  ExpandableSegment* expandable_segment_{nullptr};

device、size、ptr：这几个参数比较好理解，由于是内存管理，就一定需要涉及到分配的设备、分配的内存大小以及在设备上的地址；
requested_size：分配时请求的实际大小；
stream：代表分配时的stream，默认所有操作会在该流进行；
stream_uses：使用多流时记录使用过的流；
prev、next：说明Block本质是一个链表，这也是为什么可以很方便的在Block内部执行切分、重组操作；
event_count：未完成的cuda_event计数；

还剩一些涉及到虚拟内存、可扩展段以及资源管理的描述，我们放到后文去深究。总之，这几个属性就已经决定了如何申请、释放一个block，以及存在异步冲突怎么解决。

PyTorch内存管理：Caching Allocator（一）

Sat, 31 Jan 2026 13:18:00 +0800

概述

当Pytorch显存溢出时，通常会抛出如下的错误：

torch.OutOfMemoryError: CUDA out of memory. Tried to allocate xxxx MB. GPU 0 has a total capacity of xxxx GiB of which xxxx bytes is free. Of the allocated memory xxxx GiB is allocated by PyTorch, and xxxx MB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation. See documentation for Memory Management (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)

阅读本文，你将理解下面这句报错中的显存大小

GPU 0 has a total capacity of xxxx GiB of which xxxx bytes is free

和下面这句报错中的显存大小的关系

CachingAllocator on YuZhouye's Blog

PyTorch内存管理：Caching Allocator（二）一文读懂Block

前言

Block

PyTorch内存管理：Caching Allocator（一）

概述