Pytorch 显存管理机制与显存占用分析方法

文章预览

【点击】 加入大模型技术交流群 原文：https://zhuanlan.zhihu.com/p/699254132 1、显存管理机制 GPU 作为一种通用的数据处理设备，为了满足更广泛客户的需求且保证更小的维护成本，其 API 在设计的时候比较开放，尽管 CUDA 生态中也有高阶 API，但并没有针对某个深度学习框架做设计优化，其中显存的精细管理留给上层的深度学习框架去完成。 cudaMalloc（CUDA API）是从 GPU 申请显存最常用的方式，给定指针和数据大小即可进行 API 调用，其调用有着不小的时间开销，且是 stream 内的同步操作。当深度学习框架使用的数据非常零碎且数量多时，需要反复调用 cudaMalloc，该行为会直接影响程序的整体性能，因此深度学习框架的显存管理机制在设计时要尽量降低 cudaMalloc 的调用频次。 PyTorch 框架基于 CUDA API 实现了一套显存管理逻辑/机制，可更好地满足框架的日常使用 ………………………………