PyTorch

c10/cuda/CUDACachingAllocator.hc10/cuda/CUDACachingAllocator.cpp

更多内…

c10/cuda/CUDACachingAllocator.hc10/cuda/CUDACachingAllocator.cpp

更多内容请参考

PyTorch【翻译】pytorch/CONTRIBUTING.md【翻译】Pytorch机制源代码分析与内存管理调研深度学习框架与静态/动态计算图【笔记】PyTorch

源码学习阅读经验

release_available_cached_blocks

bool

缓存分配器。

其目的是通过缓存和复用内存块来减少频繁的cudaMalloc和cudaFree操作提升

CUDA

原生显存分配接口逻辑比较简单但显存分配开销无法接受显存分配碎片也需要优化。

显存管理机制与显存占用分析方法这篇文章对显存管理机制的总结写得很好。

GPU

作为一种通用的数据处理设备为了满足更广泛客户的需求且保证更小的维护成本其

API

API但并没有针对某个深度学习框架做设计优化其中显存的精细管理留给上层的深度学习框架去完成。

cudaMallocCUDA

内的同步操作。

当深度学习框架使用的数据非常零碎且数量多时需要反复调用cudaMalloc该行为会直接影响程序的整体性能因此深度学习框架的显存管理机制在设计时要尽量降低cudaMalloc的调用频次。

PyTorch

实现了一套显存管理逻辑/机制可更好地满足框架的日常使用需求相比原生的

CUDA

可做到管理细化、使用相对高效其采用动态申请与二次分配的设计思路

GPU

发出请求最大优点是不会占用过量的显存方便多人同时使用一个设备与之相对的是

TensorFlow

上的大部分显存都申请到然后再去分配使用。

二次分配将显存的申请与使用进行分离即显存申请后会进行二次分配。

显存管理机制会先通过cudaMalloc向

GPU

{public:std::vectorstd::unique_ptrDeviceCachingAllocator

DeviceCachingAllocator以下内容来自PyTorch显存管理介绍与源码解析二

显存管理代码中的类的大体关系如下图所示。

cudaCachingAllocator类的结构相对来说比较清晰代码不是特别清晰DeviceCachingAllocator管理一个设备显存THCCachingAllocator管理一个进程上所有的DeviceCachingAllocatorCudaCachingAllocator继承自Allocator方便框架调用。

DeviceCachingAllocator类是显存管理机制真正实现的地方它负责完成

device

1.13.0因为在该版本的代码注释中提到THCCachingAllocator这种说法已经是

old-style

了。

作者这里提到的THCCachingAllocator在最新的版本里对应的是NativeCachingAllocator而CudaCachingAllocator对应的是CUDAAllocator。

但上面的内容仍值得参考。

1.6

设备卡都维护一个这样的结构用于对该设备进行显存管理THCCachingAllocator维护一个

Pytorch

默认实现的分配器包装器向上层暴露一些有用的接口核心逻辑还是委托给

THCCachingAllocator

有一个与之对应的DeviceCachingAllocator实例管理同时记录下进程中创建的全部的

class

pointerska::flat_hash_mapvoid*,

Block*

实例用于管理该设备上的显存分配std::vectorstd::unique_ptrDeviceCachingAllocator

device_allocator;//

{allocated_blocks.erase(it);}return

block;}//

static_castint64_t(device_allocator.size());if

(size

{device_allocator.resize(device_count);for

(const

std::make_uniqueDeviceCachingAllocator();}}}//

allocated_blocks

检查设备号是否有效TORCH_INTERNAL_ASSERT(0

device

device_allocator[device]-malloc(device,

size,

中add_allocated_block(block);}//

free(void*

释放显存块device_allocator[block-device]-free(block);}//

分配显存并返回一个

1152921504606846976ULL;TORCH_CHECK_WITH(OutOfMemoryError,size

one_exa_bytes,CUDA

device;C10_CUDA_CHECK(cudaGetDevice(device));void*

nullptr;//

capturing.C10_CUDA_CHECK(cudaMalloc(r,

size));return

const!?const_castNativeCachingAllocator*(this)-malloc(r,

device,

cuda::getCurrentCUDAStream(device));}return

{r,

device;C10_CUDA_CHECK(cudaGetDevice(device));void*

nullptr;//

cuda::getCurrentCUDAStream(device));return

r;}//

