大模型之不同精度下显存占用与相互转换实践

文章预览

【点击】 加入大模型技术交流群 原文：https://zhuanlan.zhihu.com/p/658343628 之前讲述了 大模型精度（FP16，FP32，BF16）详解与实践 ，本文围绕1.模型不同精度下显存占用情况；2.模型不同精度之间如何转换，两个问题进行一个简单的实践。 模型不同精度下显存占用 以显卡NVIDIA A40 48G，模型用llama-2-7b-hf为例 [1] ，这个模型的保存的精度通过查看模型文件的congfig.json可以看到是"torch_dtype": "float16"。 首先打印相关的版本和显卡信息： import transformers from transformers import AutoModelForCausalLM, AutoTokenizer import torch # 打印版本号 print ( "transformers version:" , transformers . __version__) print ( "torch version:" , torch . __version__) # 检查系统中是否有可用的 GPU if torch . cuda . is_available(): # 获取可用的 GPU 设备数量 num_devices = torch . cuda . device_count() print ( "可用 GPU 数量:" , num_devi ………………………………