当你发现训练一个70亿参数模型需要8张A100时，不是显卡不够，而是你的显存优化还不到位。我是maoku，今天带你掌握大模型训练的“显存瘦身术”。
2024年，大型语言模型的参数规模已突破万亿级别，但与此同时，训练这些模型所需的显存资源却成为普通研究者和中小企业难以逾越的门槛。一张80G显存的A100显卡市场价超过10万元，训练一个千亿参数模型通常需要数十张这样的显卡——成本高达数百万甚至上千万元。
但真相是，通过精密的显存优化技术，完全可以用远少于常规配置的显卡完成相同模型的训练。本文将为你揭示大模型训练中显存消耗的秘密，并提供一套从理论到实践的完整优化方案。
第一章：显存去哪了？——大模型训练的“内存账簿”
直观比喻：把GPU显存想象成你的工作台
想象你要组装一台精密仪器，你的工作台空间有限（GPU显存），上面需要放置：
- 设计图纸和零件（模型权重）：必须始终放在工作台上
- 组装说明书（优化器状态）：告诉你如何调整零件位置
- 临时拆下的部件（激活值/计算图）：组装过程中需要暂存的部分
- 工具和包装材料（临时变量）：用完就可以扔掉
下面这张图清晰地展示了训练过程中显存的主要消耗构成：
pie title 大模型训练显存消耗分布（以BF16混合精度为例）
“模型权重 (Weights)” : 1
“梯度 (Gradients)” : 1
“优化器状态 (Optimizer States)” : 6
“激活值/计算图 (Activations)” : 3
“临时变量 (Temporary)” : 1
详细分解：每个部分的具体构成
1. 模型权重（Weight） - 占比约12.5%
- 这是模型的核心知识，以浮点数矩阵形式存储
- 一个70亿参数模型，使用BF16精度需要约14GB显存（7B × 2字节）
2. 梯度（Gradient） - 占比约12.5%
- 反向传播时计算得到，指示每个参数应该调整的方向和幅度
- 与模型权重完全相同的尺寸
3. 优化器状态（Optimizer States） - 占比约75%
- 这是显存的最大消耗者，通常包括：
- 一阶动量（m）：如Adam优化器中的梯度指数移动平均
- 二阶动量（v）：梯度平方的指数移动平均
- 主权重（Master Weight）：混合精度训练中保持的高精度权重副本
- 对于Adam优化器，每个参数需要额外存储8字节（BF16训练时）
4. 激活值/计算图（Activations） - 变量最大
- 前向传播中每一层的输出结果，反向传播时需要用于梯度计算
- 消耗与批次大小、序列长度、模型层数成正比
- 通常是除优化器状态外的第二大显存消耗
5. 临时变量（Temporary Variables）
- 计算过程中的中间结果，生命周期短
- 良好的编程实践可以显著减少这部分消耗
关键计算公式
总显存消耗 ≈
模型权重 +
梯度 +
优化器状态 +
激活值 +
临时变量
简化估算（Adam优化器，BF16精度）：
总显存 ≈ 模型权重 × (1 + 1 + 6) + 激活值
≈ 模型权重 × 8 + 激活值
第二章：为什么必须优化显存？——不只是省钱那么简单
经济性：让更多人用得起AI训练
以主流的70亿参数模型为例，对比优化前后的显存需求：
| 配置方案 | 显卡需求（80G A100） | 硬件成本 | 适用团队 |
|---|---|---|---|
| 原始方案 | 8-16张 | 80-160万元 | 大型企业 |
| 优化后方案 | 2-4张 | 20-40万元 | 中小企业/研究机构 |
| 极致优化 | 1张（部分场景） | 10万元 | 个人研究者 |
成本降低75%以上，这使得AI训练从“土豪游戏”变成了更多人可以参与的技术探索。
技术灵活性：获得更多策略选择空间
显存充足时，你可以自由选择更适合任务的技术策略：
- 减少张量并行（TP）规模：降低通信开销，提升训练速度
- 减小流水线并行（PP）阶段：减少训练“气泡时间”，提高效率
- 避免完全复制（CP）策略：节省显存，简化实现复杂度
- 增大批次大小：更充分利用计算核心，提升硬件利用率
效率提升：更高的MFU（模型浮点运算利用率）
MFU是衡量训练效率的关键指标，优化显存可以直接提升MFU：
MFU = 实际浮点运算量 / (显卡峰值算力 × 训练时间)
通过优化显存，你可以：
- 减少通信等待时间
- 提高计算核心利用率
- 减少闲置资源
- 最终实现更快的训练速度和更低的成本
第三章：实战指南——九大显存优化技巧
技巧1：算子融合——减少不必要的中间变量
问题：某些连续操作会产生大量临时张量，占用显存后立即释放
解决方案：将多个操作融合为单个内核函数
# 优化前：产生两个大型临时张量
# 假设vocab_size=100,000，hidden_size=4096
logits = torch.matmul(hidden_states, lm_head_weight.T) # [batch, seq, vocab_size]
loss = F.cross_entropy(logits.view(-1, vocab_size), labels.view(-1))
# 优化后：自定义融合算子
# 使用Triton或其他内核融合技术
class FusedLMHeadWithLoss(torch.autograd.Function):
@staticmethod
def forward(ctx, hidden_states, weight, labels):
# 直接在核函数中计算logits和loss
# 避免存储完整的logits矩阵
loss = fused_lm_head_loss_forward(hidden_states, weight, labels)
ctx.save_for_backward(hidden_states, weight, labels)
return loss
@staticmethod
def backward(ctx, grad_loss):
# 反向传播时重新计算必要中间结果
hidden_states, weight, labels = ctx.saved_tensors
grad_hidden, grad_weight = fused_lm_head_loss_backward(
grad_loss, hidden_states, weight, labels
)
return grad_hidden, grad_weight, None
# 使用融合算子
loss = FusedLMHeadWithLoss.apply(hidden_states, lm_head_weight, labels)
适用场景：LM Head + CrossEntropy Loss、RMSNorm等连续操作
技巧2：避免不必要的张量拷贝
常见陷阱：PyTorch中的某些操作会隐式创建张量副本
# 不良实践：产生拷贝
x = torch.randn(1024, 4096, device="cuda")
# 非连续张量的reshape会产生拷贝
y = x.t().reshape(-1, 1024) # 隐式拷贝！
# 良好实践：避免拷贝
# 方法1：使用view代替reshape（当可能时）
y = x.t().view(-1, 1024) # 要求原始张量是连续的
# 方法2：使用in-place操作
x.div_(2.0) # 原地除法，不创建新张量
x.add_(y) # 原地加法
# 方法3：使用permute代替transpose（某些情况下）
# permute不会改变数据布局，只是改变视图
诊断工具：使用PyTorch的内存分析器
import torch
torch.cuda.memory._record_memory_history()
# ...运行你的代码...
torch.cuda.memory._dump_snapshot("snapshot.pickle")
技巧3：混合精度训练——BF16/FP8的威力
原理：使用低精度格式存储和计算，显著减少显存占用
| 精度格式 | 字节/参数 | 显存节省 | 适用场景 |
|---|---|---|---|
| FP32 | 4字节 | 基准 | 传统训练 |
| BF16 | 2字节 | 50% | 现代大模型训练 |
| FP8 | 1字节 | 75% | 最新硬件支持 |
PyTorch自动混合精度示例：
from torch.cuda.amp import autocast, GradScaler
scaler = GradScaler() # 梯度缩放，防止下溢
for inputs, labels in dataloader:
optimizer.zero_grad()
# 前向传播使用自动混合精度
with autocast(dtype=torch.bfloat16):
outputs = model(inputs)
loss = criterion(outputs, labels)
# 反向传播，scaler自动处理精度转换
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
技巧4：梯度检查点——时间换空间
原理：不保存所有中间激活值，反向传播时重新计算部分前向传播
# 常规训练：存储所有层激活
# 显存消耗 ∝ 层数 × 激活大小
# 梯度检查点：只存储关键点激活
from torch.utils.checkpoint import checkpoint_sequential
# 将模型分成若干段
segments = 4 # 将模型分成4段
def forward_with_checkpoint(x):
return checkpoint_sequential(model.layers, segments, x)
# 或者手动设置检查点
def custom_forward(layer_module, hidden_states):
def create_custom_forward(module):
def custom_forward(*inputs):
return module(*inputs)
return custom_forward
return torch.utils.checkpoint.checkpoint(
create_custom_forward(layer_module),
hidden_states,
use_reentrant=False # 新版本推荐
)
# 使用策略：对计算密集度低的层使用检查点
for layer in model.layers:
if layer.is_compute_intensive:
hidden_states = layer(hidden_states) # 常规前向
else:
hidden_states = custom_forward(layer, hidden_states) # 检查点
经验法则：
- 激活检查点可将激活显存从 O(层数) 降低到 O(√层数)
- 对注意力机制层使用检查点效果显著
- 重新计算带来的时间开销通常为20-30%
技巧5：优化器状态分片与卸载
ZeRO优化器原理：将优化器状态、梯度、参数分片到不同设备
# DeepSpeed ZeRO配置示例（deepspeed_config.json）
{
"train_batch_size": 32,
"train_micro_batch_size_per_gpu": 4,
"zero_optimization": {
"stage": 3, # ZeRO阶段：1-优化器状态分片，2-+梯度分片，3-+参数分片
"offload_optimizer": {
"device": "cpu", # 优化器状态卸载到CPU
"pin_memory": true
},
"offload_param": {
"device": "cpu", # 参数卸载到CPU
"pin_memory": true
},
"overlap_comm": true, # 重叠通信和计算
"contiguous_gradients": true, # 连续梯度缓冲区
"stage3_max_live_parameters": 1e9,
"stage3_max_reuse_distance": 1e9,
"stage3_prefetch_bucket_size": 5e8
},
"fp16": {
"enabled": true,
"loss_scale": 0,
"loss_scale_window": 1000,
"initial_scale_power": 16
}
}
训练脚本调整：
import deepspeed
# 初始化模型
model_engine, optimizer, _, _ = deepspeed.initialize(
args=args,
model=model,
model_parameters=model.parameters(),
config="deepspeed_config.json"
)
# 训练循环
for batch in dataloader:
loss = model_engine(batch)
model_engine.backward(loss)
model_engine.step()
技巧6：解决显存碎片化问题
问题现象：已分配显存远小于预留显存，大量零碎空间无法利用
诊断方法：
import torch
# 检查显存碎片情况
print(f"已分配: {torch.cuda.memory_allocated()/1e9:.2f} GB")
print(f"已预留: {torch.cuda.memory_reserved()/1e9:.2f} GB")
print(f"碎片率: {(torch.cuda.memory_reserved()-torch.cuda.memory_allocated())/torch.cuda.memory_reserved()*100:.1f}%")
# 如果碎片率>30%，说明存在严重碎片问题
解决方案：
- 大张量分块处理
# 将大张量拆分为多个小张量处理
def process_large_tensor_in_chunks(tensor, chunk_size=1024):
results = []
for i in range(0, tensor.size(0), chunk_size):
chunk = tensor[i:i+chunk_size]
# 处理chunk
result_chunk = expensive_operation(chunk)
results.append(result_chunk)
# 及时释放不再需要的中间变量
del chunk
torch.cuda.empty_cache() # 谨慎使用，可能有性能开销
return torch.cat(results)
- 统一数据布局：确保张量在内存中是连续的
# 确保张量连续
if not tensor.is_contiguous():
tensor = tensor.contiguous() # 可能会产生拷贝，但减少碎片
- 使用可扩展段分配器（PyTorch 2.0+）
# 启动时设置环境变量
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:true python train.py
技巧7：流水线并行中的负载均衡
问题：流水线并行的不同阶段显存消耗不均
解决方案：手动分配层到不同设备
# 自定义层分配策略
def balanced_pipeline_assignment(total_layers, num_stages):
"""均衡分配层到各个流水线阶段"""
# 考虑不同层类型的显存消耗差异
# 通常：嵌入层和输出层消耗大，中间层相对均匀
layers_per_stage = total_layers // num_stages
remainder = total_layers % num_stages
assignment = []
start = 0
# 为前几个阶段多分配一层（如果有余数）
for stage in range(num_stages):
end = start + layers_per_stage + (1 if stage < remainder else 0)
assignment.append((start, end))
start = end
return assignment
# 示例：24层模型，4个流水线阶段
# 传统平均分配：[0-6], [6-12], [12-18], [18-24]
# 均衡分配（考虑首尾层较重）：[0-4], [4-10], [10-16], [16-24]
技巧8：激活值卸载（Activation Offload）
高级技术：将激活值临时卸载到CPU内存
# 概念示例（实际实现更复杂）
class ActivationOffload:
def __init__(self, layer, offload_device='cpu'):
self.layer = layer
self.offload_device = torch.device(offload_device)
def forward(self, x):
# 正常计算
output = self.layer(x)
# 将激活值移动到CPU（非阻塞传输）
if self.training:
# 使用pinned memory加速传输
offloaded_activation = output.to(
self.offload_device,
non_blocking=True,
memory_format=torch.pinned_memory
)
# 保存引用，反向传播时使用
self.saved_for_backward = offloaded_activation
return output
def backward(self, grad_output):
# 从CPU取回激活值
activation = self.saved_for_backward.to(
grad_output.device,
non_blocking=True
)
# 计算梯度
grad_input = self.layer.backward(grad_output, activation)
return grad_input
注意事项：
- 需要pinned memory支持
- CPU-GPU传输可能成为瓶颈
- 适用于激活值很大但计算量不大的情况
技巧9：自定义内存分配器
高级主题：重写PyTorch内存分配器（仅推荐高级用户）
// 简化的自定义分配器概念
class CustomAllocator {
public:
void* allocate(size_t size) {
// 实现更智能的分配策略
// 1. 合并小块内存
// 2. 预分配大块内存池
// 3. 智能缓存管理
}
void free(void* ptr) {
// 实现更有效的释放策略
// 及时合并相邻空闲块
}
};
// 在PyTorch中注册自定义分配器
torch::cuda::CUDACachingAllocator::set_allocator(
std::make_unique<CustomAllocator>()
);
对于不想深入底层实现的研究者，可以考虑使用【LLaMA-Factory Online】这类集成平台，它们通常内置了经过优化的显存管理策略，让用户能更专注于模型设计和实验。
第四章：效果评估与验证
监控指标：你需要关注的关键数据
- 显存使用效率
def calculate_memory_efficiency():
allocated = torch.cuda.memory_allocated()
reserved = torch.cuda.memory_reserved()
max_reserved = torch.cuda.max_memory_reserved()
efficiency = allocated / reserved if reserved > 0 else 0
peak_usage = max_reserved / torch.cuda.get_device_properties(0).total_memory
return {
'allocated_gb': allocated / 1e9,
'reserved_gb': reserved / 1e9,
'efficiency_percent': efficiency * 100,
'peak_usage_percent': peak_usage * 100,
'fragmentation_percent': (reserved - allocated) / reserved * 100 if reserved > 0 else 0
}
-
训练速度对比
- 优化前后的每秒处理样本数（samples/sec）
- 每个epoch的训练时间
- 达到相同准确率所需的总时间
-
模型质量验证
- 在验证集上的准确率/损失变化
- 避免因优化引入的数值不稳定
实用调试技巧
- 显存使用时间线分析
# 使用PyTorch Profiler
python -m torch.profiler \
--activities=cuda \
--schedule=repeat=1 \
--on_trace_ready=trace_handler \
train.py
-
渐进式优化策略
- 先确保模型能运行（即使很慢）
- 逐一应用优化技巧，每次验证效果
- 记录每次优化带来的显存节省和性能影响
-
常见问题诊断表
| 问题症状 | 可能原因 | 解决方案 |
|---|---|---|
| OOM但显存未满 | 严重碎片化 | 使用连续张量，避免大张量频繁分配释放 |
| 训练速度突然下降 | 频繁缓存清理 | 调整allocator配置，增大缓存大小 |
| 梯度爆炸/消失 | 混合精度配置不当 | 调整loss scaling，检查梯度裁剪 |
| 通信开销过大 | ZeRO阶段设置不当 | 调整stage，启用overlap_comm |
第五章：成功案例与最佳实践
实际案例：在8卡80G上训练720亿参数模型
挑战：720亿参数模型，常规训练需要32+张A100
优化策略组合：
- ZeRO Stage 3：参数、梯度、优化器状态全部分片
- 梯度检查点：对前40层使用激活检查点
- BF16混合精度：减少一半的权重存储
- 自定义流水线分配：首尾stage分配较少层数
- CPU Offload：将部分优化器状态卸载到CPU
最终配置：
- 显卡：8×A100 80GB
- 批次大小：每卡微批次2，全局批次128
- 优化器：AdamW（β1=0.9，β2=0.95）
- 学习率：1e-4，余弦退火
- 达到的目标：与32卡配置相比，MFU仅降低8%，成本降低75%
最佳实践总结
- 从简单开始：先让模型跑起来，再逐步优化
- 测量优先：任何优化前先建立基准性能指标
- 组合使用：单一技术效果有限，组合多个技术效果显著
- 平衡取舍：显存优化往往以时间为代价，找到适合的平衡点
- 保持可复现：记录每次优化的配置和结果
总结与展望
通过本文介绍的九大显存优化技术，我们看到了即使是资源有限的研究者也能训练大型语言模型的可能性。关键要点总结如下：
- 理解显存消耗分布是优化的基础，优化器状态通常是最大消耗者
- 混合精度训练是最简单有效的入门优化
- 梯度检查点在激活值显存优化中作用显著
- ZeRO优化器是分布式训练显存优化的核心技术
- 综合应用多种技术才能达到极致优化效果
未来趋势：
- 硬件协同设计：新一代GPU针对大模型训练优化显存架构
- 更智能的分配器：机器学习驱动的动态显存管理
- 算法创新：从源头减少显存需求的训练算法
- 标准化工具链：一键式显存优化配置
给初学者的行动建议：
- 从小的模型和数据集开始实验
- 先掌握混合精度和梯度检查点这两个基础技术
- 使用成熟的框架（如DeepSpeed、ColossalAI）而非从头实现
- 建立自己的优化工具箱，记录每种技术在不同场景下的效果
- 加入相关社区，跟进最新优化技术
显存优化不是一次性的工作，而是贯穿整个大模型训练生命周期的重要实践。通过持续优化，我们不仅能降低成本，还能更深入地理解模型训练的内在机制。
希望这篇指南能帮助你开启大模型训练的高效之旅。如果你在实践过程中遇到具体问题，或有自己的优化经验想要分享，欢迎在评论区交流讨论。我是maoku，我们下次再见！