device;C10_CUDA_CHECK(cudaGetDevice(device));void*

nullptr;//

{this-free(ptr);}};这里要区分下DataPtr

allocate(size_t

add_allocated_block。

在分配显存后会将显存块添加到allocated_blocks中并返回一个

DataPtr

allocate和raw_alloc两者的使用场景还是有很大区别的因本人精力有限暂不深入研究了。

一些线索

allocate(n);AT_ASSERT(dptr.get()

raw_deallocate(void*

raw_deleter();AT_ASSERT(d);d(ptr);}

};2class

Block不论是否为空闲只要是由Allocator::malloc得来的都被组织在一个双向链表里便于在释放某一个

Block

device(device),stream(stream),stream_uses(),size(size),requested_size(0),pool(pool),ptr(ptr)

{}//

device(device),stream(stream),stream_uses(),size(size),requested_size(0)

{}//

是通过一次cudaMalloc调用分配的一块连续的显存区域。

一个内存段

(Segment)

allocation。

std::vectorBlockInfo

blocks;

用户实际请求的大小。

分配器通常会将请求大小对齐到某个块大小。

int32_t

gc_counter

在管理机制中将显存的申请与使用过程进行了分离即显存申请后会进行二次分配其过程是先通过cudaMalloc申请一个显存块

segment然后从

内存管理机制里对显存分了两级segment、blocks显存池分了两类active

pool、remaining

结构体来完成承上启下的功能显存池结构体blockPool主要是记录

block

使用的块即torch.cuda.memory_allocated()的值cached

Pytorch

进程管理的所有显存块即torch.cuda.memory_reserved()的值reserved

allocated

BlockPool用于管理多个Block可以分为小块池和大块池。

PyTorch

都是空闲的。

DeviceCachingAllocator中维护两种

BlockPool

small_blocks)分别存放较小的块和较大的块为了分别加速小需求和大需求简单地将

1MB

目前未使用的块管理有两种类型的池子large_blockssmall_blocks。

分类阈值默认值设置为

1M小于

ska::flat_hash_setBlock*基于哈希表的集合。

O(1)

O(1)适合频繁的动态操作。

active_blocks目的是快速跟踪所有活跃的内存块从代码来看active_blocks主要用于

large_blocks和small_blocks是BlockPool类型的对象内部使用std::setBlock*,

Comparison来存储和排序

O(logn)适合需要排序的场景。

large_blocks和small_blocks目的是高效管理空闲内存块。

1.6

Pool笔者有一些和上述博客不同的看法在这里根据笔者的理解做一些补充如果有不对之处还请指出。

Small

但从alloc_block这个函数的具体实现来看其实并没有将新创建的

Block

owner_PrivatePool(private_pool)

};关于

计数都降为零时我们可以从图内存池graph_pools中删除这个私有内存池PrivatePool。

BlockPool

small_blocks;//

use_count(1),cudaMalloc_count(0),//

初始化时传入比较函数。

large_blocks(BlockComparator,

/*is_small*/false,

this),small_blocks(BlockComparator,

/*is_small*/true,

delete;PrivatePool(PrivatePool)

delete;PrivatePool

AllocParams封装分配内存时的参数和状态用于寻找或创建合适的内存块。

struct

size),pool(pool),alloc_size(alloc_size),block(nullptr),err(cudaSuccess)

{}//

search_key.device;}cudaStream_t

stream()

默认构造函数DeviceCachingAllocator():

初始化大块内存池small_blocks(BlockComparator,

/*is_small*/true),//

CachingAllocatorConfig::max_split_size();context_recorder_.store(nullptr);}}结构体和类的区别Struct

Class

block如果有直接跳入步骤3如果没有跳入步骤2创建调用cudaMalloc创建新的

segment如果满足

根据需求进行切分如果有剩余跳入步骤4如果没有剩余跳入步骤5保存将切分后的

block

尝试直接从内存池中获取空闲块。

get_free_block(params)//

Trigger

(trigger_free_memory_callbacks(params)

Cant

CachingAllocatorConfig::garbage_collection_threshold()

0.0))

清理不再使用的缓存块。

garbage_collect_cached_blocks();}//

Attempt

