DeepSpeed
一、项目简介
DeepSpeed 是微软开源的大规模训练优化库。其 ZeRO 系列通过切分优化器状态、梯度与参数,配合 CPU/NVMe 卸载,让在有限 GPU 上训练千亿参数模型成为现实,是超大模型训练的核心工具之一。
二、核心特性
- ZeRO-1/2/3:逐级切分显存,显存占用大幅下降;
- Offload / Infinity:把状态卸载到 CPU 内存甚至 NVMe;
- 3D 并行:数据并行 + 张量并行 + 流水线并行;
- 混合精度与稀疏注意力;
- 与 HF Trainer / Accelerate 集成。
三、适用场景
- 大模型预训练与全参数微调;
- 多机多卡分布式训练;
- 显存极度受限下的训练探索。
四、快速上手
pip install deepspeed
deepspeed --num_gpus=4 train.py --deepspeed ds_config.json
五、注意事项
- 配置较复杂,需按集群拓扑与模型规模调优;
- 通信开销可能成为瓶颈,需高速网络;
- 建议从小规模验证开始再扩展。