CPU Offloading在PyTorch中的实践：当GPU显存不够时把优化器状态放回内存

CPU Offloading在PyTorch中的实践当GPU显存不够时把优化器状态放回内存一、GPU显存的三座大山——参数、梯度与优化器状态在混合精度训练AMP中GPU显存的占用可以分解为三个主要部分模型参数FP16每个参数2字节。对7B模型约为14GB。梯度FP16同样每个参数2字节约14GB。优化器状态FP32Adam优化器维护两份动量m和v均为FP32。每个参数448字节约56GB。三者合计约84GB远超单张A100-80GB的容量。即使模型参数和梯度可以使用FP16降低一半优化器状态仍然是FP32存储。优化器状态占训练显存的60-70%——这是CPU Offloading策略之所以有效的核心数据。CPU Offloading的基本思想将优化器状态Adam的m和v和FP32的参数副本存放在CPU内存中仅在优化器步骤optimizer.step时将梯度传输到CPU完成更新再将更新后的FP32参数传回GPU。graph TB subgraph GPU显存 A[FP16 模型参数br/14GB] B[FP16 梯度br/14GB] C[FP16 激活值br/可变] end subgraph CPU内存 D[FP32 优化器状态 mbr/28GB] E[FP32 优化器状态 vbr/28GB] F[FP32 参数副本br/28GB] end A -.-|前向使用| G[计算] B -.-|反向产生| G D --|step时传输| A E --|step时传输| A F --|更新后回传| A二、Offloading的两种粒度与同步策略PyTorch FSDP和DeepSpeed都支持CPU Offloading但在实现粒度上有差异。FSDP的CPU Offloading基于FSDP单元通常是一个Transformer层进行offloading。在每个FSDP单元的反向传播完成后立即将优化器状态offload到CPU。这种细粒度offloading的优点是GPU和CPU之间的D2H传输可以与下一个FSDP单元的前向计算overlap。DeepSpeed ZeRO-Offload采用更粗的粒度——在整个反向传播完成后统一将优化器状态offload到CPU。优势是实现更简单劣势是D2H传输无法与前向计算overlap。关键的性能瓶颈是PCIe带宽。A100的PCIe 4.0 x16带宽约32GB/s而GPU内部的HBM带宽为2TB/s——差距约60倍。优化器状态在CPU和GPU之间的每次传输都会经过这个瓶颈。# FSDP CPU Offloading的配置 # 设计思路通过FSDP的cpu_offload参数将优化器状态迁移到CPU import torch import torch.nn as nn from torch.distributed.fsdp import ( FullyShardedDataParallel as FSDP, CPUOffload, ShardingStrategy, MixedPrecision, BackwardPrefetch, ) from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy def configure_fsdp_with_cpu_offload( model: nn.Module, offload_params: bool False, # 是否offload参数更激进 offload_optimizer: bool True, # 是否offload优化器状态 ) - FSDP: 配置FSDP的CPU Offloading 参数: offload_params: 将模型参数也offload到CPU显存最小但最慢 offload_optimizer: 将优化器状态offload到CPU推荐配置 设计决策 - 默认仅offload优化器状态。参数offload会导致每次前向传播 都需要D2H传输延迟增加显著。 - 配合FULL_SHARD策略使用因为分片后的参数已经降低了显存 再叠加CPU offload可以达到最大的显存节省效果。 # CPU Offload配置 cpu_offload CPUOffload( offload_paramsoffload_params ) # 注FSDP的CPUOffload类同时控制参数和优化器的offload # 混合精度进一步降低参数和梯度的显存 mixed_precision MixedPrecision( param_dtypetorch.bfloat16, reduce_dtypetorch.float32, # 梯度聚合仍用FP32 buffer_dtypetorch.bfloat16, keep_low_precision_gradsFalse, # 保持FP32梯度用于优化器 ) # Transformer层的自动包装策略 # 每个Transformer层作为一个FSDP单元实现细粒度offloading auto_wrap_policy transformer_auto_wrap_policy model FSDP( model, sharding_strategyShardingStrategy.FULL_SHARD, cpu_offloadcpu_offload, mixed_precisionmixed_precision, auto_wrap_policyauto_wrap_policy, backward_prefetchBackwardPrefetch.BACKWARD_PRE, # 前向prefetch在当前层计算时预取下一层参数 # 对CPU offload场景尤为重要——预取可以隐藏D2H延迟 forward_prefetchTrue, # 限制同时all-gather的FSDP单元数防止显存峰值 limit_all_gathersTrue, ) return model三、CPU Offloading的性能代价测量在LLaMA-7B的训练中实测不同offloading配置的性能影响配置显存占用/GPU训练吞吐相对DDPDDP (无offload)OOM--FSDP FULL_SHARD38GB3200 t/s基准FSDP 优化器offload22GB2100 t/s-34%FSDP 参数优化器offload14GB980 t/s-69%34%的吞吐下降主要来源于每次optimizer.step时的D2H传输。一个优化设计是使用pinned memory——将CPU侧的优化器状态分配在页锁定内存中加速D2H传输。实测表明pinned memory可以回收约8-12%的吞吐损失。# Pinned Memory优化器状态 # 设计思路使用CUDA的pinned memory加速CPU-GPU传输 import torch class PinnedOptimizerState: 使用pinned memory存储优化器状态 Pinned memory页锁定内存不会被操作系统换出到磁盘 CUDA可以使用DMA直接访问传输速度约为普通内存的2-3倍。 代价是pinned memory的分配和释放比普通内存慢 且占用的是物理内存不能swap因此仅应用于优化器状态 这种固定大小的长期分配。 def __init__(self, param_tensor: torch.Tensor): # 在CPU上分配pinned内存 # pin_memory()将已分配的张量移动到页锁定内存 self.m torch.zeros_like( param_tensor, devicecpu ).pin_memory() # ← 关键pinned memory self.v torch.zeros_like( param_tensor, devicecpu ).pin_memory() def transfer_to_gpu(self): 传输优化器状态到GPU 使用non_blockingTrue允许传输与计算overlap。 前提是后续计算不立即依赖这些数据。 self.m_gpu self.m.to(cuda, non_blockingTrue) self.v_gpu self.v.to(cuda, non_blockingTrue) def transfer_to_cpu(self): 更新后的状态回传CPU self.m.copy_(self.m_gpu, non_blockingTrue) self.v.copy_(self.v_gpu, non_blockingTrue)四、CPU Offloading的适用边界CPU Offloading的收益随模型规模和硬件配置而变化- 显存刚好不够时缺10-20%CPU Offloading是最优解。性能损失可控15-25%避免了购买更多GPU的硬件成本。- 显存严重不足时缺50%仅靠CPU Offloading可能不够需要叠加gradient checkpointing。但两者的叠加会使吞吐下降50%以上此时应考虑模型切分到更多GPU。- CPU内存不足时如果CPU内存也不足以存放优化器状态例如8节点训练100B模型需要使用NVMe OffloadingDeepSpeed ZeRO-Infinity将优化器状态卸载到NVMe SSD。graph TD A[显存不足程度] -- B{显存缺口} B --|20%| C[FSDP FULL_SHARDbr/无CPU Offload] B --|20-40%| D[FSDP 优化器Offloadbr/性能损失15-25%] B --|40-60%| E[FSDP 优化器参数Offloadbr/Gradient Checkpointing] B --|60%| F[增加GPU数br/或DeepSpeed ZeRO-Infinity]五、总结CPU Offloading通过将占显存60%以上的优化器状态和FP32参数副本迁移到CPU内存以15-35%的吞吐损失换取30-50%的显存节省。这是一个典型的时间换空间策略。关键优化技巧使用pinned memory加速传输、启用forward_prefetch隐藏传输延迟、优先offload优化器而非参数前者的性能代价远小于后者。对于单卡80GB训练7B-13B级别模型的场景FSDP优化器offload是目前最具成本效益的方案。