DeepSpeed Integration

TRL supports training with DeepSpeed, a library that implements advanced training optimization techniques. These include optimizer state partitioning, offloading, gradient partitioning, and more.

DeepSpeed integrates the Zero Redundancy Optimizer (ZeRO), which allows you to scale the model size proportional to the number of devices with sustained high efficiency.

Installation

To use DeepSpeed with TRL, install it using the following command:

pip install deepspeed

Running Training Scripts with DeepSpeed

No modifications to your training script are required. DeepSpeed is enabled through the training arguments, like for any Trainer: set deepspeed to the path of a DeepSpeed configuration file, and launch with torchrun.

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

torchrun --nproc_per_node 8 train.py

Scripts that parse their arguments with TrlParser, such as the ones behind the trl CLI, take it from the command line:

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

We provide ready-to-use DeepSpeed configuration files in the examples/deepspeed_configs directory. For example, to run training with ZeRO Stage 2:

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

Values set to "auto" in these files (batch size, gradient accumulation, precision, …) are filled in from the training arguments.

Additional Resources

Consult the 🤗 Transformers DeepSpeed documentation for more information about the DeepSpeed integration.