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

文章预览

作者 | 不理不理  编辑 | 极市平台 原文链接： https://zhuanlan.zhihu.com/p/699254132 点击下方 卡片 ，关注“ 自动驾驶Daily ”公众号 戳我->  领取近15个自动驾驶方向路线 >> 点击进入→ 自动驾驶Daily技术交流群 本文只做学术分享，如有侵权，联系删文 近期在研究大模型推理加速框架 VLLM 源码的过程中，对 Pytorch 的显存占用和分配机制十分感兴趣，因此花了一些时间研究和测试。写作本文，既是笔记，也是分享。 1. 前言 1.1 设备及版本 操作系统：Ubuntu 22.04 驱动版本：535.161.08 GPU：A800-SXM4-80GB CUDA：12.1 Pytorch：2.3.0 Python：3.10.6 得益于社区的不懈努力，PyTorch 的显存管理机制一直在不断优化。尽管不同版本的显存管理机制在核心思路上保持一致，但在细节上可能会略有差异。本文关于显存管理机制的内容基于 Pytorch 2.3.0 版本，文章的最后也提供了显存管理机 ………………………………