微调大模型需要多少GPU显存?

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. 总结与建议:

  1. 起点估算:

    • 全参数微调 (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 + 激活。

  2. 关键变量batch_size 和 sequence_length 对激活值影响巨大。如果显存不足,优先减小这两个值,或者加强梯度检查点的使用。

  3. 梯度检查点: 对于大模型微调(无论是全参数还是 LoRA),几乎是必需的技术,用计算换显存。

  4. 优化器: 如果显存极其紧张,可以考虑显存优化型的优化器(如 Adafactor, 8-bit Adam),但这可能会影响收敛效果。

  5. 分布式训练 (DeepSpeed ZeRO): 对于全参数微调或者超大模型的 LoRA 微调,单卡显存往往不够。DeepSpeed ZeRO (特别是 Stage 2 和 3) 可以将优化器状态和梯度分片到多张 GPU 上,极大降低单卡显存压力。

  6. 实际监控: 最好的方法是在目标硬件上用小 batch_size 跑一个测试批次,并使用 nvidia-smi 或 PyTorch 的 torch.cuda.memory_summary() / torch.cuda.max_memory_allocated() 来监控实际峰值显存占用,然后根据需要调整参数。

记住,这些都是估算值,实际需求会因具体的模型架构、代码实现、库版本等因素有所浮动。

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值