72B大模型微调显存优化实战：从理论到工程实践

1次阅读

没有评论

共计 1855 个字符，预计需要花费 5 分钟才能阅读完成。

在微调 72B 参数大模型时，显存消耗主要来自三个方面：模型参数、梯度和优化器状态。我们需要分别计算它们的显存占用。

模型参数：假设使用 FP32 精度，每个参数占用 4 字节，72B 参数需要 72 * 10^9 * 4 / 1024^3 ≈ 268GB 显存
梯度：与模型参数相同大小，也需要约 268GB 显存
优化器状态：以 Adam 优化器为例，需要存储动量和方差，每个参数额外占用 8 字节，共 72 * 10^9 * 8 / 1024^3 ≈ 536GB

总计理论最小需求：268 + 268 + 536 ≈ 1072GB 显存，这远超当前单卡 GPU 的容量。

梯度检查点通过在前向传播时只保存部分激活值，反向传播时重新计算中间结果，显著减少显存占用。

实现原理：
– 前向传播时不保存所有中间结果
– 反向传播时按需重新计算部分前向结果
– 显存占用从 O(n) 降低到 O(√n)

混合精度训练结合 FP16 和 FP32 计算，减少显存占用并保持数值稳定性。

配置方法：

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

with autocast():
    outputs = model(inputs)
    loss = criterion(outputs, targets)

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

技术	可训练参数量	显存节省	性能保持
Full FT	100%	0%	100%
LoRA	0.1-1%	90-99%	95-99%
Adapter	1-5%	80-95%	90-98%
Prefix-tuning	0.5-2%	85-98%	92-98%

import torch
from transformers import AutoModelForCausalLM
from torch.utils.checkpoint import checkpoint

# 初始化模型
model = AutoModelForCausalLM.from_pretrained("bigscience/bloom-72b")
model.gradient_checkpointing_enable()

# 自定义前向函数以支持梯度检查点
def custom_forward(*inputs):
    outputs = model(*inputs, use_cache=False)
    return outputs

# 训练循环
for batch in dataloader:
    inputs, labels = batch

    # 混合精度训练
    with torch.cuda.amp.autocast():
        # 使用梯度检查点
        outputs = checkpoint(custom_forward, inputs)
        loss = criterion(outputs.logits, labels)

    # 反向传播
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()
    optimizer.zero_grad()

在 8×A100 40GB 上测试不同优化策略的效果：