推理框架学习笔记

序言

下面主要是关于训练推理框架中资源管理的一些内容，先是介绍了一下缓存管理的必要性，其次是介绍一些框架的资源缓存分配策略，最后是一个关于框架资源分配的简易实现逻辑。

资源管理类的设计：

一.必要性

1. device上存储管理api中cudaFree()会导致同步操作,这是因为 cudaFree() 函数需要等待所有与该 CUDA 内存缓冲区相关的 CUDA 操作完成后,才能释放该内存缓冲区。这个同步过程是必要的,因为 CUDA 运行时需要确保在释放内存之前,所有与该内存相关的操作都已完成。否则,如果有正在使用该内存的 CUDA 操作仍在运行,就可能会导致内存访问错误或其他问题。
   因此,在使用 cudaFree() 时需要格外小心,因为它会导致 CUDA 流水线停滞,无法依赖Host和Device的异步执行特性来隐藏Host端的代码耗时，从而影响应用程序的整体性能。
   对大部分训练来说，在迭代稳定的情况下，卡上的显存分配也是稳定的，因此可以通过一些缓存手段来空间换时间，从而避免上述打断异步的情况。

二.资源分配策略

1.对于请求的size，首先要rounding到一个合适的size，通常是2的幂次方之间的一个值，通常先是round to 512B，或者通过环境变量设置roundup_power2_divisions,例如需要分配1200，划分数为4，大小1200位于2的幂次方1024和2048之间，在两者之间进行4次除法，它们的值为1024,1280,1536,和1792，因此返回1280作为round之后的size用来请求分配，然后按照最佳的拟合策略来查找适合分配的最小快。如果找不到可用的，说明还未稳定，需要cudaMalloc一下，这时也会触发round_size，只向cuda请求2MB的倍数，并有可能将其分割更小的块进行多次分配，从而减少cudaMalloc的次数，增加框架再分配次数，
当OOM时，比如说给一个声明周期较长的小tensor分配了较大的块，动态计算的时候后向需要一些前向的数据，不得不将相应的生命周期拉长，导致存储被切割，为了解决溢出问题，会尝试释放缓存，然后再通过cudaMalloc申请对于分配的大小大于需求，对于小于1MB和1MB-10MB的申请，分配的块可能会需要再split，split是按stream记录的，避免复用造成踩踏。同一stream分配和释放是顺序执行的，stream之间为了顺序保证，需要调用record_stream()让分配器知道使用情况，为了剩余的块会添加到对应的池中并可用于另一次分配。释放会将块返回到pool中，当返回一个块时，如果相邻块也空闲，就会将它们合并成一个更大的块。NVIDIA GPU通过页表机制在逻辑上能够获得更好的合并，碎片化程度也会降低。

三.框架资源分配简易实现逻辑

1. 由于要支持模型既可以在上推理，也可以在GPU上推理，所以资源的分配以及释放抽象为一个顶级父类，提供统一的接口满足在不同设备上的资源管理；
2. DeviceAllocator为父类抽象类，两个子类CPUDeviceAllocator、CUDA -DeviceAllocator分别实现allocate()和release()接口来满足CPU、GPU上的内存分配和释放，CPUDeviceAllocatorFactory和CDeviceAllocatorFactory用来返回CPUDeviceAllocator、CUDADeviceAllocator全局单例，CudaMemoryBuffer为片上不同设备id之间池化出来的内存单元，用一个vector容器来装载buffer，buffer类中包括三个成员属性，一个void*指针用来指向显存中的位置，以及表示cudabuffer多大，是否正被占用标志。
3. CUDADeviceAllocator里维护了下面三个数据结构：
   3.1 map<int, size_t> no_busy_cnt_;用来表示不同设备之间空闲内存大小，如果超过1G大小就要进行清理。
   3.2 map<int, std::vector> big_buffers_map_;一个大内存池，当待分配大小超过1M时，优先从大内存池中选取合适大小的buffer(大内存池中空闲buffer大小大于待分配大小且两者不超过1M大小），如果没有合适的或者池中buffer大小小于待分配大小，则进行cudamalloc重新开辟一段内存并放入池中。
   3.3 map<int, std::vector> cuda_buffers_map_;一个小内存池，当待分配大小较小时，从小内存池中寻找未使用的buffer，从而避免内存碎片，如果不成功则会cudaMalloc请求更多的存储。

四.一些辅助资料

1.PyTorch显存可视化与Snapshot数据分析
2.CUDA Memory Management