大模型训练\部署显存占用估算

1.经验评估

• 推理显存估算：7B-float 是 28 GB，7B-BF16 是 14GB，7B-int8 是 7GB；其他版本以此类推即可。
• 训练的参数类型，只能是 float / BF16
• 训练 所需显存 保守估算 是 同参数同类型llm 推理 的 4倍。
  • 例子：7B-float 训练 显存：28 * 4 = 112 GB

方法	bits	7B	13B	30B	65B	8*7B
全参数微调	16	160GB	320GB	600GB	1200GB	900GB
Freeze	16	20GB	40GB	120GB	240GB	200GB
LoRA	16	16GB	32GB	80GB	160GB	120GB
QLoRA	8	10GB	16GB	40GB	80GB	80GB
QLoRA	4	6GB	12GB	24GB	48GB	32GB

2.精确评估

2.1 在线评估

accelerate estimate-memory 是 huggingface 的 accelerate 开发库中提供的一个工具。可网页在线访问

https://huggingface.co/spaces/hf-accelerate/model-memory-usage选择相应模型进行评估

2.2 本地评估

• 安装 accelerate, transformers

pip install accelerate
pip install transformers

• 使用方法举例

# 基本使用方法
accelerate estimate-memory mistralai/Mistral-7B-v0.1

# 只显示指定的数据类型
accelerate estimate-memory mistralai/Mistral-7B-v0.1 --dtypes float16

# 指定开发库(针对本地模型，Hub上存储的模型不需要指定)
accelerate estimate-memory mistralai/Mistral-7B-v0.1 --dtypes float32 float16 --library_name transformers

# 设置 trust_remote_code=True
accelerate estimate-memory Qwen/Qwen1.5-7B #正常
accelerate estimate-memory Qwen/Qwen-7B #报错
accelerate estimate-memory Qwen/Qwen-7B --trust_remote_code #可以运行

# 其他模型
accelerate estimate-memory google/gemma-7b
accelerate estimate-memory baichuan-inc/Baichuan2-7B-Base --trust_remote_code

3.训练显存估算

对于 LLaMA 7B 模型 (70亿参数) 使用混合精度 (FP16/FP32) 和 AdamW 优化器，这些组件的计算是：

1. 模型参数 (FP16): 用于前向和反向传播计算。
   • 7B 参数 * 2 bytes/参数 = 14 GB
2. 模型参数 (FP32 Copy): 由混合精度训练维护，用于更精确的参数更新。
   • 7B 参数 * 4 bytes/参数 = 28 GB
3. 梯度 (Gradients): 在反向传播过程中计算。虽然计算可能在 FP16 进行，但在传递给 FP32 优化器更新前，梯度通常会被转换为 FP32（或者累加到 FP32 缓冲区）。因此，需要为 FP32 梯度预留空间以确保更新精度。
   • 7B 参数 * 4 bytes/参数 = 28 GB
4. 优化器状态 (AdamW): AdamW 需要为每个参数维护两个状态：一阶矩 (m) 和二阶矩 (v)。这些状态通常以 FP32 存储。
   • m: 7B 参数 * 4 bytes/参数 = 28 GB
   • v: 7B 参数 * 4 bytes/参数 = 28 GB
   • AdamW 状态总计: 28 GB + 28 GB = 56 GB

总计显存占用 (不含激活值):

将以上各项相加： 14 GB (FP16 参数) + 28 GB (FP32 参数 Copy) + 28 GB (FP32 梯度) + 56 GB (AdamW 状态 m+v) = 126 GB

重要补充:

• 激活值 (Activations): 激活值的占用取决于模型结构、batch size 和序列长度。对于 LLM 来说，激活值占用的显存可以非常大，有时甚至超过参数和优化器状态的总和，特别是在使用大 batch size 或长序列时。你的计算中没有包含这一部分是明确说明的，但在实际训练中，它是一个必须考虑的重要因素。
• 优化技术: 这个计算是基于标准的混合精度和 AdamW 优化器，没有使用显存优化技术，如：
  • Gradient Checkpointing: 显著减少激活值占用，但会增加计算时间。
  • ZeRO (Zero Redundancy Optimizer): 可以将参数、梯度、优化器状态分片存储在不同的 GPU 上，从而允许在单卡显存不足的情况下训练大模型。ZeRO Stage 1 分片优化器状态，ZeRO Stage 2 分片优化器状态和梯度，ZeRO Stage 3 分片所有状态包括参数。
  • 8-bit/4-bit 优化器状态: 有些优化器（如 bitsandbytes 的 AdamW）可以将优化器状态甚至梯度存储在更低精度，大幅减少显存占用。