DeepSpeed 集成

TRL 支持使用 DeepSpeed 进行训练，这是一个实现了高级训练优化技术的库。这些技术包括优化器状态分区、卸载、梯度分区等。

DeepSpeed 集成了 Zero Redundancy Optimizer (ZeRO)，它允许模型规模随设备数量成比例扩展，同时保持高效性能。

安装

要在 TRL 中使用 DeepSpeed，请使用以下命令进行安装：

pip install deepspeed

使用 DeepSpeed 运行训练脚本

您的训练脚本无需任何修改。DeepSpeed 通过训练参数启用，与任何 Trainer 一样：将 deepspeed 设置为 DeepSpeed 配置文件的路径，并使用 torchrun 启动。

training_args = SFTConfig(..., deepspeed="<DEEPSPEED_CONFIG_FILE.json>")

torchrun --nproc_per_node 8 train.py

使用 TrlParser 解析参数的脚本（例如 trl CLI 背后的脚本）会从命令行中获取它：

trl sft ... --deepspeed <DEEPSPEED_CONFIG_FILE.json>

我们在 examples/deepspeed_configs 目录中提供了可直接使用的 DeepSpeed 配置文件。例如，要使用 ZeRO Stage 2 运行训练：

trl sft ... --deepspeed examples/deepspeed_configs/zero2.json

在这些文件中设置为 "auto" 的值（批大小、梯度累积、精度等）将从训练参数中填充。

其他资源

有关 DeepSpeed 集成的更多信息，请参阅 🤗 Transformers 的 DeepSpeed 文档。