deepspeedai/DeepSpeed

DeepSpeed – 可扩展的深度学习训练框架

它是什么 DeepSpeed 是一个开源库（基于 PyTorch 构建），可以在多种硬件（NVIDIA、AMD、Intel Gaudi、Xeon CPU 等）上高效训练超大规模深度学习模型——从数亿到数千亿参数。它提供系统级优化、内存节省技术和并行性原语，使研究人员和工程师无需编写底层 CUDA/HIP 代码即可扩展模型规模和速度。

关键能力

• ZeRO 系列 – 内存优化优化器（ZeRO、ZeRO-Infinity），在 GPU 间分片模型状态，支持万亿参数训练。
• 3D 并行 – 结合数据、张量和流水线并行，最大化硬件利用率。
• 卸载引擎 – SuperOffload、ZeRO-Offload、ZenFlow、Ulysses-Offload 将数据移至 CPU、GPU 内存或其他加速器，以减少 GPU 内存压力。
• 专用内核 – 自定义 C++/CUDA/HIP 操作（例如 AMD 的 System DMA、DeepNVMe I/O 扩展），可重叠通信和计算。
• 集成层 – 即用型适配器，支持 Hugging Face Transformers、Accelerate、PyTorch-Lightning、MosaicML Composer、Determined AI、MMEngine 等。
• 支持多种硬件后端 – NVIDIA、AMD、Intel Gaudi、Intel Xeon/XPU、华为昇腾 NPU 以及社区贡献的加速器。

典型工作流程

1. 安装 deepspeed pip 包（或从源码构建以支持自定义操作）。
2. 编写标准 PyTorch 训练脚本。
3. 添加一个小型 JSON 配置，告诉 DeepSpeed 使用哪些并行和内存卸载策略。
4. 使用 deepspeed train.py --deepspeed_config cfg.json 启动。
5. 使用 ds_report 验证编译了哪些操作以及哪些硬件功能处于活动状态。

为什么重要 DeepSpeed 已用于训练一些世界上最大的语言模型——例如 Microsoft 的 5300 亿参数 MT-NLG、BLOOM（1760 亿）、Jurassic-1（1780 亿）——证明该库可以扩展到 AI 研究的前沿。其开源特性让任何人都能复现这些结果或尝试新的并行思想。

入门与资源

• 安装 – pip install deepspeed（需要 PyTorch ≥ 2.0 和 CUDA/ROCm 编译器）。Windows 支持但有一些功能限制。
• 文档 – https://www.deepspeed.ai/（入门指南、JSON 配置参考、API 文档）。
• 社区 – 定期在 Zoom 上举办公开办公时间、Slack 工作区以及活跃的 GitHub 讨论。
• 示例 – 仓库附带示例脚本，涵盖 ZeRO-Offload、SDMA all-gather、长序列训练（Arctic Long Sequence Training）等。

谁应该使用

• 构建或微调超大型 Transformer 模型的研究人员。
• 需要从多 GPU 或异构集群中榨取最大性能的工程师。
• 希望为现有 PyTorch 代码库添加即插即用加速层的团队。

---

以上所有信息均直接取自项目 README。