deepspeedai / DeepSpeed

deepspeedai/DeepSpeed

open_in_new前往仓库

DeepSpeed是一个深度学习优化库，让分布式训练和推理变得简单、高效且有效。

项目概览

项目概述

DeepSpeed 是一个由微软开发并维护的开源深度学习优化库，旨在让大规模分布式训练和推理变得简单、高效且有效。它通过一系列系统创新，如 ZeRO、ZeRO-Infinity、3D 并行、Ulysses 序列并行、DeepSpeed-MoE 等，突破了深度学习模型在规模、速度和易用性方面的限制。DeepSpeed 已成功用于训练多个世界级的大规模语言模型，如 Megatron-Turing NLG (530B)、BLOOM (176B) 等，并支持与 Hugging Face Transformers、PyTorch Lightning 等主流框架集成。

核心功能

  • ZeRO 优化器：通过分区优化器状态、梯度和参数，显著降低内存占用，支持训练超大模型。
  • 3D 并行：结合数据并行、模型并行和流水线并行，支持训练万亿参数模型。
  • ZeRO-Offload：将模型状态卸载到 CPU 或 NVMe，实现单 GPU 训练超大模型。
  • DeepSpeed-Inference：为 Transformer 模型提供高效的推理优化，包括模型并行、内核注入和量化。
  • DeepSpeed-MoE：支持混合专家模型的训练和推理，提升模型容量和效率。
  • 压缩库：提供多种模型压缩技术，如量化、剪枝和层缩减，以减小模型大小并加速推理。
  • 数据效率库：通过课程学习和随机层 token 丢弃，提高训练数据效率。
  • 通信效率优化：如 1-bit Adam、0/1 Adam 等，减少分布式训练中的通信开销。
  • 自动调优：自动搜索最佳 ZeRO 配置和微批大小，简化性能优化。
  • 性能分析工具：提供 Flops Profiler、通信日志等工具，帮助分析和调试训练性能。

适用与不适用场景

适用场景：

  • 训练大规模深度学习模型（数十亿至数万亿参数）。
  • 在有限 GPU 资源下训练超大模型（利用 CPU/NVMe 卸载）。
  • 需要混合精度训练（FP16/BF16）以加速训练。
  • 需要模型并行、流水线并行或混合并行策略。
  • 需要高效的模型推理，尤其是大型 Transformer 模型。
  • 需要模型压缩（量化、剪枝）以减小模型大小和推理成本。
  • 需要优化数据加载和训练效率（课程学习、token 丢弃）。

不适用场景：

  • 小型模型或简单任务，可能不需要如此复杂的优化。
  • 对训练确定性要求极高且无法接受随机性（但可选用确定性内核）。
  • 需要与特定框架深度绑定，而该框架未与 DeepSpeed 集成。
  • 硬件不支持所需特性（如 BF16 需要 A100 等）。
  • 需要 Windows 上的异步 I/O 或 GDS 功能（当前不支持）。

技术架构与依赖

DeepSpeed 基于 PyTorch 构建，是一个轻量级的封装库。其核心架构包括：

  • ZeRO 优化器：实现模型状态分区，支持多种阶段（0/1/2/3）。
  • 并行策略：支持数据并行、模型并行（张量并行）、流水线并行及其组合。
  • 内核优化：提供高度优化的 CUDA 内核（如 Transformer 内核、稀疏注意力内核）。
  • 通信优化：实现高效的通信原语，如 1-bit Adam 等。
  • 卸载机制：支持将参数、优化器状态卸载到 CPU 或 NVMe。
  • 推理引擎：支持模型并行推理和内核注入。

主要依赖：

  • Python 3.x
  • PyTorch（建议 >= 2.0）
  • CUDA 或 ROCm 编译器（如 nvcc、hipcc）
  • 可选：NVIDIA Apex（用于 FP16 训练）
  • 可选：tensorboard、wandb、comet_ml（用于监控）

安装与快速开始

安装：

pip install deepspeed

安装后，可通过 ds_report 验证安装并查看兼容的扩展。

快速开始：

  1. 准备一个 PyTorch 模型。
  2. 创建 DeepSpeed 配置文件（JSON）。
  3. 使用 deepspeed.initialize 初始化引擎。
  4. 使用 engine.backward 和 engine.step 进行训练。

示例：

import deepspeed

model = MyModel()
model_engine, optimizer, _, _ = deepspeed.initialize(
    args=args,
    model=model,
    model_parameters=model.parameters()
)

for batch in dataloader:
    loss = model_engine(batch)
    model_engine.backward(loss)
    model_engine.step()

典型使用方法

1. 使用 ZeRO 训练大模型：

在配置文件中设置 zero_optimization 的 stage 为 1、2 或 3，并指定相关参数。

2. 使用 ZeRO-Offload 在单 GPU 上训练大模型：

配置 offload_optimizer 和 offload_param 将状态卸载到 CPU 或 NVMe。

3. 使用 DeepSpeed-Inference 进行推理：

加载模型后，调用 deepspeed.init_inference 并指定模型并行度等参数。

4. 使用模型压缩：

在配置文件中启用 compression_training，并设置量化、剪枝等参数。

5. 使用自动调优：

设置 autotuning 为 enabled: true，DeepSpeed 将自动搜索最佳配置。

配置与部署要点

  • 配置文件：DeepSpeed 使用 JSON 配置文件，包含训练参数、优化器、调度器、ZeRO 设置等。
  • 多节点部署：通过 hostfile 指定节点和 GPU 资源。
  • 混合精度：支持 FP16、BF16 和 AMP，需根据硬件选择。
  • 检查点：支持保存和加载检查点，包括 ZeRO 分区状态。
  • 监控：可配置 TensorBoard、WandB、Comet 等监控后端。
  • 硬件支持：主要支持 NVIDIA GPU，也支持 AMD、Intel 等加速器（部分功能）。

限制、风险与许可证

限制：

  • 部分功能（如异步 I/O）在 Windows 上不可用。
  • 某些优化可能不适用于所有模型架构。
  • 需要一定的分布式训练知识。

风险：

  • 使用随机性内核可能影响训练的可复现性。
  • 硬件兼容性问题可能导致某些功能无法使用。

许可证：

Apache-2.0 许可证。

官方链接

信息来源和分析时间

本指南基于 GitHub 仓库 deepspeedai/DeepSpeed 的 README 和文档（截至 2026 年 8 月 10 日）编写。仓库资料未提供更多详细信息。