大语言模型LLM推理及训练显存计算方法

发布时间 2023-10-03 20:39:03作者: tommickey

一、推理:显存计算

推理的显存大头就是:参数量,参数类型版本一般有以下四种: 

  1. float 32位 浮点数 4 字节
  2. half / BF16 16位 浮点数 2 字节
  3. int8 8位 整数 1 字节
  4. int4 4位 整数 0.5 字节

 以 7B-BF16 版本为例,需要显存 = 数量 * 类型大小 = 70亿 * 2字节 = 140亿字节

 所以 140亿 字节 = 14 * 1000 * 1000 * 1000 / 1024 / 1024 / 1024 = 14 * (1000 / 1024) ^3 = 14 * 0.93 = 13 GB

 

注1: (1000/1024)^3 = 0.93...

注2: 估算的话,干脆直接记 1就好了,也就是 7B-BF 就是 7 * 2 = 14 GB 的显存,因为除了参数,推理还需要另外的显存,溢出一点没坏处。

按估算,各个 llama2-13B 版本推理大概用到多少显存(假设存在对应类型)

 

7B参数的4种常见类型所需推理显存计算如下:

  1. float 7 * 4 = 28 GB
  2. half / BF16 7 * 2 = 14 GB
  3. int8 7 * 1 = 7 GB
  4. int4 7 * 0.5 = 3.5 GB

13B参数的4种常见类型所需推理显存计算如下:

  1. float 13 * 4 = 52 GB
  2. half / BF16 13 * 2 = 26 GB
  3. int8 13 * 1 = 13 GB
  4. int4 13 * 0.5 = 6.5 GB

 

二、训练:显存

为了让模型收敛,训练的参数类型不能是 int8 / int4;一般是 float,效果差点的 就用 BF16;

 因为 反向传播 / Adam-优化 / Transformer架构 等因素,一般来说,训练需要的显存,是 同样规模 LLM 推理 的 3-4倍;

 

保守估计,按 4倍 计算。

 例如:7B 训练需要的显存,估算结果:

 参数类型 所需显存

float 7 * 4 * 4 = 112 GB

half / BF16 7 * 2 * 4 = 56 GB

 

梯度:参数量的1倍

优化器的状态:参数量的2倍

用 AdamW 优化器,需要2倍 参数量;

SGD 优化器,需要 1倍 参数量;

关于 LoRA / QLoRA 显存,以 LoRA 为例: 

LoRA 只需要给原始模型做推理,训练一个更小的模型来实现和训练原始参数差不多的效果。

例:原本需要微调1024×512的参数,用了LoRA之后,如果选择的 Rank=8,则只需要微调这么多参数:1024×8+512×8

 在这个基础上跑一次原始参数量的推理(不需要梯度和优化器状态,但仍然需要一些显存来存储运算过程的数据),合起来就是LoRA的所需要的显存。