How to Enable Distributed Training with LlamaFactory: FSDP, DeepSpeed, and Multi-Node Setups

Enable distributed training in LlamaFactory by setting the distributed field in your YAML configuration to fsdp2, deepspeed, or another supported backend, then launch the job with torchrun or the llamafactory-cli wrapper.

LlamaFactory provides a unified interface to scale large language model fine-tuning across multiple GPUs and nodes. Whether you are working with a single multi-GPU server or a cluster of machines, you can enable distributed training with LlamaFactory by selecting a backend-specific plugin and providing the appropriate configuration.

Architecture Overview

The distributed training stack in LlamaFactory is built around a plugin-based architecture that abstracts backend-specific implementations behind a common interface.

Component	Role	Source Location
DistributedInterface	Detects distributed context, initializes process groups, and manages DeviceMesh for parallelism.	src/llamafactory/v1/accelerator/interface.py
DistributedPlugin	Hub that routes to the correct engine based on the distributed config value.	src/llamafactory/v1/plugins/trainer_plugins/distributed/hub.py
FSDP2Engine	Implements Fully Sharded Data Parallel v2 with per-layer sharding, CPU offloading, and FP8 support.	src/llamafactory/v1/plugins/trainer_plugins/distributed/fsdp2.py
DeepSpeedEngine	Wrapper around Hugging Face Accelerate's DeepSpeedPlugin for ZeRO-3 and optimizer offloading.	src/llamafactory/v1/plugins/trainer_plugins/distributed/deepspeed.py
BaseTrainer	Core training loop that invokes DistributedPlugin(...).shard_model(model) when distributed config is present.	src/llamafactory/v1/core/base_trainer.py
Launcher	Handles torchrun re-launch, sets MASTER_ADDR, MASTER_PORT, and WORLD_SIZE, then starts training.	src/llamafactory/v1/launcher.py

When you enable distributed training with LlamaFactory, the BaseTrainer checks for a dist_config object. If present, it loads the DistributedPlugin hub, which instantiates the appropriate engine (FSDP2, DeepSpeed, etc.) to shard the model before the training loop begins.

Supported Distributed Backends

LlamaFactory supports multiple distributed training strategies, each optimized for different hardware configurations and model sizes.

Backend	Config Key	Best For
FSDP	fsdp	Traditional PyTorch FSDP (v1) with configurable sharding strategies.
FSDP2	fsdp2	Modern FSDP implementation with per-layer sharding, CPU offloading, and FP8 support.
DeepSpeed	deepspeed	ZeRO-3 optimization for very large models with optimizer state sharding.
Megatron-LM	megatron	NVIDIA's Megatron for tensor and pipeline parallelism.
Ray	ray	Multi-node clusters managed by Ray Train.

All backends require the distributed field in your YAML configuration to activate the corresponding plugin in src/llamafactory/v1/plugins/trainer_plugins/distributed/hub.py.

Configuration Guide

To enable distributed training with LlamaFactory, you must provide a YAML configuration file that specifies the backend and its parameters.

FSDP2 Configuration

FSDP2 is the recommended backend for most multi-GPU training scenarios. It provides fine-grained control over sharding and memory optimization.

# examples/accelerate/fsdp2_config.yaml

distributed: fsdp2
dist_config:
  mixed_precision: bf16        # Options: bf16, fp16, fp32

  offload_params: true         # Offload parameters to CPU

  reshard_after_forward: true  # Reshard parameters after forward pass

  pin_memory: true             # Pin memory for CPU offloading

  fp8_enable_fsdp_float8_all_gather: false  # Enable FP8 all-gather (requires H100)

model_name_or_path: meta-llama/Llama-2-7b-hf
train_dataset: ./data/train.jsonl
output_dir: ./outputs/fsdp2_run

The FSDP2Engine in src/llamafactory/v1/plugins/trainer_plugins/distributed/fsdp2.py processes these options to configure FullyShardedDataParallel with the appropriate sharding policies.

DeepSpeed Configuration

For training very large models that do not fit in a single GPU's memory, use DeepSpeed with ZeRO-3 optimization.

# examples/v1/train_full/train_full_deepspeed.yaml

distributed: deepspeed
deepspeed: examples/deepspeed/ds_z3_config.json  # Path to DeepSpeed JSON config

model_name_or_path: meta-llama/Llama-2-70b-hf
train_dataset: ./data/train.jsonl
output_dir: ./outputs/deepspeed_z3

The DeepSpeedEngine wrapper in src/llamafactory/v1/plugins/trainer_plugins/distributed/deepspeed.py integrates with Hugging Face Accelerate to apply the ZeRO-3 partitioning strategy defined in your JSON configuration.

