大模型训练的三大显存优化策略
单卡场景
混合精度训练
核心思想：降低模型权重的数值精度，从 FP32 → FP16/BF16 → INT8 → INT4，大幅减少显存占用
FP16 / BF16 半精度推理
显存减半（FP32 7B 模型约 28GB → FP16 约 14GB）
使用
torch_dtype设置：
梯度检查点-以时间换空间
核心思想：不保存所有中间激活值，反向传播时重新计算，大幅降低训练时的显存占用
适用于微调（SFT、DPO）场景
显存节省 50%~80%，但训练速度降低 20%~30%
model.gradient_checkpointing_enable()或在from_pretrained 中设置：
model = AutoModelForCausalLM.from_pretrained(...,gradient_checkpointing=True
)多卡场景
模型并行
核心思想：将一个大模型的层或参数切分到多个 GPU 上，每张卡只存储一部分模型。
1. 张量并行（Tensor Parallelism）
将单个层的权重矩阵（如注意力头、FFN）按维度拆分
例如：将 7B 模型的
q_proj矩阵拆到 2 张卡上需要频繁的 GPU 间通信（All-Reduce）
适用场景：模型参数量 < 单卡显存容量
✅ 工具：Megatron-LM
2. 流水线并行（Pipeline Parallelism）
将模型按层拆分，每张卡负责一部分层
数据像“流水线”一样依次通过各卡
显著降低单卡显存占用
适用场景：模型参数量 > 单卡显存容量
✅ 工具：DeepSpeed、Accelerate
ZeRO(Zero Redundancy Optimizer)— DeepSpeed 的核心优化
核心思想：消除数据并行中冗余的优化器状态、梯度和参数副本
ZeRO 技术分为三个阶段：
阶段 | 优化对象 | 显存节省 |
ZeRo-1 | 将优化器状态分片到各GPU | 3~5x |
ZeRo-2 | 梯度 | 8~10x |
ZeRo-3 | 模型参数 | 15x+ |
示例：
// deepspeed_config.json
{"fp16": {"enabled": true},"zero_optimization": {"stage": 3,"offload_optimizer": {"device": "cpu"}},"train_micro_batch_size_per_gpu": 1
}
# 启动训练
deepspeed --num_gpus=4 train.py \--deepspeed deepspeed_config.json前沿优化方案
内存高效Attention
FlashAttention-2
通过 IO 感知算法减少显存读写
速度提升 2-4 倍
参数高效微调（PEFT）
LoRA
冻结原参数，训练低秩适配矩阵
显存占用降低70%
⚠️ 避坑指南
梯度检查点会导致约30%训练速度下降
FP16训练可能出现梯度下溢，需配合Loss Scaling
模型并行需要改写网络结构，调试复杂
📊 业界案例
LLaMA-2 70B：采用8路张量并行+16路流水线并行
GPT-3：梯度检查点+FP16节省显存78%