为什么建行网站打不开,广州部队网站建设费用,网络域名查询,vue网页模板免费因引入 expandable_segments 机制#xff0c;PyTorch 2.1.0 版本发生了较大变化。本文关注的是 PyTorch 原生的 GPU 内存管理机制#xff0c;故研究的 PyTorch 版本为 2.0.0。代码地址#xff1a; c10/cuda/CUDACachingAllocator.hc10/cuda/CUDACachingAllocator.cpp 更多内… 因引入 expandable_segments 机制PyTorch 2.1.0 版本发生了较大变化。本文关注的是 PyTorch 原生的 GPU 内存管理机制故研究的 PyTorch 版本为 2.0.0。代码地址 c10/cuda/CUDACachingAllocator.hc10/cuda/CUDACachingAllocator.cpp 更多内容请参考 Ubuntu 22.04 LTS 源码编译安装 PyTorch【翻译】pytorch/CONTRIBUTING.md【翻译】Pytorch机制源代码分析与内存管理调研深度学习框架与静态/动态计算图【笔记】PyTorch 源码学习阅读经验 代码结构PyTorch 源码学习从 Tensor 到 StoragePyTorch 源码学习Dispatch Autograd OperatorsPyTorch 源码学习GPU 内存管理之它山之石——TensorFlow BFC 算法 文章目录 1 PyTorch 为什么要有 CUDACachingAllocator2 核心类2.1 class NativeCachingAllocator2.2 Allocator 和 CUDAAllocator 3 全局变量4 核心数据结构4.1 struct Block 内存块4.1.1 struct SegmentInfo 内存段的信息4.1.2 struct BlockInfo 分配块的信息4.1.3 进一步区分 Segment 和 Block 4.2 struct BlockPool 内存池4.3 struct PrivatePool 私有内存池4.4 struct AllocParams 分配内存时的参数和状态4.5 class DeviceCachingAllocator 类成员变量 5 PyTorch BFC 算法6 核心函数6.1 Block* malloc 分配函数6.2 void free 释放函数 7 功能函数7.1 static bool BlockComparator 自定义比较函数7.2 static size_t round_size 对内存大小 size 进行调整7.3 BlockPool get_pool 返回对应的内存块池7.4 static size_t get_allocation_size 确定实际要进行分配的内存大小7.5 bool get_free_block 寻找满足分配需求的空闲内存块7.6 void garbage_collect_cached_blocks 释放未使用的缓存块7.7 bool release_available_cached_blocks 优先从最大的块开始释放7.8 bool release_cached_blocks 释放所有未使用的内存块7.9 void release_blocks 释放所有非拆分块7.10 void release_block 释放指定的内存块7.11 bool alloc_block 为指定的分配请求分配一个内存块7.12 cudaError_t cudaMallocMaybeCapturing 一个封装的内存分配工具7.13 bool is_split 判断 block 是否被拆分过7.14 bool should_split 判断给定的内存块 block 是否应该进行分割7.15 void free_block 将一个块移动到缓存的空闲块池中7.16 size_t try_merge_blocks 合并先前拆分的块 8 问题与挑战参考资料 1 PyTorch 为什么要有 CUDACachingAllocator
CUDACachingAllocator是 PyTorch 的 CUDA 缓存分配器。其目的是通过缓存和复用内存块来减少频繁的cudaMalloc和cudaFree操作提升 CUDA 程序的性能。
GPU 底层机制分析显存分配开销 这篇文章通过实验分析了 CUDA 原生显存分配接口的显存分配开销和显存分配碎片。结论就是虽然 CUDA 原生显存分配接口逻辑比较简单但显存分配开销无法接受显存分配碎片也需要优化。 以下内容来自Pytorch 显存管理机制与显存占用分析方法这篇文章对显存管理机制的总结写得很好。 GPU 作为一种通用的数据处理设备为了满足更广泛客户的需求且保证更小的维护成本其 API 在设计的时候比较开放尽管 CUDA 生态中也有高阶 API但并没有针对某个深度学习框架做设计优化其中显存的精细管理留给上层的深度学习框架去完成。
cudaMallocCUDA API是从 GPU 申请显存最常用的方式给定指针和数据大小即可进行 API 调用其调用有着不小的时间开销且是 stream 内的同步操作。当深度学习框架使用的数据非常零碎且数量多时需要反复调用cudaMalloc该行为会直接影响程序的整体性能因此深度学习框架的显存管理机制在设计时要尽量降低cudaMalloc的调用频次。
PyTorch 框架基于 CUDA API 实现了一套显存管理逻辑/机制可更好地满足框架的日常使用需求相比原生的 CUDA API 可做到管理细化、使用相对高效其采用动态申请与二次分配的设计思路
动态申请在使用的时候根据用量实时地向 GPU 发出请求最大优点是不会占用过量的显存方便多人同时使用一个设备与之相对的是 TensorFlow 早期版本在启动前就把 GPU 上的大部分显存都申请到然后再去分配使用。二次分配将显存的申请与使用进行分离即显存申请后会进行二次分配。显存管理机制会先通过cudaMalloc向 GPU 申请一个显存块 Segment然后从 Segment 分离出子块 Block我们使用的是分离后的 Block 显存而不直接使用 Segment。 2 核心类
Allocator - CUDAAllocator - NativeCachingAllocator - DeviceCachingAllocator
struct C10_API Allocatorclass CUDAAllocator : public Allocatorclass NativeCachingAllocator : public CUDAAllocator {public:std::vectorstd::unique_ptrDeviceCachingAllocator device_allocator;class DeviceCachingAllocator以下内容来自PyTorch显存管理介绍与源码解析二 显存管理代码中的类的大体关系如下图所示。cudaCachingAllocator类的结构相对来说比较清晰代码不是特别清晰DeviceCachingAllocator管理一个设备显存THCCachingAllocator管理一个进程上所有的DeviceCachingAllocatorCudaCachingAllocator继承自Allocator方便框架调用。 DeviceCachingAllocator类是显存管理机制真正实现的地方它负责完成 device 上的显存管理包括显存开辟、分块、合并、释放等。 说明这张图分析的应该不是作者在文中提到的 PyTorch 1.13.0因为在该版本的代码注释中提到THCCachingAllocator这种说法已经是 old-style 了。作者这里提到的THCCachingAllocator在最新的版本里对应的是NativeCachingAllocator而CudaCachingAllocator对应的是CUDAAllocator。但上面的内容仍值得参考。 以下内容来自Pytorch 1.6 显存管理CudaCachingAllocator剖析 Pytorch 通过 Allocator 实现显存管理实现如下图 DeviceCachingAllocator每个 GPU 设备卡都维护一个这样的结构用于对该设备进行显存管理THCCachingAllocator维护一个 DeviceCachingAllocator 列表及一些全局的状态核心逻辑委托给 DeviceCachingAllocator 实现CudaCachingAllocator这是 Pytorch 默认实现的分配器包装器向上层暴露一些有用的接口核心逻辑还是委托给 THCCachingAllocator 来实现用户也可以自定义一个实现分配器替代 CudaCachingAllocator。 在 PyTorch 1.6.0 里确实是上面这样的类关系。 2.1 class NativeCachingAllocator 以下内容来自PyTorch显存管理介绍与源码解析二 NativeCachingAllocator该类主要是管理 CPU 主进程调用的 GPU 的显存使用状况。通过创建一个数组每个 GPU 有一个与之对应的DeviceCachingAllocator实例管理同时记录下进程中创建的全部的 blocks。 NativeCachingAllocator的代码实现
class NativeCachingAllocator : public CUDAAllocator {private:// 存储已分配的显存块// allocated blocks by device pointerska::flat_hash_mapvoid*, Block* allocated_blocks;// 将分配的显存块添加到 allocated_blocks 中void add_allocated_block(Block* block) {allocated_blocks[block-ptr] block;}public:// 每个设备GPU都有一个 DeviceCachingAllocator 实例用于管理该设备上的显存分配std::vectorstd::unique_ptrDeviceCachingAllocator device_allocator;// 根据设备指针获取对应的显存块Block* get_allocated_block(void* ptr, bool remove false) {auto it allocated_blocks.find(ptr); // 查找 ptr 对应的 Block*if (it allocated_blocks.end()) {return nullptr;}Block* block it-second;if (remove) {allocated_blocks.erase(it);}return block;}// 初始化 device_allocator为每个设备创建一个 DeviceCachingAllocator 实例void init(int device_count) override {const auto size static_castint64_t(device_allocator.size());if (size device_count) {device_allocator.resize(device_count);for (const auto i : c10::irange(size, device_count)) {device_allocator[i] std::make_uniqueDeviceCachingAllocator();}}}// 在指定设备上分配显存并将分配的显存块添加到 allocated_blocks 中/** allocates a block which is safe to use from the provided stream */void malloc(void** devPtr, int device, size_t size, cudaStream_t stream) {// 1) 检查设备号是否有效TORCH_INTERNAL_ASSERT(0 device static_castsize_t(device) device_allocator.size(), Allocator not initialized for device ,device,: did you call init?);// 2) 调用 6.1 Block* malloc 分配显存块Block* block device_allocator[device]-malloc(device, size, stream);// 3) 将分配的显存块添加到 allocated_blocks 中add_allocated_block(block);}// 释放指定的显存块void free(void* ptr) {if (!ptr) {return;}// 1) 从 allocated_blocks 中获取对应的 Block*并将其从 allocated_blocks 中移除Block* block get_allocated_block(ptr, true /* remove */);if (!block) {TORCH_CHECK(false, invalid device pointer: , ptr);}// 2) 调用 6.2 void free 释放显存块device_allocator[block-device]-free(block);}// 分配显存并返回一个 DataPtr 对象DataPtr allocate(size_t size) const override {constexpr size_t one_exa_bytes 1152921504606846976ULL;TORCH_CHECK_WITH(OutOfMemoryError,size one_exa_bytes,CUDA out of memory. Tried to allocate more than 1EB memory.);int device;C10_CUDA_CHECK(cudaGetDevice(device));void* r nullptr;// 关于 forceUncachedAllocator() 有下面这样的注释// 其实就是直接不使用 CUDACachingAllocator直接调用 cudaMalloc 分配内存// Returns whether to force all allocations to bypass the caching allocator and// go straight to cudaMalloc. This setting is useful when debugging GPU memory// errors, since the caching allocator foils cuda-memcheck.if (forceUncachedAllocator()) {// Deliberately dont use cudaMallocMaybeCapturing here, to force an error// if someone tries to use forceUncachedAllocator while capturing.C10_CUDA_CHECK(cudaMalloc(r, size));return {r, r, uncached_delete, Device(DeviceType::CUDA, device)};}// 使用 CUDACachingAllocator 分配内存// 调用了 void malloc(void** devPtr, int device, size_t size, cudaStream_t stream)if (size ! 0) {// Allocator declars allocate const!?const_castNativeCachingAllocator*(this)-malloc(r, device, size, cuda::getCurrentCUDAStream(device));}return {r, r, local_raw_delete, Device(DeviceType::CUDA, device)};}// 返回删除器函数指针DeleterFnPtr raw_deleter() const override {if (forceUncachedAllocator()) {return uncached_delete;} else {return local_raw_delete;}}// 分配指定大小的显存void* raw_alloc(size_t nbytes) override {if (nbytes 0) {return nullptr;}int device;C10_CUDA_CHECK(cudaGetDevice(device));void* r nullptr;// 调用了 void malloc(void** devPtr, int device, size_t size, cudaStream_t stream)malloc(r, device, nbytes, cuda::getCurrentCUDAStream(device));return r;}// 在指定的 CUDA 流上分配显存void* raw_alloc_with_stream(size_t nbytes, cudaStream_t stream) override {if (nbytes 0) {return nullptr;}int device;C10_CUDA_CHECK(cudaGetDevice(device));void* r nullptr;// 调用了 void malloc(void** devPtr, int device, size_t size, cudaStream_t stream)malloc(r, device, nbytes, stream);return r;}// 释放指定的显存void raw_delete(void* ptr) override {this-free(ptr);}};这里要区分下DataPtr allocate(size_t size)和void* raw_alloc(size_t nbytes)
allocate调用链allocate - malloc - add_allocated_block。在分配显存后会将显存块添加到allocated_blocks中并返回一个 DataPtr 对象。raw_alloc调用链raw_alloc - malloc。仅分配显存并返回指针。
TODO: allocate和raw_alloc两者的使用场景还是有很大区别的因本人精力有限暂不深入研究了。
一些线索
1struct C10_API Allocator的代码实现
struct C10_API Allocator {virtual ~Allocator() default;virtual DataPtr allocate(size_t n) const 0;// If this returns a non nullptr, it means that allocate()// is guaranteed to return a unique_ptr with this deleter attached;// it means the rawAllocate and rawDeallocate APIs are safe to use.// This function MUST always return the same BoundDeleter.virtual DeleterFnPtr raw_deleter() const {return nullptr;}void* raw_allocate(size_t n) {auto dptr allocate(n);AT_ASSERT(dptr.get() dptr.get_context());return dptr.release_context();}void raw_deallocate(void* ptr) {auto d raw_deleter();AT_ASSERT(d);d(ptr);}
};2class CUDAAllocator : public Allocator的代码实现
class CUDAAllocator : public Allocator {public:virtual void* raw_alloc(size_t nbytes) 0;virtual void* raw_alloc_with_stream(size_t nbytes, cudaStream_t stream) 0;virtual void raw_delete(void* ptr) 0;// ……
};2.2 Allocator 和 CUDAAllocator
不是本文的重点就不深入分析了。前文的“一些线索”也展示了 Allocator 和 CUDAAllocator 两个类的部分代码感兴趣可以直接学习源码。
本文剩余内容都实现于 DeviceCachingAllocator 类中这也是与设备打交道的最底层的一个类。
3 全局变量
全局变量的代码实现
constexpr size_t kMinBlockSize 512;
// all sizes are rounded to at least 512 bytes
// 为了优化内存管理显存块的大小通常会对齐到特定的值。
// 表示显存分配的最小粒度为 512 字节。constexpr size_t kSmallSize 1048576;
// largest small allocation is 1 MiB
// 定义了 小分配 的最大阈值为 1 MiB1024 * 1024 字节。
// 小于或等于 kSmallSize 的分配被视为 小分配。constexpr size_t kMinLargeAlloc 10485760;
// allocations between 1 and 10 MiB may use kLargeBuffer
// 定义了 大分配 的最小阈值为 10 MiB10 * 1024 * 1024 字节。
// 大于或等于 kMinLargeAlloc 的分配被视为 大分配。constexpr size_t kSmallBuffer 2097152;
// small allocations are packed in 2 MiB blocks
// 定义了 小分配 的默认打包大小为 2 MiB2 * 1024 * 1024 字节。constexpr size_t kLargeBuffer 20971520;
// large allocations may be packed in 20 MiB blocks
// 定义了 大分配 的默认打包大小为 20 MiB20 * 1024 * 1024 字节。constexpr size_t kRoundLarge 2097152;
// round up large allocations to 2 MiB
// 定义了 大分配 的对齐单位为 2 MiB2 * 1024 * 1024 字节。
// 所有大分配的大小都会被向上取整到 2 MiB 的倍数以简化分配管理并减少碎片。constexpr size_t kRoundUpPowerOfTwoIntervals 16;4 核心数据结构
4.1 struct Block 内存块
Block表示内存池中的一个内存块封装了内存分配的信息。 以下内容来自一文读懂 PyTorch 显存管理机制 Block
分配 / 管理内存块的基本单位(stream_id, size, ptr) 三元组可以特异性定位一个 Block即 Block 维护一个 ptr 指向大小为 size 的内存块隶属于 stream_id 的 CUDA Stream。所有地址连续的 Block不论是否为空闲只要是由Allocator::malloc得来的都被组织在一个双向链表里便于在释放某一个 Block 时快速检查前后是否存在相邻碎片若存在可以直接将这三个 Block 合成为一个。Block 在 Allocator 内有两种组织方式 一种是显式地组织在 BlockPool红黑树中按照大小排列另一种是具有连续地址的 Block 隐式地组织在一个双向链表里通过结构体内的 prev, next 指针可以以 O(1) 时间查找前后 Block 是否空闲便于在释放当前 Block 时合并碎片。
一段连续空间内部由双向链表组织的 Block 们如下图所示 Block的代码实现
struct Block {int device; // gpu 表示分配内存的 GPU 设备 ID。cudaStream_t stream; // allocation stream // 分配内存时关联的 CUDA 流。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 标志是否正在使用。Block* prev{nullptr}; // prev block if split from a larger allocationBlock* next{nullptr}; // next block if split from a larger allocation// 若该块是从更大的块分割而来指向前后关联块。int event_count{0}; // number of outstanding CUDA eventsint gc_count{0}; // counter for prioritizing older / less useful blocks for garbage collection// 用于垃圾回收计数优先回收不活跃的块。// 构造函数1用于创建实际的 Block 对象。Block(int device,cudaStream_t stream,size_t size,BlockPool* pool,void* ptr): device(device),stream(stream),stream_uses(),size(size),requested_size(0),pool(pool),ptr(ptr) {}// constructor for search key// 构造函数2用于搜索时的键。Block(int device, cudaStream_t stream, size_t size): device(device),stream(stream),stream_uses(),size(size),requested_size(0) {}// 检查该块是否是通过分割其他内存块生成的。bool is_split() const {return (prev ! nullptr) || (next ! nullptr);}
};注意以下 4.1.1 和 4.1.2 中的 Segment 和 Block 其实都是逻辑上的概念本质上其实都是 4.1 中的 Block。可以理解为Segment 是cudaMalloc直接申请得到的一段内存而 Block 是实际分配给 Tensor 的内存块一个 Segment 可能包含多个 Block。 SegmentInfo、BlockInfo 记录 segment 和 block 的信息用于 Snapshot 数据记录可用于显存可视化。 4.1.1 struct SegmentInfo 内存段的信息
Segment 是通过一次cudaMalloc调用分配的一块连续的显存区域。一个内存段 (Segment) 可以包含多个分配块 (Block) 。
SegmentInfo的代码实现
// Struct containing info of a memory segment (i.e. one contiguous cudaMalloc).
struct SegmentInfo {int64_t device 0; // 分配此段的 CUDA 设备。int64_t address 0; // 内存段的起始地址。int64_t total_size 0; // 整个内存段的总大小。int64_t requested_size 0; // 用户请求的总大小。int64_t allocated_size 0; // 段中实际被分配的内存大小。int64_t active_size 0; // 段中当前正在使用的内存大小。cudaStream_t stream 0; // 与此段相关联的 CUDA 流。bool is_large false; // 标记此段是否是一个大段分配large allocation。std::vectorBlockInfo blocks; // 此段内的分配块列表。
};4.1.2 struct BlockInfo 分配块的信息
Block 是 CUDA 显存分配的基本单元。一个分配块 (Block) 可以是通过cudaMalloc获取的内存段 (Segment) 的一部分。
BlockInfo的代码实现
// Struct containing info of an allocation block (i.e. a fractional part of a cudaMalloc)..
struct BlockInfo {int64_t size 0; // 当前块的大小以字节为单位。int64_t requested_size 0; // 用户实际请求的大小。分配器通常会将请求大小对齐到某个块大小。int32_t gc_counter 0; // 垃圾回收计数器记录该块的回收次数。bool allocated false; // 标记该块是否已被分配给用户。bool active false; // 标记该块是否正在被某个 CUDA 流使用例如分配或计算。std::vectorHistory history;
};4.1.3 进一步区分 Segment 和 Block 以下内容来自PyTorch显存机制分析 直观点来说PyTorch 的显存管理是一个层级结构。 然后他们又是包含与被包含的关系即 PyTorch Allocated memory 使用的是 PyTorch Cached Memory 里的显存PyTorch Cached Memory 则用的是 GPU 的显存。 这里 PyTorch Allocated memory 指的就是实际分配给上层应用的 Block 占用的内存而 PyTorch Cached Memory 指的就是直接通过cudaMalloc申请得到并缓存在内存池的 Block也就是逻辑上的 Segment占用的内存。 以下内容来自PyTorch显存管理介绍与源码解析一和 PyTorch显存管理介绍与源码解析二 在管理机制中将显存的申请与使用过程进行了分离即显存申请后会进行二次分配其过程是先通过cudaMalloc申请一个显存块 segment然后从 segment 分离出来的子显存块 block框架的上层应用使用的是分离后的 block 显存上层应用不直接使用 segment。
进一步通过池化的方式将 block 按照块放入不同的显存池中进行分类管理。 PyTorch 内存管理机制里对显存分了两级segment、blocks显存池分了两类active pool、remaining pool。在结构体的实现上面主要是靠块 (block) 结构体来完成承上启下的功能显存池结构体blockPool主要是记录 block 的归属segment 是一个从cudaMalloc创建出来的显存块承载依然是 block 结构体segment 再次切分就是能够提供给用户使用的显存。 以下内容来自Pytorch 1.6 显存管理CudaCachingAllocator剖析 下图为多种进程使用显存的布局 allocated 是 Pytorch 分配给 Tensor 使用的块即torch.cuda.memory_allocated()的值cached 是 Pytorch 空闲块reserved 是 Pytorch 进程管理的所有显存块即torch.cuda.memory_reserved()的值reserved allocated cached 逻辑上来看这里的 reserved 就是所有 Segment 占用的内存allocated 就是所有分配出去的 Block 占用的内存cached 就是所有未分配出去的 Block 占用的内存。 4.2 struct BlockPool 内存池
BlockPool用于管理多个Block可以分为小块池和大块池。 以下内容来自一文读懂 PyTorch 显存管理机制 BlockPool
内存池用std::set存储 Block 的指针按照 (cuda_stream_id - block size - addr) 的优先级从小到大排序。所有保存在 BlockPool 中的 Block 都是空闲的。DeviceCachingAllocator中维护两种 BlockPool (large_blocks, small_blocks)分别存放较小的块和较大的块为了分别加速小需求和大需求简单地将 1MB 的 Block 归类为小块 1MB 的为大块。
直观理解 Block、BlockPool 见下图 以下内容来自PyTorch显存管理介绍与源码解析一和 PyTorch显存管理介绍与源码解析二 DeviceCachingAllocator类创建两个 BlockPoolLarge_blocks Small_blocks将未使用的 blocks 存入其中。对于在使用的 block通过哈希集合active_blocks存储指向 block 的指针。
目前未使用的块管理有两种类型的池子large_blockssmall_blocks。分类阈值默认值设置为 1M小于 1M 放入small_blocks大于 1M 放入large_blocks。 关于active_blocks见 4.5 class DeviceCachingAllocator 类成员变量 一节的内容。另外在 PyTorch显存管理介绍与源码解析二 文中作者提到 “显存池分了两类active pool、remaining pool”这里的 active pool 指的就是active_blocks。 active_blocks 是一个 ska::flat_hash_setBlock*基于哈希表的集合。 哈希表的特性是插入、删除和查找的平均时间复杂度为 O ( 1 ) O(1) O(1)适合频繁的动态操作。active_blocks目的是快速跟踪所有活跃的内存块从代码来看active_blocks主要用于 Snapshot。 large_blocks和small_blocks是BlockPool类型的对象内部使用std::setBlock*, Comparison来存储和排序 Block是一个基于红黑树的集合。 红黑树的特性是元素有序插入、删除和查找的时间复杂度为 O ( l o g n ) O(log n) O(logn)适合需要排序的场景。large_blocks和small_blocks目的是高效管理空闲内存块。 以下内容来自Pytorch 1.6 显存管理CudaCachingAllocator剖析 空闲块管理规则如下 ❗关于 Small Pool 和 Large Pool笔者有一些和上述博客不同的看法在这里根据笔者的理解做一些补充如果有不对之处还请指出。
Small Pool 和 Large Pool 的作用是管理空闲 Block他们本质上是红黑树集合。
PyTorch 根据请求的大小去这两个 Pool 里寻找空闲内存 Block对于 ≤ 1 MB 的请求首先去 Small Pool 寻找空闲 Block对于1 MB 的请求首先去 Large Pool 寻找空闲 Block。对应代码 L817-L820 和 L824-L829。
当在 Small Pool 和 Large Pool 找不到空闲 Block 时PyTorch 会去申请新的 Block。对应代码 L839-L847。新申请的 Block 确实也有一个 Pool 的属性这里还是对应代码 L817-L820其中
≤ 1MB 的请求对应 Small Pool会去申请一个 2MB 的 Cache Block1MB 的请求对应 Large Pool会按规则申请更大的 Cache Block。
但从alloc_block这个函数的具体实现来看其实并没有将新创建的 Block 插入到 Small Pool 或 Large Pool。对应代码 L626-L646。
而把空闲 Block 插入 Small Pool 或 Large Pool 的只有下面两处
split对应代码 L959这里会把拆分后的 remaining Block 插入到与被拆分的 Block 一样的 Pool 中。并且根据should_split函数插入 Small Pool 的一定都是 512B ≤ size2MB 的空闲块插入 Large Pool 的一定都是 size1MB 的空闲块。对应代码 L1546-L1554。free对应代码 L1433这里会把释放后的 Block 插入到对应的 Pool 中。因为可能涉及合并操作也可能不合并所以在这个过程中插入 Small Pool 的一定都是 512B ≤ size ≤ 2MB 的空闲块插入 Large Pool 的一定都是 size1MB 的空闲块。
这么分析来看Small Pool 和 Large Pool 管理的空闲 Block 的大小其实是有一些交叉的。主要是 1MB size ≤ 2MB 这个区间。 以下内容来自Pytorch 显存管理机制与显存占用分析方法 这篇文章的说法比较妥当
显存管理机制会依据未分配 Block 所在 Segment 的大小将未分配的 Block 划入 large poolSegment 2MB或 small poolSegment ≤ 2MB。用户创建 tensor 申请显存时会先从 tensor size 对应未分配显存的 pool 中查找是否有满足 size 要求的 Block如果没有才会向 GPU 申请新的 Segment 显存块。 比如显存管理器当前有且仅有一个 2MB 的 Segment已分配了 0.5MB还剩 1.5MB用户此时需要创建一个 1.1MB 的 tensor那么显存管理器不会从这 1.5MB 的未分配 Block 中划分一部分空间给 tensor而是额外申请一个 20MB 的 Segment 再进行分配。 BlockPool的代码实现
struct BlockPool {std::setBlock*, Comparison blocks;// 一个 std::set 容器用于存储和排序 Block 指针。const bool is_small;// 标识该池是否用于小块管理。PrivatePool* owner_PrivatePool;// 指向关联的 PrivatePool可选。// 构造函数BlockPool(Comparison comparator,// 初始化时需要自定义比较函数bool small,PrivatePool* private_pool nullptr): blocks(comparator), is_small(small), owner_PrivatePool(private_pool) {}
};关于 std::set 容器【C】set 容器最全解析什么是 set? set容器的常用接口有那些 4.3 struct PrivatePool 私有内存池
PrivatePool是一个支持 CUDA 图CUDA Graphs的私有内存池用于管理显存块。
PrivatePool的代码实现
// CUDA graphs helper
struct PrivatePool {int use_count;// Number of live graphs using this pool// 记录使用该池的 CUDA 图数量。int cudaMalloc_count;// Number of unfreed cudaMallocs made for this pool.// 记录未释放的 cudaMalloc 调用计数。// When use_count and cudaMalloc_count drop to zero, we can delete this PrivatePool from graph_pools.// 当引用计数use_count和 cudaMalloc 计数都降为零时我们可以从图内存池graph_pools中删除这个私有内存池PrivatePool。BlockPool large_blocks;BlockPool small_blocks;// BlockPool 类型的成员分别管理大块和小块。// 构造函数PrivatePool(): use_count(1),cudaMalloc_count(0),// 初始化时传入比较函数。large_blocks(BlockComparator, /*is_small*/false, this),small_blocks(BlockComparator, /*is_small*/true, this) {}PrivatePool(const PrivatePool) delete;PrivatePool(PrivatePool) delete;PrivatePool operator(const PrivatePool) delete;
};4.4 struct AllocParams 分配内存时的参数和状态
AllocParams封装分配内存时的参数和状态用于寻找或创建合适的内存块。
AllocParams的代码实现
struct AllocParams {Block search_key;// 存储查找所需的信息BlockPool* pool;// 目标内存池。size_t alloc_size;// 分配大小可能大于 size因为需要对齐。Block* block;// 指向找到或创建的内存块。// 构造函数AllocParams(int device,size_t size,cudaStream_t stream,BlockPool* pool,size_t alloc_size,DeviceStats stats): search_key(device, stream, size),pool(pool),alloc_size(alloc_size),block(nullptr),err(cudaSuccess) {}// 提供访问设备、流和大小的接口。int device() const {return search_key.device;}cudaStream_t stream() const {return search_key.stream;}size_t size() const {return search_key.size;}
};4.5 class DeviceCachingAllocator 类成员变量
DeviceCachingAllocator的代码实现
class DeviceCachingAllocator {private:BlockPool large_blocks;BlockPool small_blocks;// unallocated cached blocks larger than 1 MB// unallocated cached blocks 1 MB or smaller// 管理未分配的大内存块 1 MB、小内存块≤ 1 MB。ska::flat_hash_setBlock* active_blocks;// allocated or in use by a stream. Holds all active allocations,// whether they came from graph_pools or one of the BlockPools above.// 用于跟踪当前被分配活跃的内存块。// 无论内存块来自 large_blocks、small_blocks还是 CUDA 图专用内存池它都会被记录在这里。size_t total_allocated_memory 0;// record used memory.// 记录当前分配的总显存大小。size_t allowed_memory_maximum 0;// 最大允许的显存分配量。bool set_fraction false;// 标记是否已经设置了允许的最大内存。public:// 默认构造函数DeviceCachingAllocator(): large_blocks(BlockComparator, /*is_small*/false),// 初始化大块内存池small_blocks(BlockComparator, /*is_small*/true),// 初始化小块内存池alloc_trace(new std::vectorTraceEntry()) {stats.max_split_size CachingAllocatorConfig::max_split_size();context_recorder_.store(nullptr);}}结构体和类的区别Struct 和 Class 的区别以及使用场景_class和struct区别 5 PyTorch BFC 算法 以下内容来自PyTorch显存管理介绍与源码解析一 整体的运行逻辑 查找从显存池里面查找看是否有满足当前条件的 block如果有直接跳入步骤3如果没有跳入步骤2创建调用cudaMalloc创建新的 segment如果满足 split 条件跳入步骤3不满足直接跳入步骤5切分搜索/创建到满足条件的 block 根据需求进行切分如果有剩余跳入步骤4如果没有剩余跳入步骤5保存将切分后的 block 放入显存池中执行完成后再执行步骤5返回将创建好的 block 数据指针返回给上层API回收当 block 释放时会先去池子里面查找是否有可以与之合并的 block如果有先进行合并然后将合并后的 block 存入显存池中。释放当 block 为一个 segment 时可以触发cudaFree操作释放显存。 6 核心函数
6.1 Block* malloc 分配函数 以下内容来自PyTorch CUDA backend作者这种作图方式值得学习。 以下内容来自Pytorch 显存管理机制与显存占用分析方法 malloc的代码实现 Block* malloc(int device, size_t orig_size, cudaStream_t stream) {// 对齐请求的内存大小。size_t size round_size(orig_size);// 获取内存池。auto pool get_pool(size, stream);// 实际的分配大小。const size_t alloc_size get_allocation_size(size);// 创建一个 AllocParams 对象封装所有分配参数。AllocParams params(device, size, stream, pool, alloc_size, stats);// First, try to get a block from the existing pool.// 1、从现有池中尝试分配bool block_found // Search pool// 尝试直接从内存池中获取空闲块。get_free_block(params)// Trigger callbacks and retry search// 如果上一步失败触发释放内存的回调然后再次尝试获取空闲块。|| (trigger_free_memory_callbacks(params) get_free_block(params));// Cant reuse an existing block; try to get a new one.// 2、如果复用失败尝试分配新的内存块if (!block_found) {// Do garbage collection if the flag is set.// 如果启用了垃圾回收if (C10_UNLIKELY(set_fraction CachingAllocatorConfig::garbage_collection_threshold() 0.0)) {// 清理不再使用的缓存块。garbage_collect_cached_blocks();}// Attempt allocate// 尝试分配新内存块block_found alloc_block(params, false)// Free enough available cached blocks to satisfy alloc and retry alloc.// 释放足够的可用缓存块以满足 alloc 并重试 alloc。|| (release_available_cached_blocks(params) alloc_block(params, false))// Free all non-split cached blocks and retry alloc.// 释放所有未拆分的缓存块并重试 alloc。|| (C10_LIKELY(captures_underway 0) release_cached_blocks() alloc_block(params, true));}// 如果多次尝试仍然失败报告内存不足OOM错误if (!block_found) {// ……}// 检查 params 参数的有效性确保// CUDA 操作成功params.err cudaSuccess。// 内存块params.block和其指针params.block-ptr不为空。TORCH_INTERNAL_ASSERT(params.err cudaSuccess params.block ! nullptr params.block-ptr ! nullptr);// 当前要操作的内存块。Block* block params.block;// 用于存储分割后的剩余内存块。Block* remaining nullptr;// 检查 block 是否已经是分割块。const bool already_split block-is_split();// 判断是否需要分割if (should_split(block, size)) {// 内存块分割处理// 1.保存 remaining 为原始块。remaining block;// 2.创建新的 block它使用原始块的前 size 内存地址。block new Block(device, stream, size, pool, block-ptr);// 3.block 链接到前面的块。block-prev remaining-prev;if (block-prev) {block-prev-next block;}// 4.block 链接到 remainingblock-next remaining;remaining-prev block;// 5.调整 remaining 起始地址和大小。remaining-ptr static_castchar*(remaining-ptr) size;remaining-size - size;// 6.将剩余块重新插入到内存池通过断言确保插入成功。bool inserted pool.blocks.insert(remaining).second;TORCH_INTERNAL_ASSERT_DEBUG_ONLY(inserted);// 7.更新分割块的统计信息// ……} else if (already_split) {// An already-split block is becoming active// 如果块不需要分割但已经是分割块则直接更新统计信息表示该块变为“活动”块。// ……}// 标记内存块为“已分配”。block-allocated true;// 设置用户请求的大小。block-requested_size orig_size;// 将分配完成的块插入到 active_blocks 集合中。bool inserted active_blocks.insert(block).second;TORCH_INTERNAL_ASSERT_DEBUG_ONLY(inserted);// 更新统计数据// ……return block;}malloc依赖的函数
7.2 static size_t round_size(size_t size)7.3 BlockPool get_pool(size_t size, cudaStream_t stream)7.4 static size_t get_allocation_size(size_t size)7.5 bool get_free_block(AllocParams p)7.6 void garbage_collect_cached_blocks()7.7 bool release_available_cached_blocks(const AllocParams p)7.8 bool release_cached_blocks()7.11 bool alloc_block(AllocParams p, bool isRetry)7.13 bool is_split()7.14 bool should_split(const Block* block, size_t size)
6.2 void free 释放函数 以下内容来自PyTorch CUDA backend作者这种作图方式值得学习。 free的代码实现 void free(Block* block) {// 标记内存块为未分配状态block-allocated false;if (!block-stream_uses.empty()) {if (C10_UNLIKELY(captures_underway)) {// Its forbidden to cudaEventQuery an event recorded during CUDA graph// capture. We conservatively defer recording end-of-life events until// the next call to process_events() (which wont happen until no// captures are underway)needs_events_deferred_until_no_capture.push_back(block);} else {insert_events(block);}} else {// 如果没有活动事件则释放内存块并返回到池中。free_block(block);}// 报告内存使用情况// ……}free依赖的函数7.15 void free_block(Block* block) → 调用 7.16 size_t try_merge_blocks(Block* dst, Block* src, BlockPool pool)
7 功能函数
7.1 static bool BlockComparator 自定义比较函数
用于在BlockPool的std::set中对Block指针排序。
BlockComparator的代码实现
static bool BlockComparator(const Block* a, const Block* b) {// 按 stream 指针的地址排序升序。if (a-stream ! b-stream) {return (uintptr_t)a-stream (uintptr_t)b-stream;}// 若流相同按块大小排序升序。if (a-size ! b-size) {return a-size b-size;}// 若大小也相同按内存地址排序升序。return (uintptr_t)a-ptr (uintptr_t)b-ptr;
}7.2 static size_t round_size 对内存大小 size 进行调整
对size调整到合适的大小确保其满足块大小的最小值或接近划分点。
roundup_power2_next_division将一个数值size调整为靠近的 2 的次幂划分点多个划分点中的一个。例子
输入: size 1200, divisions 4过程: power2_floor 1024区间划分点: [1024, 1280, 1536, 1792]size 1200 在划分点 [1024, 1280) 内向上调整到 1280。 输出: 1280
kMinBlockSize * ((size kMinBlockSize - 1) / kMinBlockSize)。例子
输入: size 700, kMinBlockSize 512过程: size 大于 kMinBlockSize但划分点不足。对齐为 512 的倍数size 1024。 输出: 1024
round_size的代码实现 static size_t round_size(size_t size) {// 如果 size 小于最小块大小 kMinBlockSize直接返回 kMinBlockSize。512Bif (size kMinBlockSize) {return kMinBlockSize;} else {auto divisions CachingAllocatorConfig::roundup_power2_divisions(size);// 调用配置函数 roundup_power2_divisions(size) 获取划分点的数量。if (divisions 0 size (kMinBlockSize * divisions)) {return roundup_power2_next_division(size, divisions);// 如果划分点数大于 0 且 size 大于划分点范围调用 roundup_power2_next_division 调整。} else {// 否则将 size 向上舍入到最接近 kMinBlockSize512字节的倍数。// 用于确保分配的内存块大小是满足最小粒度要求的合适值。return kMinBlockSize * ((size kMinBlockSize - 1) / kMinBlockSize);}}}7.3 BlockPool get_pool 返回对应的内存块池
get_pool的代码实现
BlockPool get_pool(size_t size, cudaStream_t stream) {// kSmallSize 1 MiBif (size kSmallSize) {return small_blocks;} else {return large_blocks;}}7.4 static size_t get_allocation_size 确定实际要进行分配的内存大小 以下内容来自PyTorch显存管理介绍与源码解析一 显存管理机制是根据申请 size 来决定从 GPU 创建多大的 segment以及是否要进行切分split。基本执行逻辑如下图所示 以下内容来自Pytorch 1.6 显存管理CudaCachingAllocator剖析 Pytorch 向 cuda 申请显存时并不是严格按实际所需大小申请的而是按块大小申请的。块是一段地址连续的显存但块与块之间地址不一定连续。
块计算规则 Pytorch 按上述大小的块向 cuda 申请显存。如果直接将这些规格的块分给 Tensor 使用将会产生大量的内部碎片造成浪费。因此Pytorch 会对这些大小的块先尝试进行切分split。
Pytorch 向 cuda 申请块的方式是惰性的。当只有从管理的块池子中找不到满足要求的空闲块时才向 cuda 申请新的块。新申请的块会根据实际请求大小切分成更小的块。 文中还给了一些具体的例子可以有助于理解。 get_allocation_size的代码实现 static size_t get_allocation_size(size_t size) {if (size kSmallSize) { // 1 MiBreturn kSmallBuffer; // 2 MiB} else if (size kMinLargeAlloc) { // 10 MiBreturn kLargeBuffer; // 20 MiB} else {return kRoundLarge * ((size kRoundLarge - 1) / kRoundLarge); // 2 MiB}}7.5 bool get_free_block 寻找满足分配需求的空闲内存块
get_free_block尝试从指定内存池中BlockPool找到一个满足分配需求的空闲内存块Block并将其分配给传入的参数对象AllocParams。
get_free_block的代码实现 bool get_free_block(AllocParams p) { // 包含当前内存分配请求的参数// 获取当前操作的内存块池的引用后续操作都基于这个内存池。BlockPool pool *p.pool;if (C10_UNLIKELY(set_fraction CachingAllocatorConfig::garbage_collection_threshold() 0.0)) {// Track block reuse interval only when garbage collection is enabled.// 如果启用了垃圾回收GC每个空闲块的 gc_count 自增用于跟踪未使用的时间间隔。for (auto b : pool.blocks) {b-gc_count;}}// 从内存池 pool.blocks 中找到第一个大小不小于 p.search_key 的内存块。auto it pool.blocks.lower_bound(p.search_key);// 没有找到合适的块或者找到的块的流 (stream) 不匹配当前请求流 p.stream()。if (it pool.blocks.end() || (*it)-stream ! p.stream())return false;// 当前请求的块大小 p.size() 小于允许的最大拆分大小但找到的块太大超出 CachingAllocatorConfig::max_split_size()。// Do not return an oversized block for a large request// 目的避免浪费内存尽量分配更接近需求的块。if ((p.size() CachingAllocatorConfig::max_split_size()) ((*it)-size CachingAllocatorConfig::max_split_size()))return false;// 如果当前请求的大小较大 (p.size() max_split_size())但找到的块超出需求太多大于 p.size() kLargeBuffer。// Allow oversized block size to be rounded up but within a limit// 允许超大块大小向上取整但在一定范围内。if ((p.size() CachingAllocatorConfig::max_split_size()) ((*it)-size p.size() kLargeBuffer)) // 20 MiBreturn false;// 找到合适的块p.block *it;(*it)-gc_count 0; // Denote this block has been usedpool.blocks.erase(it);return true;}7.6 void garbage_collect_cached_blocks 释放未使用的缓存块
garbage_collect_cached_blocks的代码实现 void garbage_collect_cached_blocks() {// Unlike release_cached_blocks(), this does not enforce synchronization and// therefore should be of less overheads.// 避免了强制同步操作以降低性能开销// 计算垃圾回收的阈值size_t gc_threshold static_castsize_t(CachingAllocatorConfig::garbage_collection_threshold() *allowed_memory_maximum);// No need to trigger GC yet// 内存使用尚未超出阈值无需回收。if (total_allocated_memory gc_threshold) {return;}// 计算需要释放的目标内存量const auto target_size total_allocated_memory - gc_threshold;// 初始化已回收内存size_t gc_reclaimed 0;// Calculate the total age of the free-able blocks. Well use it later to get avg age threshold.// 计算可释放块的总年龄。稍后我们将使用它来获取“平均年龄”阈值。double total_age 0.0;int freeable_block_count 0;// 遍历当前的大块内存块列表找出可以回收的块for (auto b : large_blocks.blocks) {// 只有未拆分的大块才可回收。if (!b-is_split()) {// b-gc_count 表示该块未使用的累计时间用作垃圾回收的优先级。// 将所有可回收块的 gc_count 累加表示总的未使用时间。total_age b-gc_count; // 这里如果用真实时间呢// 记录可以回收的块总数。freeable_block_count;}}// No free-able blocks?// 没有可回收块直接返回if (freeable_block_count 0) {return;}// 初始化为 true表示可能会有块被释放。bool block_freed true;// Repeat GC until we reach reclaim target size.// 继续执行回收直到以下条件之一成立// 已回收的内存量gc_reclaimed达到或超过目标释放量target_size。// 当前循环中没有任何块被释放block_freed false。// 没有剩余的可回收块freeable_block_count 0。while (gc_reclaimed target_size block_freed true freeable_block_count 0) {// Free blocks exceeding this age threshold first.// 首先释放超过此年龄阈值的区块 - 优先释放未使用时间最长的块。double age_threshold total_age / freeable_block_count;// Stop iteration if we can no longer free a block.// 如果某轮循环中没有任何块被释放退出循环避免陷入无效操作。block_freed false;// Free blocks of avg age. Dont stop upon reaching the target_size,// we dont want this GC to be triggered frequently.// 即使达到了目标释放量target_size仍然尝试多释放一些内存减少垃圾回收频率。// 遍历当前的内存块列表检查每个块。auto it large_blocks.blocks.begin();while (it ! large_blocks.blocks.end()) {Block* block *it;it;// 块必须未拆分块的未使用时间需要大于等于年龄阈值。if (!block-is_split() block-gc_count age_threshold) {block_freed true;// 累计回收的内存量gc_reclaimed block-size;// Decrement the age// 从总未使用时间中减去该块的 gc_count。total_age - block-gc_count; // One less block that can be freed// 减少可回收块的计数。freeable_block_count--; // 释放内存块release_block(block);}}}}7.7 bool release_available_cached_blocks 优先从最大的块开始释放
release_available_cached_blocks的代码实现 // Free one or more oversize blocks to the system allocator. // But only enough to satisfy the target size.// 释放一个或多个超大块到系统分配器。但只足以满足目标大小。bool release_available_cached_blocks(const AllocParams p) {// 检查是否允许释放超大块if (CachingAllocatorConfig::max_split_size() std::numeric_limitssize_t::max())return false;// 提取目标内存池BlockPool pool *p.pool;// because of std::unique_ptr, block cannot be trivially copied// 由于std::unique_ptr块不能简单地复制// 创建一个 Block 对象 key用于搜索与当前需求匹配的块。Block key(p.search_key.device,p.search_key.stream,p.search_key.size,p.search_key.pool,p.search_key.ptr);// 调整目标大小 key.size避免释放小于超大块的块。key.size (key.size CachingAllocatorConfig::max_split_size())? CachingAllocatorConfig::max_split_size(): key.size;// 在内存池的块集合找到第一个大小大于或等于 key.size 的块。auto it pool.blocks.lower_bound(key);if (it pool.blocks.end() || (*it)-stream ! p.stream()) {// No single block is large enough; free multiple oversize blocks, starting with the largest// 没有一个块足够大释放多个oversize块从最大的开始if (it pool.blocks.begin())return false;// 累计已释放的内存大小。size_t totalReleased 0;// Back up one item. Now on the largest block for the correct stream// 将迭代器退回到前一个块确保从最大的块开始释放。--it; // 在以下条件满足时继续释放// 已释放的总大小 totalReleased 小于目标大小 key.size。// 当前块是超大块(*it)-size max_split_size。// 当前块的流与请求流匹配。 while ((totalReleased key.size) ((*it)-size CachingAllocatorConfig::max_split_size()) ((*it)-stream p.stream())) {auto cur it;totalReleased (*it)-size;// 如果当前块不是第一个块迭代器退回到前一个块继续释放。if (it ! pool.blocks.begin()) {--it;release_block(*cur);} else {// 如果已经到达集合起始位置退出循环。release_block(*cur);break;}}// 如果释放的总内存仍不足以满足目标大小返回 false。if (totalReleased key.size)return false;} else {// 如果搜索结果是一个单个块且满足需求直接调用 release_block 释放该块。release_block(*it);}// 如果成功释放了足够的内存无论是单个块还是多个块返回 true。return true;}7.8 bool release_cached_blocks 释放所有未使用的内存块
release_cached_blocks的代码实现 bool release_cached_blocks() {// First ensure that all blocks that cant currently be allocated due to// outstanding events are returned to the pool.// 有些内存块可能被异步 CUDA 操作如内核执行锁定。这些块不能立即被释放因为尚未完成的操作需要这些内存。调用 synchronize_and_free_events 确保所有与事件相关的操作完成后这些块可以被安全释放并返回到内存池中。synchronize_and_free_events();// Free all non-split cached blocks to system allocator// 释放所有未拆分的缓存块release_blocks(large_blocks);release_blocks(small_blocks);// 遍历并释放 CUDA 图的可释放内存池for (auto it graph_pools_freeable.begin();it ! graph_pools_freeable.end();) {// See notifyCaptureDestroy for the strategy here.TORCH_INTERNAL_ASSERT(it-second-use_count 0);release_blocks(it-second-small_blocks);release_blocks(it-second-large_blocks);if (it-second-cudaMalloc_count 0) {auto erase_count graph_pools.erase(it-first);TORCH_INTERNAL_ASSERT(erase_count 1);it graph_pools_freeable.erase(it);} else {it;}}return true;}函数调用关系release_cached_blocks → release_blocks → release_block
7.9 void release_blocks 释放所有非拆分块
release_blocks的代码实现 void release_blocks(BlockPool pool) {// Frees all non-split blocksauto it pool.blocks.begin();while (it ! pool.blocks.end()) {Block* block *it;it;// 非分割块可以直接释放if (!block-prev !block-next) {release_block(block);}}}7.10 void release_block 释放指定的内存块
free函数把 block 标记为未使用把 block 从active_blocks中清除之后将 block 回收到Large_blocks /Small_blocks中。这个过程不会触发cudaFree真正要释放掉一个 block 需要在release_block中完成。
release_block的代码实现 void release_block(Block* block) {// 释放该内存块的指针 block-ptr 所占用的 GPU 内存。C10_CUDA_CHECK(cudaFree((void*)block-ptr));// 减少已分配的总内存量。total_allocated_memory - block-size;// 如果块属于私有内存池释放后减少分配计数。auto* pool block-pool;if (pool-owner_PrivatePool) {// The cudaFreed block belonged to a CUDA graphs PrivatePool.TORCH_INTERNAL_ASSERT(pool-owner_PrivatePool-cudaMalloc_count 0);pool-owner_PrivatePool-cudaMalloc_count--;}// …… 更新内存统计信息// 从块所属的内存池pool-blocks中移除该块。pool-blocks.erase(block);// 释放 Block 对象本身占用的内存。// 这是分配器的元数据而非 GPU 内存。delete block;}7.11 bool alloc_block 为指定的分配请求分配一个内存块
alloc_block的代码实现 bool alloc_block(AllocParams p, bool isRetry) {// Defensively checks for preexisting CUDA error state.// 防御性地检查是否有未捕获的 CUDA 错误状态确保代码运行前的环境干净。C10_CUDA_CHECK(cudaGetLastError());size_t size p.alloc_size; // 需要分配的内存大小以字节为单位。void* ptr; // 分配后的内存指针初始值为未定义。// 记录重试分配的次数便于统计和调试。if (isRetry) {stats.num_alloc_retries 1;}// 检查当前总分配的内存是否会超过允许的最大内存。if (set_fraction // 启用内存限制机制total_allocated_memory size allowed_memory_maximum) {p.err cudaErrorMemoryAllocation;return false;} else {// 分配 CUDA 设备内存。p.err cudaMallocMaybeCapturing(ptr, size);if (p.err ! cudaSuccess) {if (p.err cudaErrorMemoryAllocation) {// If this is the first attempt (!isRetry), we can forgive and clear// CUDAs internal error state.//// If this is the second attempt (isRetry), mallocs TORCH_CHECK_WITH// will take over to throw a helpful exception. The user can choose// to catch the exception, free some stuff in their script, and// attempt the allocation again. In this case, we can also forgive and// clear CUDAs internal error state.cudaGetLastError();} else {// If the errors unrelated to memory allocation, we should throw immediately.// 如果错误与内存分配无关立即抛出异常。C10_CUDA_CHECK(p.err);}return false;}}// 如果分配的内存块属于某个 CUDA 图CUDA Graph的私有内存池。if (p.pool-owner_PrivatePool) {// The block is for a CUDA graphs PrivatePool.// 记录分配的调用次数。p.pool-owner_PrivatePool-cudaMalloc_count;}// 记录总分配的内存量。total_allocated_memory size;// 为分配的内存创建一个 Block 对象并存储在 p.block 中。p.block new Block(p.device(), p.stream(), size, p.pool, (char*)ptr);// 更新内存分配的统计信息。for_each_selected_stat_type(p.stat_types, [](size_t stat_type) {// 分配的内存块数量。update_stat(stats.segment[stat_type], 1);// 分配的总字节数。update_stat(stats.reserved_bytes[stat_type], size);});// 记录过大的分配。if (size CachingAllocatorConfig::max_split_size())update_stat(stats.oversize_segments, 1);// p.block came from new, not cudaMalloc. It should not be nullptr here.// 验证内存块 p.block 和其指针 p.block-ptr 是否有效。TORCH_INTERNAL_ASSERT(p.block ! nullptr p.block-ptr ! nullptr);return true;}7.12 cudaError_t cudaMallocMaybeCapturing 一个封装的内存分配工具
cudaMallocMaybeCapturing的代码实现
cudaError_t cudaMallocMaybeCapturing(void** p, size_t size) {
#if !defined(USE_ROCM) || ROCM_VERSION 50300if (at::cuda::currentStreamCaptureStatusMayInitCtx() at::cuda::CaptureStatus::None) {
#endif// 如果返回值为 None表示未进入 CUDA 图捕获模式执行普通内存分配。return C10_CUDA_ERROR_HANDLED(cudaMalloc(p, size));
#if !defined(USE_ROCM) || ROCM_VERSION 50300} else {// Its ok to capture cudaMallocs, as long as we never cudaFree those// addresses before replay.// Capturing cudaMalloc behaves nicely: it gives the graph new VA,// but is ignored (wont leakily allocate new memory) in replays.// 在捕获模式下CUDA 图会为 cudaMalloc 提供新的虚拟地址Virtual AddressVA但在图重放时不会实际重新分配物理内存。at::cuda::CUDAStreamCaptureModeGuard g{cudaStreamCaptureModeRelaxed};return C10_CUDA_ERROR_HANDLED(cudaMalloc(p, size));}
#endif
}7.13 bool is_split 判断 block 是否被拆分过 以下内容来自一文读懂 PyTorch 显存管理机制 当 Block 被释放时会检查其 prev、next 指针是否为空及若非空是否正在被使用。若没有在被使用则会使用try_merge_blocks合并相邻的 Block。由于每次释放 Block 都会检查因此不会出现两个相邻的空闲块于是只须检查相邻的块是否空闲即可。 is_split的代码实现 bool is_split() const {return (prev ! nullptr) || (next ! nullptr);}7.14 bool should_split 判断给定的内存块 block 是否应该进行分割 以下内容来自Pytorch 1.6 显存管理CudaCachingAllocator剖析 在决定了申请的块由哪个池子管理之后Pytorch 通过should_split()来判断申请的块是否应该切分。规则如下
对于 2M 的块如果说切分后剩余的大小小于 512B那么就不进行切分了直接把整个块给用户好了对于大于 2M 的块如果说切分后剩余的大小小于了 1M 了那么就不进行切分了直接把整个块给用户好了。
这样做的好处是可以保证 Small BlockPool 中所有空闲的块的大小至少是 512BLarge BlockPool 中所有空闲的块至少是 1M 大小。 should_split的代码实现 bool should_split(const Block* block, size_t size) {// 计算出如果进行分配操作后该内存块剩余的空间大小size_t remaining block-size - size;if (block-pool-is_small) {return remaining kMinBlockSize; // 512B} else {// 要分配的内存大小 size 必须小于最大分割尺寸// 剩余空间大小要大于 kSmallSizereturn (size CachingAllocatorConfig::max_split_size()) (remaining kSmallSize); // 1 MiB}}7.15 void free_block 将一个块移动到缓存的空闲块池中
free_block的代码实现 /** moves a block into a pool of cached free blocks */void free_block(Block* block) {// 确保 block 是可释放的。TORCH_INTERNAL_ASSERT(!block-allocated block-event_count 0 block-stream_uses.empty());// 获取其所属的内存池auto pool *block-pool;// 用于记录内存块合并等操作导致的非活跃拆分内存块数量和大小净变化int64_t net_change_inactive_split_blocks 0;int64_t net_change_inactive_split_size 0;// 数组 merge_candidates 包含了要释放的内存块 block 的前一个内存块和后一个内存块const std::arrayBlock*, 2 merge_candidates {block-prev, block-next};for (Block* merge_candidate : merge_candidates) {// 尝试进行合并操作并返回被合并的内存块大小const int64_t subsumed_size try_merge_blocks(block, merge_candidate, pool);// 成功合并了一个内存块更新前面定义的统计净变化变量。if (subsumed_size 0) {net_change_inactive_split_blocks - 1;net_change_inactive_split_size - subsumed_size;}}// 从活跃块列表中移除active_blocks.erase(block);// 插入到空闲池中bool inserted pool.blocks.insert(block).second;TORCH_INTERNAL_ASSERT(inserted);// 如果 block 是分裂块增加分裂块的计数。增加分裂块的总大小。if (block-is_split()) {net_change_inactive_split_blocks 1;net_change_inactive_split_size block-size;}// 更新统计信息// ……}7.16 size_t try_merge_blocks 合并先前拆分的块 以下内容来自PyTorch显存管理介绍与源码解析一 当用户不需要使用某个显存块时显存管理机制并不会直接从 GPU 设备上删除free该块而是将其先回收到 BlockPool。
在 block 块释放后会触发一种块的融合机制逻辑如下所示当释放一个 Block_A 显存块时去池子里面寻找是否有空闲的、地址与之连续的 Block当匹配到符合条件的 Block_C 时触发了合并操作这样 Block_A 和 Block_C 融合成了 Block_D。这种机制能够降低显存碎片问题。 try_merge_blocks的代码实现 // returns the size of the subsumed block, or 0 on failure.// 返回包含块的大小失败时返回 0。size_t try_merge_blocks(Block* dst, Block* src, BlockPool pool) {// 检查可合并条件dst 是 blocksrc 是准备和 block 合并的块 if (!src || src-allocated || src-event_count 0 ||!src-stream_uses.empty()) {return 0;}// 进一步确保 dst 和 src 两个内存块都是拆分块AT_ASSERT(dst-is_split() src-is_split());if (dst-prev src) { // [src dst] 合并 block 前的块// dst 内存块指向合并后内存块应该指向的内存起始位置实现内存区域的合并效果。dst-ptr src-ptr;// 构建合并后内存块在链表中的前向连接关系dst-prev src-prev;if (dst-prev) {dst-prev-next dst;}} else { // [dest src] 合并 block 后的块// 构建合并后内存块在链表中的后向连接关系dst-next src-next;if (dst-next) {dst-next-prev dst;}}const size_t subsumed_size src-size;// 两个内存块合并后的实际内存占用情况dst-size subsumed_size;// 从内存块所属的内存池中移除被合并的内存块auto erased pool.blocks.erase(src);TORCH_INTERNAL_ASSERT_DEBUG_ONLY(erased 1);// 释放被合并内存块 src 的资源delete src;// 返回被合并内存块的大小return subsumed_size;}8 问题与挑战 以下内容来自PyTorch显存管理介绍与源码解析一 当前机制下的一些问题
机制的参数为静态参数没有根据需求进行自动调整。 比如 Large_blocks 和 Small_blocks 分类的阈值、切分的阈值合并回收的触发方式单一只有显存不足或者手工调用时才会触发回收没有一个算法对回收时机进行分析当前机制下有不可避免的碎片问题。 还可以参考 深入剖析大模型训练的显存优化 提到的五种设备内存优化路径。
参考资料 说明资料内容依据的 PyTorch 版本见标题前[]最新版本的源码实现还请查看 PyTorch 仓库。 [unknown] PyTorch源码浅析(2)THC | NIUHETHCCachingAllocator 小节[1.5.1] PyTorch-CUDA端显存池函数解读[1.9.1] PyTorch显存机制分析提到 CUDA Context 开销 → 补充 [1.10.0] PyTorch显存机制分析 - Angry_Panda[1.12.0] 一文读懂 PyTorch 显存管理机制 ⭐[unknown] GPU 底层机制分析显存分配开销[unknown] PyTorch CUDACachingAllocator[1.13.0] PyTorch显存管理介绍与源码解析一 ⭐ [1.13.0] PyTorch显存管理介绍与源码解析二 ⭐ [1.6.0] Pytorch 1.6 显存管理CudaCachingAllocator剖析 ⭐提到 CUDA Context 开销还分析了为什么没有必须显式调用torch.cuda.empty_cache()[unknown] Pytorch内存管理机制小记文中分析max_spilt_size_mb的部分写得较清晰[unknown] PyTorch CUDA backend分析了 CUDA stream、CUDA event 和 CUDA graph 等[2.3.0] Pytorch 显存管理机制与显存占用分析方法 ⭐