Ray Multi-Node Configuration

For distributed training across multiple physical machines, configure the Ray backend.

distributed: ray
ray_num_workers: 4               # Total workers across cluster

master_addr: 10.0.0.1            # Head node address (auto-detected if omitted)

master_port: 29500
model_name_or_path: meta-llama/Llama-2-7b-hf
train_dataset: ./data/train.jsonl
output_dir: ./outputs/ray_run

Ray handles the process group initialization automatically, while LlamaFactory's DistributedInterface in src/llamafactory/v1/accelerator/interface.py manages the DeviceMesh creation for the Ray cluster.

Launching Distributed Jobs

Once your configuration file is ready, you must launch the training job using a distributed launcher.

Single-Node Multi-GPU

For training on a single machine with multiple GPUs, use torchrun to spawn the processes:

torchrun --nproc_per_node=4 -m torch.distributed.run \
  src/llamafactory/v1/launcher.py train path/to/config.yaml

The launcher.py script automatically sets MASTER_ADDR, MASTER_PORT, and WORLD_SIZE if they are not already defined in your environment.

Using the Built-in CLI Wrapper

Alternatively, use the provided CLI wrapper which handles the torchrun invocation internally:

llamafactory-cli train path/to/config.yaml

This command automatically detects the number of available GPUs and launches the appropriate distributed configuration based on your YAML settings.

FP8 Mixed Precision with Distributed Training

For NVIDIA H100 GPUs, you can enable FP8 mixed precision together with FSDP2 to reduce memory bandwidth and increase throughput.

Add the following to your configuration:

fp8: true                 # Enable FP8 via Accelerate

fp8_enable_fsdp_float8_all_gather: true   # Optimized all-gather for FSDP2

The fp8_utils.py module in src/llamafactory/train/fp8_utils.py propagates the FP8_ENABLE_FSDP_FLOAT8_ALL_GATHER environment variable before the FSDP2Engine wraps your model, enabling optimized FP8 communication between GPUs.

Summary

• Enable distributed training with LlamaFactory by setting the distributed field in your YAML configuration to fsdp2, deepspeed, megatron, or ray.
• The DistributedPlugin hub in src/llamafactory/v1/plugins/trainer_plugins/distributed/hub.py automatically routes to the correct backend engine based on your configuration.
• FSDP2 is recommended for most multi-GPU scenarios, offering per-layer sharding and CPU offloading via src/llamafactory/v1/plugins/trainer_plugins/distributed/fsdp2.py.
• DeepSpeed with ZeRO-3 is ideal for very large models that require optimizer state sharding.
• Launch distributed jobs using torchrun --nproc_per_node=N src/llamafactory/v1/launcher.py train config.yaml or the llamafactory-cli train wrapper.

Frequently Asked Questions

What is the difference between FSDP and FSDP2 in LlamaFactory?

FSDP refers to the original PyTorch Fully Sharded Data Parallel implementation, configured via Accelerate config files. FSDP2 is LlamaFactory's native implementation in src/llamafactory/v1/plugins/trainer_plugins/distributed/fsdp2.py, offering more granular control over per-layer sharding, CPU offloading, and FP8 support. FSDP2 is generally recommended for new projects.

Do I need to modify my model code to use DeepSpeed with LlamaFactory?

No. LlamaFactory's DeepSpeedEngine wrapper in src/llamafactory/v1/plugins/trainer_plugins/distributed/deepspeed.py handles the integration automatically. You only need to provide a valid DeepSpeed JSON configuration file path in the deepspeed field of your YAML config. The wrapper applies ZeRO-3 partitioning without requiring changes to your model architecture.

Can I use FP8 precision with any distributed backend?

FP8 support is currently optimized for FSDP2 when running on NVIDIA H100 GPUs. While you can set fp8: true with other backends, the fp8_enable_fsdp_float8_all_gather optimization specifically requires FSDP2 as implemented in src/llamafactory/v1/plugins/trainer_plugins/distributed/fsdp2.py. The fp8_utils.py module ensures the correct environment variables are set before model wrapping occurs.

How do I debug distributed training failures in LlamaFactory?

Start by checking the DistributedInterface initialization in src/llamafactory/v1/accelerator/interface.py to verify that is_distributed is correctly detecting your environment. Ensure MASTER_ADDR, MASTER_PORT, and WORLD_SIZE are properly set by the launcher in src/llamafactory/v1/launcher.py. For backend-specific issues, inspect the engine logs: FSDP2 issues typically surface in src/llamafactory/v1/plugins/trainer_plugins/distributed/fsdp2.py, while DeepSpeed problems appear in the corresponding deepspeed.py wrapper.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}