# Does LlamaFactory Support DeepSpeed for Training? Implementation Guide

> Yes LlamaFactory supports DeepSpeed training via accelerate enabling ZeRO-2 and ZeRO-3 strategies with automatic initialization and checkpoint management.

- Repository: [Yaowei Zheng/LlamaFactory](https://github.com/hiyouga/LlamaFactory)
- Tags: how-to-guide
- Published: 2026-03-04

---

**Yes, LlamaFactory fully supports DeepSpeed training through the accelerate library, enabling ZeRO-2 and ZeRO-3 distributed strategies with automatic engine initialization and checkpoint management.**

LlamaFactory integrates DeepSpeed as a first-class citizen for large language model training, allowing users to launch multi-GPU and multi-node jobs via the `accelerate` backend. The framework handles dependency validation, engine wrapper creation, and ZeRO-aware checkpoint saving automatically when you specify a DeepSpeed configuration file.

## How LlamaFactory Implements DeepSpeed Support

The DeepSpeed integration spans four layers: dependency checking, engine initialization, training step management, and checkpoint serialization.

### Dependency Validation

Before training begins, [`src/llamafactory/hparams/parser.py`](https://github.com/hiyouga/LlamaFactory/blob/main/src/llamafactory/hparams/parser.py) validates that the `deepspeed` package is installed when the `--deepspeed` flag is present. The `check_version()` function enforces this requirement strictly:

```python
if training_args.deepspeed:
    check_version("deepspeed", mandatory=True)

```

This check occurs during argument parsing, ensuring the training job fails fast with a clear message if the library is missing.

### Engine Initialization

In [`src/llamafactory/v1/core/base_trainer.py`](https://github.com/hiyouga/LlamaFactory/blob/main/src/llamafactory/v1/core/base_trainer.py), the `BaseTrainer` class detects the `"deepspeed"` distribution name and instantiates a `DistributedPlugin` wrapper. This plugin builds a `DeepSpeedEngine` around Hugging Face's `accelerate` integration:

```python
self._deepspeed_engine = DistributedPlugin("deepspeed")(...)

```

The engine automatically calls `accelerator.prepare()`, which triggers `deepspeed.initialize()` internally to return ZeRO-sharded models, optimizers, and learning rate schedulers according to your JSON configuration.

### Gradient Handling and Backward Pass

During the training loop, the DeepSpeed plugin in [`src/llamafactory/v1/plugins/trainer_plugins/distributed/deepspeed.py`](https://github.com/hiyouga/LlamaFactory/blob/main/src/llamafactory/v1/plugins/trainer_plugins/distributed/deepspeed.py) forwards loss computation to `accelerator.backward()`. This ensures proper gradient synchronization across data parallel ranks while respecting gradient accumulation boundaries:

```python
self.accelerator.backward(loss)

```

The wrapper manages `sync_gradients` internally to implement correct accumulation steps without manual intervention.

### ZeRO-3 Checkpoint Saving

For ZeRO-3 configurations where model parameters are partitioned across GPUs, the `save_model()` helper gathers sharded states before writing to disk. As implemented in [`deepspeed.py`](https://github.com/hiyouga/LlamaFactory/blob/main/deepspeed.py):

```python
state_dict = accelerator.get_state_dict(model)

```

This ensures the saved checkpoint contains the full, consolidated model weights rather than individual shards.

## Configuring DeepSpeed Training in LlamaFactory

You can configure DeepSpeed via the command line, Web UI helpers, or the Python API.

### CLI Training with Auto-Generated Configs

LlamaFactory provides a Web UI helper at [`src/llamafactory/webui/common.py`](https://github.com/hiyouga/LlamaFactory/blob/main/src/llamafactory/webui/common.py) that generates ready-to-use DeepSpeed JSON files. The `create_ds_config()` function writes [`ds_z2_config.json`](https://github.com/hiyouga/LlamaFactory/blob/main/ds_z2_config.json) and [`ds_z3_config.json`](https://github.com/hiyouga/LlamaFactory/blob/main/ds_z3_config.json) to the cache directory:

```bash

# Generate configurations (optional but recommended)

python -c "from llamafactory.webui.common import create_ds_config; create_ds_config()"

# Launch ZeRO-2 training

llamafactory-cli train \
    --model_name_or_path meta-llama/Meta-Llama-3-8B \
    --dataset path/to/dataset \
    --output_dir ./output \
    --deepspeed ./cache/default/ds_z2_config.json \
    --deepspeed_stage 2 \
    --bf16 True \
    --per_device_train_batch_size 4

```

The `--deepspeed` flag accepts any valid DeepSpeed JSON configuration file path.

### Programmatic Python API

For custom training scripts, pass the DeepSpeed configuration through the argument dictionary:

```python
from llamafactory.cli import get_train_args
from llamafactory.train.trainer import Trainer

args = get_train_args({
    "model_name_or_path": "meta-llama/Meta-Llama-3-8B",
    "dataset": "my_dataset",
    "output_dir": "./output",
    "deepspeed": "./ds_z3_config.json",
    "deepspeed_stage": 3,
    "bf16": True,
    "per_device_train_batch_size": 2,
})

trainer = Trainer(args)
trainer.train()

```

The `Trainer` automatically detects the DeepSpeed configuration and initializes the distributed backend accordingly.

### Saving ZeRO-3 Checkpoints Manually

After training, use the plugin's `save_model` function to handle parameter gathering:

```python
from llamafactory.v1.plugins.trainer_plugins.distributed.deepspeed import save_model

save_model(model, "./final_checkpoint", processor)

```

This ensures compatibility with standard Hugging Face model loading regardless of the ZeRO stage used during training.

## Limitations and Compatibility Constraints

Several features are incompatible with DeepSpeed ZeRO-3 due to the partitioning of model parameters across GPUs:

- **Generation during evaluation**: `predict_with_generate` raises a `ValueError` when ZeRO-3 is enabled because the full model weights are not available on any single rank.
- **Accuracy computation**: `compute_accuracy` metrics are blocked for the same reason.
- **Advanced optimizers**: GaLore, APOLLO, and BAdam layer-wise optimizers are disabled when DeepSpeed is active.
- **External accelerators**: Unsloth and K-Transformers integrations are incompatible with DeepSpeed.
- **PiSSA initialization**: Parameter-efficient initialization methods like PiSSA cannot run under ZeRO-3.

Additionally, distributed training must be launched using `llamafactory-cli` or `torchrun`. The [`base_trainer.py`](https://github.com/hiyouga/LlamaFactory/blob/main/base_trainer.py) enforces this with an explicit check:

```python
raise ValueError("Please launch distributed training with `llamafactory-cli` or `torchrun`.")

```

## Summary

- **Full integration**: LlamaFactory supports DeepSpeed ZeRO-2 and ZeRO-3 through the `accelerate` library with automatic engine management.
- **Simple configuration**: Use `--deepspeed path/to/config.json` or the `create_ds_config()` helper to generate optimized JSON files.
- **Automatic handling**: The framework manages `deepspeed.initialize()`, backward passes, and ZeRO-3 checkpoint gathering without manual code changes.
- **Known constraints**: ZeRO-3 blocks generation-based evaluation, certain optimizers, and must be launched via `llamafactory-cli` or `torchrun`.

## Frequently Asked Questions

### Does LlamaFactory require manual DeepSpeed configuration file creation?

No. While you can provide custom JSON files, the Web UI helper `create_ds_config()` in [`src/llamafactory/webui/common.py`](https://github.com/hiyouga/LlamaFactory/blob/main/src/llamafactory/webui/common.py) automatically generates optimized ZeRO-2 and ZeRO-3 configurations in the cache directory. You can reference these generated files directly via the `--deepspeed` flag.

### Can I use DeepSpeed ZeRO-3 with evaluation metrics that require model.generate()?

No. According to [`src/llamafactory/hparams/parser.py`](https://github.com/hiyouga/LlamaFactory/blob/main/src/llamafactory/hparams/parser.py), the framework raises a `ValueError` if `predict_with_generate` is enabled alongside ZeRO-3. This is because the model parameters are partitioned across GPUs, making full forward passes for generation impossible without expensive parameter gathering.

### What launch commands are compatible with LlamaFactory's DeepSpeed integration?

You must use either `llamafactory-cli train` or `torchrun` to launch distributed jobs. The `BaseTrainer` in [`src/llamafactory/v1/core/base_trainer.py`](https://github.com/hiyouga/LlamaFactory/blob/main/src/llamafactory/v1/core/base_trainer.py) explicitly checks the launch method and aborts with an error if it detects an unsupported distributed launcher.

### Are there any optimizer restrictions when using DeepSpeed in LlamaFactory?

Yes. The parser blocks GaLore, APOLLO, and BAdam layer-wise optimizers when DeepSpeed is enabled. Standard PyTorch optimizers like AdamW and those provided by DeepSpeed's built-in optimizer zoo remain fully compatible.