(release_available_cached_blocks(params)

alloc_block(params,

内存块params.block和其指针params.block-ptr不为空。

TORCH_INTERNAL_ASSERT(params.err

cudaSuccess

static_castchar*(remaining-ptr)

size;remaining-size

pool.blocks.insert(remaining).second;TORCH_INTERNAL_ASSERT_DEBUG_ONLY(inserted);//

……}

如果块不需要分割但已经是分割块则直接更新统计信息表示该块变为“活动”块。

//

……}//

active_blocks.insert(block).second;TORCH_INTERNAL_ASSERT_DEBUG_ONLY(inserted);//

更新统计数据//

garbage_collect_cached_blocks()7.7

bool

release_available_cached_blocks(const

AllocParams

(C10_UNLIKELY(captures_underway))

{//

underway)needs_events_deferred_until_no_capture.push_back(block);}

else

如果没有活动事件则释放内存块并返回到池中。

free_block(block);}//

报告内存使用情况//

用于在BlockPool的std::set中对Block指针排序。

static

对size调整到合适的大小确保其满足块大小的最小值或接近划分点。

roundup_power2_next_division将一个数值size调整为靠近的

输入:

CachingAllocatorConfig::roundup_power2_divisions(size);//

调用配置函数

roundup_power2_next_division(size,

divisions);//

用于确保分配的内存块大小是满足最小粒度要求的合适值。

return

kMinBlockSize

segment以及是否要进行切分split。

基本执行逻辑如下图所示

1.6

申请显存时并不是严格按实际所需大小申请的而是按块大小申请的。

块是一段地址连续的显存但块与块之间地址不一定连续。

块计算规则

申请块的方式是惰性的。

当只有从管理的块池子中找不到满足要求的空闲块时才向

cuda

get_free_block尝试从指定内存池中BlockPool找到一个满足分配需求的空闲内存块Block并将其分配给传入的参数对象AllocParams。

bool

获取当前操作的内存块池的引用后续操作都基于这个内存池。

BlockPool

pool

CachingAllocatorConfig::garbage_collection_threshold()

0.0))

pool.blocks.lower_bound(p.search_key);//

(stream)

CachingAllocatorConfig::max_split_size()。

//

not

CachingAllocatorConfig::max_split_size())

((*it)-size

CachingAllocatorConfig::max_split_size()))return

false;//

CachingAllocatorConfig::max_split_size())

((*it)-size

usedpool.blocks.erase(it);return

true;}7.6

garbage_collect_cached_blocks的代码实现

void

garbage_collect_cached_blocks()

{//

static_castsize_t(CachingAllocatorConfig::garbage_collection_threshold()

need

计算可释放块的总年龄。

稍后我们将使用它来获取“平均年龄”阈值。

double

total_age

记录可以回收的块总数。

freeable_block_count;}}//

free-able

已回收的内存量gc_reclaimed达到或超过目标释放量target_size。

//

(gc_reclaimed

如果某轮循环中没有任何块被释放退出循环避免陷入无效操作。

block_freed

false;//

即使达到了目标释放量target_size仍然尝试多释放一些内存减少垃圾回收频率。

//

large_blocks.blocks.begin();while

(it

减少可回收块的计数。

freeable_block_count--;

释放内存块release_block(block);}}}}7.7

bool

release_available_cached_blocks

release_available_cached_blocks的代码实现

Free

释放一个或多个超大块到系统分配器。

但只足以满足目标大小。

bool

release_available_cached_blocks(const

AllocParams

(CachingAllocatorConfig::max_split_size()

std::numeric_limitssize_t::max())return

false;//

key(p.search_key.device,p.search_key.stream,p.search_key.size,p.search_key.pool,p.search_key.ptr);//

调整目标大小

CachingAllocatorConfig::max_split_size())?

CachingAllocatorConfig::max_split_size():

key.size;//

pool.blocks.lower_bound(key);if

(it

CachingAllocatorConfig::max_split_size())

((*it)-stream

如果已经到达集合起始位置退出循环。

release_block(*cur);break;}}//

