> ## Documentation Index
> Fetch the complete documentation index at: https://mintlify.com/huggingface/trl/llms.txt
> Use this file to discover all available pages before exploring further.
# DeepSpeed Integration
> Scale TRL training across multiple GPUs and nodes using DeepSpeed ZeRO optimizer state, gradient, and parameter partitioning.
TRL supports training with [DeepSpeed](https://github.com/microsoft/DeepSpeed), a library that implements advanced distributed training optimizations including optimizer state partitioning, gradient partitioning, parameter offloading, and more.
DeepSpeed integrates the [Zero Redundancy Optimizer (ZeRO)](https://huggingface.co/papers/1910.02054), which allows scaling model size proportionally to the number of devices while maintaining high efficiency.
## Installation
```bash theme={null}
pip install deepspeed
```
No modifications to your training script are required to use DeepSpeed. Simply launch with an Accelerate config file.
## ZeRO stages
DeepSpeed ZeRO has three stages, each partitioning progressively more state across devices:
Each GPU holds the full model parameters and gradients, but optimizer states (e.g., Adam momentum and variance) are sharded across GPUs.
* Memory reduction: \~4x for mixed precision training (optimizer states typically consume the most memory)
* Communication overhead: minimal
* Best for: reducing optimizer memory with low communication cost
```bash theme={null}
accelerate launch --config_file examples/accelerate_configs/deepspeed_zero1.yaml train.py
```
Optimizer states and gradients are both sharded across GPUs. Model parameters remain replicated.
* Memory reduction: \~8x for mixed precision training
* Communication overhead: similar to DDP
* Best for: most multi-GPU training scenarios
```bash theme={null}
accelerate launch --config_file examples/accelerate_configs/deepspeed_zero2.yaml train.py
```
Optimizer states, gradients, **and model parameters** are all sharded across GPUs. Parameters are gathered on demand during the forward and backward passes.
* Memory reduction: proportional to the number of GPUs
* Communication overhead: higher than Stage 2
* Best for: very large models that do not fit on a single GPU even with Stage 2
```bash theme={null}
accelerate launch --config_file examples/accelerate_configs/deepspeed_zero3.yaml train.py
```
## Running training with DeepSpeed
Use `accelerate launch` with a DeepSpeed config file. TRL provides ready-to-use config files in [`examples/accelerate_configs/`](https://github.com/huggingface/trl/tree/main/examples/accelerate_configs).
```bash theme={null}
acceleate launch --config_file train.py
```
For example, to train with ZeRO Stage 2:
```bash theme={null}
acceleate launch --config_file examples/accelerate_configs/deepspeed_zero2.yaml train.py
```
## Example accelerate configs
The following configs are provided in the TRL repository and can be used directly or adapted for your setup.
```yaml theme={null}
compute_environment: LOCAL_MACHINE
debug: false
deepspeed_config:
deepspeed_multinode_launcher: standard
gradient_accumulation_steps: 1
zero3_init_flag: false
zero_stage: 1
distributed_type: DEEPSPEED
downcast_bf16: 'no'
machine_rank: 0
main_training_function: main
mixed_precision: 'bf16'
num_machines: 1
num_processes: 8
rdzv_backend: static
same_network: true
use_cpu: false
```
```yaml theme={null}
compute_environment: LOCAL_MACHINE
debug: false
deepspeed_config:
deepspeed_multinode_launcher: standard
offload_optimizer_device: none
offload_param_device: none
zero3_init_flag: false
zero_stage: 2
distributed_type: DEEPSPEED
downcast_bf16: 'no'
machine_rank: 0
main_training_function: main
mixed_precision: 'bf16'
num_machines: 1
num_processes: 8
rdzv_backend: static
same_network: true
use_cpu: false
```
```yaml theme={null}
compute_environment: LOCAL_MACHINE
debug: false
deepspeed_config:
deepspeed_multinode_launcher: standard
offload_optimizer_device: none
offload_param_device: none
zero3_init_flag: true
zero3_save_16bit_model: true
zero_stage: 3
distributed_type: DEEPSPEED
downcast_bf16: 'no'
machine_rank: 0
main_training_function: main
mixed_precision: bf16
num_machines: 1
num_processes: 8
rdzv_backend: static
same_network: true
use_cpu: false
```
## Multi-node setup
For multi-node training, update `num_machines`, `machine_rank`, and `rdzv_backend` in your config file. The `deepspeed_multinode_launcher: standard` field is already set in the provided configs.
Consult the [Accelerate DeepSpeed documentation](https://huggingface.co/docs/accelerate/usage_guides/deepspeed) for detailed guidance on multi-node configuration.
## Additional resources
Full documentation for the DeepSpeed plugin in Accelerate.
Ready-to-use Accelerate config files for all ZeRO stages.
Zero Redundancy Optimizer — the foundational research paper.