0. 引言
如何估算微调(Fine-tuning)一个X B(比如X=1,即十亿)参数的大模型所需显存(VRAM)?精确估计比较难,因为它受到多种因素的影响。这里我们分全参数微调和LoRA微调两种情况来讨论,并给出估算方法和一些经验法则。更多大模型相关欢迎关注微信公众号:"小窗幽记机器学习"。
1. 如何估计
核心影响因素:
1、 模型参数量 (X B) :模型的规模是基础。
2、 参数精度 (Precision) :
-
FP32 (单精度):每个参数占 4 字节。
-
FP16 (半精度) / BF16 (脑浮点16):每个参数占 2 字节。(训练常用)
-
INT8 (8位整数):每个参数占 1 字节。(常用于推理或特定量化训练技术如QLoRA)
3、 优化器状态 :像 Adam 或 AdamW 这样的优化器需要存储额外的状态信息(如动量和方差)。 这是显存消耗大户。
-
标准 Adam/AdamW 通常需要存储两倍于模型参数量的状态(动量和方差),且通常以 FP32 存储(即使模型用 FP16 训练,优化器状态也常是 FP32 以保证稳定性)。因此,每个参数需要额外的
2 * 4 = 8
字节。 -
一些优化器(如 Adafactor 或 8-bit Optimizers)会显著减少这部分显存。
4、 梯度 (Gradients) :反向传播计算出的梯度需要存储,其大小与模型参数量相同,精度通常与训练精度一致(如 FP16/BF16)。
5、 激活值 (Activations) :前向传播过程中产生的中间输出。其大小与 batch_size * sequence_length * hidden_dimension * num_layers
相关。 这是非常动态的部分,受批次大小和序列长度影响很大。可以使用梯度检查点 (Gradient Checkpointing / Activation Checkpointing) 技术大幅减少这部分显存,但会增加计算时间。
6、 Batch Size 和 Sequence Length:直接影响激活值的显存占用。
7、 框架和库的开销:如 PyTorch, TensorFlow, CUDA kernels 等会占用一部分固定或可变的显存。
2. 全参数微调 (Full Fine-tuning)
在全参数微调中,模型的所有参数都需要计算梯度并由优化器更新。
估算公式(以常用配置为例:FP16/BF16 混合精度训练 + AdamW 优化器):
-
模型参数显存:
X B * 2 bytes/param
=2X
GB -
梯度显存:
X B * 2 bytes/param
=2X
GB -
优化器状态显存 (AdamW, FP32 状态):
X B * 8 bytes/param
=8X
GB -
总计 (核心部分):
(2 + 2 + 8) * X
GB =12X
GB
考虑激活值和其他开销:
-
激活值显存: 这部分非常依赖
batch_size
和sequence_length
。对于大模型,如果不使用梯度检查点,这部分可能非常大(数十甚至上百 GB)。使用梯度检查点后,这部分可以显著减少,但仍会占用数 GB 到数十 GB 不等的显存,且难以简单地用X
来线性估算。
激活显存粗略估计 ≈ batch_size × seq_len × hidden_dim × n_layers × bytes_per_element × fhidden_dim:模型隐藏层维度n_layers:模型层数bytes_per_element:每个元素占用的显存(通常是 2 字节 FP16)f 系数:考虑每层是否保留多个激活(通常取 2~3)
-
其他开销 (框架、CUDA 等): 通常需要预留几 GB。
经验估算规则 (全参数微调, FP16/BF16, AdamW):
-
非常粗略的下限 (假设梯度检查点有效且 batch_size/seq_len 适中):
~14X
GB 到18X
GB。 -
更实际的考虑 (包含一定的激活值和开销): 可能需要
~20X
GB 或更多。例如,一个 7B 模型(X=7)全参数微调,通常需要7 * 20 = 140
GB 左右的显存,但这仍取决于具体配置。单个 A100/H100 (80GB) 通常是不够的,需要多卡并行(如使用 DeepSpeed ZeRO)。
示例:
-
微调一个 7B 模型 (X=7):
-
核心部分约
12 * 7 = 84
GB。 -
加上激活值(即使有检查点)和开销,可能轻松超过 100 GB。使用 2 * A100 (80GB) 并配合 DeepSpeed ZeRO Stage 2 或 3 是常见的配置。
-
-
微调一个 70B 模型 (X=70):
-
核心部分约
12 * 70 = 840
GB。 -
总需求会远超 1 TB,需要大规模的 GPU 集群。
-
3. LoRA 微调 (Low-Rank Adaptation)
LoRA 只训练一小部分注入到模型中的“适配器”参数,而原始模型的绝大部分参数保持冻结。
显存组成:
-
冻结的基础模型参数: 仍然需要加载到显存中进行前向传播。通常使用 FP16/BF16 加载。
-
显存 ≈
X B * 2 bytes/param
=2X
GB
-
-
LoRA 参数: 数量远小于
X B
,通常只有几百万到几千万(假设为Y M
)。-
LoRA 参数显存:
Y M * 2 bytes/param
(FP16/BF16) - 通常只有几十到几百 MB,相对基础模型可以忽略不计。
-
-
LoRA 参数的梯度:
Y M * 2 bytes/param
- 同上,很小。 -
LoRA 参数的优化器状态 (AdamW):
Y M * 8 bytes/param
- 仍然很小。 -
激活值: 这是关键! 即使只训练 LoRA 参数,前向传播仍然需要经过整个模型,因此激活值的显存占用与全参数微调类似,同样受
batch_size
和sequence_length
影响巨大。梯度检查点同样适用且非常推荐。 -
其他开销: 与全参数微调类似。
估算公式 (LoRA, FP16/BF16):
-
总显存 ≈ (基础模型显存) + (激活值显存) + (其他开销)
-
总显存 ≈
2X
GB + 激活值显存 + 几 GB 开销
经验估算规则 (LoRA, FP16/BF16):
-
LoRA 的主要显存节省来自于不需要存储庞大的梯度和优化器状态。
-
显存瓶颈通常是基础模型本身的大小和激活值。
-
粗略估算:
~2X
GB + (梯度检查点下的激活值显存) + 少量开销。 -
如果
batch_size
和sequence_length
控制得当,并且使用了梯度检查点,总显存可能在(2.5 ~ 4) * X
GB 的范围内,具体取决于激活值的大小。
示例:
-
微调一个 7B 模型 (X=7) 使用 LoRA:
-
基础模型约
2 * 7 = 14
GB。 -
激活值(假设使用梯度检查点,适中 batch/seq_len)可能需要 5-15 GB。
-
总计可能在
14 + (5~15) + few
GB ≈20 ~ 30+
GB。一张 24GB (如 RTX 3090/4090) 或 40GB/80GB (A100/H100) 的卡通常足够。
-
-
微调一个 70B 模型 (X=70) 使用 LoRA:
-
基础模型约
2 * 70 = 140
GB。 -
激活值和开销会增加几十 GB。
-
总计可能需要
140 + 20~40 + few
GB ≈160 ~ 190+
GB。需要多张高端 GPU(如 2-3 张 A100/H100 80GB)。
-
QLoRA (Quantized LoRA):
-
QLoRA 是一种更节省显存的技术,它将基础模型以 4-bit 加载 (NF4 类型)。
-
基础模型显存 ≈
X B * 0.5 bytes/param
(近似值,因为量化有额外开销) =~0.5X
GB。 -
这使得基础模型的显存占用大幅降低。例如,7B 模型基础部分只需要约 3.5-5 GB。
-
QLoRA 微调 7B 模型,可能在 10-16 GB 显存的 GPU 上就能运行(取决于 batch size 和序列长度)。70B 模型的 QLoRA 可能在单张 80GB GPU 上运行(但 batch size 和序列长度会受限)。
4. 总结与建议:
-
起点估算:
-
全参数微调 (FP16, AdamW): 考虑
~20X
GB 或更多。 -
LoRA 微调 (FP16): 考虑
~(2.5 - 4)X
GB,主要看基础模型2X
GB + 激活。 -
QLoRA 微调 (4-bit base, LoRA): 考虑
~(0.7 - 1.5)X
GB,主要看基础模型~0.5X
GB + 激活。
-
-
关键变量:
batch_size
和sequence_length
对激活值影响巨大。如果显存不足,优先减小这两个值,或者加强梯度检查点的使用。 -
梯度检查点: 对于大模型微调(无论是全参数还是 LoRA),几乎是必需的技术,用计算换显存。
-
优化器: 如果显存极其紧张,可以考虑显存优化型的优化器(如 Adafactor, 8-bit Adam),但这可能会影响收敛效果。
-
分布式训练 (DeepSpeed ZeRO): 对于全参数微调或者超大模型的 LoRA 微调,单卡显存往往不够。DeepSpeed ZeRO (特别是 Stage 2 和 3) 可以将优化器状态和梯度分片到多张 GPU 上,极大降低单卡显存压力。
-
实际监控: 最好的方法是在目标硬件上用小
batch_size
跑一个测试批次,并使用nvidia-smi
或 PyTorch 的torch.cuda.memory_summary()
/torch.cuda.max_memory_allocated()
来监控实际峰值显存占用,然后根据需要调整参数。
记住,这些都是估算值,实际需求会因具体的模型架构、代码实现、库版本等因素有所浮动。