大模型显存计算指南 - 推理与训练显存计算详解

📚 在部署和训练大模型时,显存往往是最大的瓶颈。本文将帮助你理解显存计算的关键要素,让你能够更好地规划硬件资源。

一、推理显存计算 - 如何评估模型部署所需显存?

1. 影响推理显存的关键因素

(1) 模型参数

• 基础显存占用 = 参数量 * 精度字节数
• 例如:7B模型在不同精度下的参数占用
  • FP32(4字节): 7B * 4 = 28GB
  • FP16(2字节): 7B * 2 = 14GB
  • INT8(1字节): 7B * 1 = 7GB
  • INT4(0.5字节): 7B * 0.5 = 3.5GB

(2) 注意力缓存(Attention Cache)

• KV Cache大小 = batch_size * num_layers * 2 * seq_length * hidden_size * precision
• 对于长文本生成,注意力缓存可能占用大量显存

(3) 激活值(Activations)

• 模型推理过程中的中间计算结果
• 通常占用基础参数量的10%-20%显存

2. 实际显存估算示例

以Qwen-7B为例(hidden_size=4096):

def calculate_inference_memory(
    batch_size=1,
    seq_length=2048,
    model_size_b=7,
    precision="fp16"
):
    # 精度映射
    precision_map = {
        "fp32": 4,
        "fp16": 2,
        "int8": 1,
        "int4": 0.5
    }
    
    # 基础参数显存
    param_memory = model_size_b * precision_map[precision]
    
    # KV缓存显存(假设32层)
    kv_cache = (batch_size * 32 * 2 * seq_length * 4096 * 
                precision_map[precision]) / (1024 ** 3)
    
    # 激活值显存(假设15%)
    activation_memory = param_memory * 0.15
    
    total_memory = param_memory + kv_cache + activation_memory
    return total_memory

💡 小贴士: 实际运行时的显存占用可能会比理论计算略高,建议预留20%左右的显存余量。

二、训练（微调）显存计算 - 为什么训练需要更多显存?

1. 训练阶段的额外显存开销

(1) 梯度(Gradients)

• 每个参数都需要存储对应的梯度
• 梯度显存 ≈ 模型参数显存

(2) 优化器状态(Optimizer States)

不同优化器的显存占用:

• SGD: 基础参数量 * 1
• Adam: 基础参数量 * 2 (需存储一阶矩和二阶矩)
• AdamW: 基础参数量 * 2
• Lion: 基础参数量 * 1
• Adafactor: 基础参数量 * 1.5

(3) 前向计算缓存

• 用于反向传播计算梯度
• 通常与序列长度成正比

2. 训练显存估算示例

以7B模型为例,使用AdamW优化器:

def calculate_training_memory(
    batch_size=1,
    seq_length=2048,
    model_size_b=7,
    precision="fp16",
    optimizer="adamw"
):
    # 基础推理显存
    inference_memory = calculate_inference_memory(
        batch_size, seq_length, model_size_b, precision
    )
    
    # 优化器状态映射
    optimizer_multiplier = {
        "sgd": 1,
        "adam": 2,
        "adamw": 2,
        "lion": 1,
        "adafactor": 1.5
    }
    
    # 梯度显存
    gradient_memory = model_size_b * 2  # fp16/bf16
    
    # 优化器状态显存
    optimizer_memory = (model_size_b * 
                       optimizer_multiplier[optimizer] * 
                       2)  # fp16/bf16
    
    total_memory = (inference_memory + 
                   gradient_memory + 
                   optimizer_memory)
    return total_memory

3. 显存优化技术

(1) 梯度检查点(Gradient Checkpointing)

• 以计算时间换显存空间
• 可节省30%-50%显存
• 训练速度降低20%-30%

(2) 混合精度训练

• 使用FP16/BF16进行计算
• 使用FP32存储主要权重
• 可节省40%-50%显存

(3) 零冗余优化器(ZeRO)

• 分布式训练中的显存优化
• 可实现接近线性的显存缩放
• 常用实现:DeepSpeed、Megatron-LM

🔍 实践建议:

1. 优先考虑使用混合精度训练
2. 如显存仍然不足,开启梯度检查点
3. 多卡训练时,使用ZeRO等分布式优化技术

参考

#监听显卡，每 1 秒刷新一次：
watch -n -1 -d nvidia-smi

下面给大家粗略估计推理和训练的显存需求情况，相信大家通过上面的学习，已经初步掌握了具体的计算方式，也可以按照下面的表格进行对照推理。

模型精度	7B	13B	70B
FP32全精度	28GB	52GB	280GB
FP16半精度	14GB	26GB	140GB
Int8 精度	7GB	13GB	70GB
Int4 精度	3.5GB	6.5GB	35GB

训练方法	模型精度	7B	13B	70B
全参数	FP32全精度	120GB	240GB	1200GB
全参数	FP16半精度	60GB	120GB	600GB
部分参数	FP16半精度	20GB	40GB	200GB
LoRA	FP16半精度	16GB	32GB	160GB
QLoRA	Int8 精度	10GB	20GB	80GB
QLoRA	Int4 精度	6GB	12GB	48GB

总结

合理估算和优化显存使用是成功部署和训练大模型的关键:

1. 推理阶段主要考虑参数量、注意力缓存和激活值
2. 训练阶段需额外考虑梯度、优化器状态和前向计算缓存
3. 合理使用显存优化技术可以突破硬件限制

希望这篇文章能帮助你更好地理解和规划大模型的显存使用！如果觉得有帮助,请点赞支持~ 😊