如何估算LLM推理和训练所需的GPU内存？

作者：孙鹏飞，南京大学 · 计算机科学与技术
原文：https://zhuanlan.zhihu.com/p/716317173

在实际工作中，经常有人问，7B、14B或70B的模型需要多大的显存才能推理？如果微调他们又需要多大的显存呢？为了回答这个问题整理一份训练或推理需要显存的计算方式。如果大家对具体细节不感兴趣，可以直接参考经验法则评估推理或训练所需要的资源。更简单的方式可以通过这个工具[1]或者huggface官网计算推理/训练需要的显存工具[2]在线评估。

数据精度

开始介绍之前，先说一个重要的概念——数据精度。数据精度指的是信息表示的精细程度，在计算机中是由数据类型和其位数决定的。如果想要计算显存，从“原子”层面来看，就需要知道我们的使用数据的精度，因为精度代表了数据存储的方式，决定了一个数据占多少bit。目前，精度主要有以下几种：

• • 4 Bytes: FP32 / float32 / 32-bit

• • 2 Bytes: FP16 / float16 / bfloat16 / 16-bit

• • 1 Byte: int8 / 8-bit

• • 0.5 Bytes: int4 / 4-bit

经验法则

• • 推理: 参数量 * 精度。

例如，假设模型都是16-bit权重发布的，也就是说一个参数消耗16-bit或2 Bytes的内存，模型的参数量为70B，基于上述经验法则，推理最低内存需要70B * 2Bytes = 140G。

• • 训练: 4 - 6 倍的推理资源。

推理

在模型推理阶段，需要的资源主要有三部分：模型的权重、KV Cache和激活（在推理过程中创建的张量）。

模型权重

加载模型权重（即模型大小）占用资源主要依赖于模型的参数量和