CPU Offload 与异构训练
CPU Offload 将部分训练数据（优化器状态、参数、梯度）卸载到 CPU 内存，利用 CPU 大容量内存弥补 GPU 显存不足。异构训练进一步整合 GPU、CPU、NVMe 等多种存储层级。
Offload 策略
优化器状态 Offload
最轻量的 Offload 方案，将 Adam 的 m 和 v 存储在 CPU：
python
# DeepSpeed 优化器 Offload
ds_config = {
"zero_optimization": {
"stage": 2,
"offload_optimizer": {
"device": "cpu",
"pin_memory": True,
"buffer_count": 4,
"fast_init": False,
},
},
}参数 Offload
将模型参数也卸载到 CPU，仅在计算时传输到 GPU：
python
ds_config = {
"zero_optimization": {
"stage": 3,
"offload_optimizer": {"device": "cpu", "pin_memory": True},
"offload_param": {"device": "cpu", "pin_memory": True},
},
}NVMe Offload
将数据卸载到 NVMe SSD，突破 CPU 内存限制：
python
ds_config = {
"zero_optimization": {
"stage": 3,
"offload_optimizer": {
"device": "nvme",
"nvme_path": "/local_nvme",
"pin_memory": True,
},
"offload_param": {
"device": "nvme",
"nvme_path": "/local_nvme",
"pin_memory": True,
},
},
"aio": {
"block_size": 1048576,
"queue_depth": 8,
"thread_count": 1,
},
}显存与性能权衡
| 配置 | GPU 显存 (7B) | CPU 内存 | 训练速度 |
|---|---|---|---|
| ZeRO-2 | 18GB | 10GB | 100% |
| ZeRO-2 + Opt Offload | 8GB | 40GB | 65% |
| ZeRO-3 + Full Offload | 4GB | 60GB | 40% |
| ZeRO-3 + NVMe Offload | 4GB | 20GB | 25% |
PyTorch 原生 Offload
python
from torch.distributed.fsdp import CPUOffload
model = FSDP(
model,
cpu_offload=CPUOffload(offload_params=True),
)Offload 最佳实践
- 始终启用
pin_memory=True，加速 CPU-GPU 传输 - 使用大
gradient_accumulation_steps隐藏传输延迟 - 优先使用优化器状态 Offload（性价比最高）
- NVMe Offload 仅在 CPU 内存也不足时使用
带宽瓶颈
CPU-GPU 带宽（PCIe 4.0: 32GB/s）远低于 GPU 显存带宽（A100: 2TB/s）。Offload 场景下，通信是主要瓶颈。使用大 batch 和梯度累积来摊薄通信开销。