(totalReleased

操作如内核执行锁定。

这些块不能立即被释放因为尚未完成的操作需要这些内存。

调用

确保所有与事件相关的操作完成后这些块可以被安全释放并返回到内存池中。

synchronize_and_free_events();//

Free

释放所有未拆分的缓存块release_blocks(large_blocks);release_blocks(small_blocks);//

遍历并释放

graph_pools_freeable.begin();it

{//

here.TORCH_INTERNAL_ASSERT(it-second-use_count

0);release_blocks(it-second-small_blocks);release_blocks(it-second-large_blocks);if

{auto

graph_pools.erase(it-first);TORCH_INTERNAL_ASSERT(erase_count

1);it

graph_pools_freeable.erase(it);}

else

true;}函数调用关系release_cached_blocks

release_blocks

/Small_blocks中。

这个过程不会触发cudaFree真正要释放掉一个

block

内存。

C10_CUDA_CHECK(cudaFree((void*)block-ptr));//

减少已分配的总内存量。

total_allocated_memory

block-size;//

PrivatePool.TORCH_INTERNAL_ASSERT(pool-owner_PrivatePool-cudaMalloc_count

0);pool-owner_PrivatePool-cudaMalloc_count--;}//

更新内存统计信息//

从块所属的内存池pool-blocks中移除该块。

pool-blocks.erase(block);//

Block

错误状态确保代码运行前的环境干净。

C10_CUDA_CHECK(cudaGetLastError());size_t

size

cudaErrorMemoryAllocation;return

false;}

如果错误与内存分配无关立即抛出异常。

C10_CUDA_CHECK(p.err);}return

false;}}//

记录分配的调用次数。

p.pool-owner_PrivatePool-cudaMalloc_count;}//

记录总分配的内存量。

total_allocated_memory

size;//

更新内存分配的统计信息。

for_each_selected_stat_type(p.stat_types,

[](size_t

分配的内存块数量。

update_stat(stats.segment[stat_type],

1);//

分配的总字节数。

update_stat(stats.reserved_bytes[stat_type],

size);});//

CachingAllocatorConfig::max_split_size())update_stat(stats.oversize_segments,

1);//

是否有效。

TORCH_INTERNAL_ASSERT(p.block

nullptr

cudaMallocMaybeCapturing(void**

size_t

(at::cuda::currentStreamCaptureStatusMayInitCtx()

#endif//

C10_CUDA_ERROR_HANDLED(cudaMalloc(p,

size));

AddressVA但在图重放时不会实际重新分配物理内存。

at::cuda::CUDAStreamCaptureModeGuard

g{cudaStreamCaptureModeRelaxed};return

C10_CUDA_ERROR_HANDLED(cudaMalloc(p,

size));}

指针是否为空及若非空是否正在被使用。

若没有在被使用则会使用try_merge_blocks合并相邻的

Block

都会检查因此不会出现两个相邻的空闲块于是只须检查相邻的块是否空闲即可。

bool

通过should_split()来判断申请的块是否应该切分。

规则如下

Small

CachingAllocatorConfig::max_split_size())

(remaining

是可释放的。

TORCH_INTERNAL_ASSERT(!block-allocated

block-event_count

用于记录内存块合并等操作导致的非活跃拆分内存块数量和大小净变化int64_t

net_change_inactive_split_blocks

0;int64_t

{net_change_inactive_split_blocks

1;net_change_inactive_split_size

subsumed_size;}}//

从活跃块列表中移除active_blocks.erase(block);//

插入到空闲池中bool

pool.blocks.insert(block).second;TORCH_INTERNAL_ASSERT(inserted);//

block

{net_change_inactive_split_blocks

1;net_change_inactive_split_size

block-size;}//

两个内存块都是拆分块AT_ASSERT(dst-is_split()

src-is_split());if

内存块指向合并后内存块应该指向的内存起始位置实现内存区域的合并效果。

dst-ptr

src-ptr;//

pool.blocks.erase(src);TORCH_INTERNAL_ASSERT_DEBUG_ONLY(erased

1);//

分类的阈值、切分的阈值合并回收的触发方式单一只有显存不足或者手工调用时才会触发回收没有一个算法对回收时机进行分析当前机制下有不可避免的碎片问题。

参考资料

开销还分析了为什么没有必须显式调用torch.cuda.empty_cache()[unknown]

Pytorch内存管理机制小记文中分析max_spilt_size_mb的部分写得较清晰[unknown]

PyTorch