DeepSpeed 旨在减少计算能力和内存使用，并在大规模分布式模型上使用现有计算机硬件进行更好的并行训练。针对低延迟、高吞吐量训练进行了优化。包括用于训练具有 1 万亿或更多参数的模型的零冗余优化器 (ZeRO)。功能包括混合精度训练、单 GPU、多 GPU 和多节点训练以及自定义模型并行。
安装
从源代码安装 DeepSpeed
克隆 DeepSpeed 代码库之后，可以通过 pip 以 JIT 模式快速安装 DeepSpeed（不会编译任何 C++/CUDA 源文件）。
单个节点安装
pip install .多节点安装
对于跨越多个节点的安装，我们建议使用代码库中的 install.sh 脚本安装 DeepSpeed。这将本地构建一个 Python wheel 文件，并将其复制到主机文件中列出的所有节点（主机文件可以通过 --hostfile 参数指定，默认为 /job/hostfile）。
首次使用 DeepSpeed 的 JIT 编译
首次使用 DeepSpeed 的代码时，它将仅自动构建运行所需的 CUDA 扩展程序，并默认将它们放置在 ~/.cache/torch_extensions/ 目录下。下次执行相同程序时，将从该目录加载这些预编译的扩展程序。
多虚拟环境注意事项
如果使用多个虚拟环境，这可能会带来问题。默认情况下，只有一个 torch_extensions 目录，但不同的虚拟环境可能使用不同的设置（例如，不同的 Python 或 CUDA 版本），然后加载另一个环境构建的 CUDA 扩展程序将失败。因此，如果需要，可以使用 TORCH_EXTENSIONS_DIR 环境变量覆盖默认位置。这样，可以在每个虚拟环境中将其指向唯一的目录，DeepSpeed 将使用它来保存和加载 CUDA 扩展程序。
可以通过以下方式仅为特定运行更改它：
TORCH_EXTENSIONS_DIR=./torch-extensions deepspeed ...使用 Conda 构建环境
如果在使用默认系统环境进行编译时遇到困难，可以尝试使用提供的 conda 环境，其中包含必要的编译工具链和 PyTorch。
conda env create -n deepspeed -f environment.yml --force激活该环境后，尝试上述安装命令。
