Multi-GPU

This guide covers advanced training configurations for multi-GPU setups using Axolotl.

Overview

When training on multiple GPUs, Axolotl supports 3 sharding/parallelism strategies. Additionally, you can layer specific optimization features on top of that strategy.

You generally cannot combine these strategies; they are mutually exclusive.

DeepSpeed: Powerful optimization library, supports ZeRO stages 1-3.
FSDP (Fully Sharded Data Parallel): PyTorch’s native sharding implementation (Recommended).
DDP (Distributed Data Parallel): PyTorch’s native parallelism implementation (Default if neither of the above are selected).

These features can often be combined with the strategies above:

Context Parallelism (CP): Splits long sequences across GPUs (Ringmaster under FSDP2; DeepSpeed Ulysses under DeepSpeed).
FSDP + QLoRA: Combines 4-bit quantization with FSDP (Specific to FSDP).

DeepSpeed

Configuration

Add to your YAML config:

deepspeed: deepspeed_configs/zero1.json

Usage

# Fetch deepspeed configs (if not already present)
axolotl fetch deepspeed_configs

# Passing arg via config
axolotl train config.yml

# Passing arg via cli
axolotl train config.yml --deepspeed deepspeed_configs/zero1.json

ZeRO Stages

We provide default configurations for:

ZeRO Stage 1 (zero1.json)
ZeRO Stage 1 with torch compile (zero1_torch_compile.json)
ZeRO Stage 2 (zero2.json)
ZeRO Stage 3 (zero3.json)
ZeRO Stage 3 with bf16 (zero3_bf16.json)
ZeRO Stage 3 with bf16 and CPU offload params(zero3_bf16_cpuoffload_params.json)
ZeRO Stage 3 with bf16 and CPU offload params and optimizer (zero3_bf16_cpuoffload_all.json)

Tip

Choose the configuration that offloads the least amount to memory while still being able to fit on VRAM for best performance.

Start from Stage 1 -> Stage 2 -> Stage 3.

Fully Sharded Data Parallel (FSDP)

FSDP allows you to shard model parameters, gradients, and optimizer states across data parallel workers.

Note

Only FSDP2 is supported. FSDP1 has been removed, and setting fsdp_version: 1 raises an error.

Configuration

fsdp_version: 2
fsdp_config:
  offload_params: false
  cpu_ram_efficient_loading: true
  auto_wrap_policy: TRANSFORMER_BASED_WRAP
  transformer_layer_cls_to_wrap: Qwen3DecoderLayer
  state_dict_type: FULL_STATE_DICT
  reshard_after_forward: true

When using auto_wrap_policy: SIZE_BASED_WRAP, set min_num_params to the minimum parameter count a module must have to be wrapped (e.g. 100000000).

With peft_ensure_weight_tying: true, tied embedding LoRA adapters and their output-head consumers stay in the root FSDP group under either wrapping policy. This keeps the shared parameters available for both uses and combines their gradients before FSDP/HSDP reduction. These weights remain gathered across the root forward, which can use more memory than separately wrapping the embedding and head. The root group follows the configured distributed mesh, including CP/EP data axes and TP on other layers. Keep the tied embedding and output head out of the tensor-parallel plan; TP-sharded adapters on those modules are rejected.

Full checkpoints and resume

fsdp_config.state_dict_type: FULL_STATE_DICT saves a single gathered model checkpoint and optimizer.bin. With expert parallelism, both checkpoints include every EP group’s experts. Restore selects each rank’s expert ownership before applying its FSDP/HSDP shard, including the different row and column packing of expert LoRA factors. LoRA/QLoRA Trainer checkpoints save only the adapter tensors; the frozen base model is loaded separately. Model-only intermediate checkpoints also write pytorch_model_fsdp.bin so Trainer uses the ownership-aware FSDP restore path.

For adamw_torch_8bit, the full optimizer file records quantized codes, block scales, lookup tables, and shard ownership. Resuming with the same layout preserves these exactly. Changing the mesh is supported when the saved blocks remain compatible with the new optimizer shards. A change that regroups blocks with different scales raises a restore-time error; it does not silently requantize optimizer state or disallow the training configuration.

Full checkpoints assemble tensors on rank 0’s CPU and use bounded GPU transfer buffers. Rank 0 therefore needs enough host memory for the complete checkpoint state. Adapter-only runs gather PEFT export state, including frozen adapter factors, saved modules, and required embeddings, without materializing other frozen base weights.

Native full bitsandbytes checkpoints preserve each EP owner’s packed weights and quantization metadata in a versioned pytorch_model_fsdp.bin payload. Resume supports rank rearrangements and compatible data-parallel resharding, but requires the same EP ownership ranges. Changing EP grouping requires an explicit checkpoint conversion. These files are training checkpoints; Hugging Face adapter exports retain their usual format. Generic Accelerator.get_state_dict() cannot represent packed EP weights and their owner-specific quantization metadata as an ordinary state dict, so it raises a collective error. Use the native save_fsdp_model() checkpoint path or Trainer adapter export instead. Ownership-aware full checkpoints support 4-bit dense weights and 4-bit or 8-bit parametrized experts; dense Int8Params in this path are rejected.

Final adapter exports use ordinary PEFT keys and gather every expert into the base model’s global layout. Models with fused expert parameters retain that adapter layout; this is separate from bitsandbytes quantization packing. Full-model Hugging Face exports use the original checkpoint conversion mapping, which unpacks fused Mixtral experts into separate w1, w2, and w3 weights. Final exports do not include the native resume file.

FSDP + QLoRA

For combining FSDP with QLoRA, see our dedicated guide.

With quantize_moe_experts: true and expert parallelism, EP slices the packed bitsandbytes expert weights and their block scales together. Double-quantized scales remain compressed when shard boundaries align. If only a nested scale block crosses a boundary, EP unpacks those scales without dequantizing or requantizing the weights. Expert slices must align with the weight quantization blocks. The packed parameters and quantization metadata stay on the same device, and FSDP wraps them on the expert mesh.

Migrating from FSDP1 to FSDP2

FSDP1 has been removed. fsdp_version: 1 and the bare top-level fsdp: list are rejected, and fsdp_version defaults to 2. To migrate an FSDP1 config, replace the fsdp: list with fsdp_config: and follow the config field mapping below to update field names.

Config mapping

FSDP1	FSDP2
fsdp_sharding_strategy	reshard_after_forward
fsdp_backward_prefetch_policy	REMOVED (dropped with a warning)
fsdp_backward_prefetch	REMOVED (dropped with a warning)
fsdp_forward_prefetch	REMOVED (rejected with an error)
fsdp_sync_module_states	REMOVED (dropped with a warning)
fsdp_limit_all_gathers	REMOVED (dropped with a warning)
fsdp_cpu_ram_efficient_loading	cpu_ram_efficient_loading
fsdp_state_dict_type	state_dict_type
fsdp_use_orig_params	REMOVED (dropped with a warning; FSDP2 always trains original parameters)
fsdp_activation_checkpointing	activation_checkpointing

For more details, please see the migration guide in the torchtitan repo. In Axolotl, if you were using the following FSDP1 config:

fsdp_version: 1
fsdp_config:
  fsdp_offload_params: false
  fsdp_cpu_ram_efficient_loading: true
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_transformer_layer_cls_to_wrap: Qwen3DecoderLayer
  fsdp_state_dict_type: FULL_STATE_DICT
  fsdp_sharding_strategy: FULL_SHARD

You can migrate to the following FSDP2 config:

fsdp_version: 2
fsdp_config:
  offload_params: false
  cpu_ram_efficient_loading: true
  auto_wrap_policy: TRANSFORMER_BASED_WRAP
  transformer_layer_cls_to_wrap: Qwen3DecoderLayer
  state_dict_type: FULL_STATE_DICT
  reshard_after_forward: true

Context parallelism

We support context parallelism (CP; other frameworks call the same thing sequence parallelism) via axolotl-ringmaster (context_parallel_size), which prefers Ulysses head-parallel attention and falls back to Ring attention for the remainder of the CP degree. This allows one to split up sequences across GPUs, which is useful in the event that a single sequence causes OOM errors during model training. Under DeepSpeed the same key maps onto DeepSpeed’s Ulysses (ALST) implementation through accelerate.

See our dedicated guide for more information.

Performance Optimization

Liger Kernel Integration

Please see docs for more info.

Troubleshooting

NCCL Issues

For NCCL-related problems, see our NCCL troubleshooting guide.

Common Problems

Reduce micro_batch_size
Reduce eval_batch_size
Adjust gradient_accumulation_steps
Consider using a higher ZeRO stage

Start with DeepSpeed ZeRO-2
Monitor loss values
Check learning rates

For more detailed troubleshooting, see our debugging guide.

--- title: "Multi-GPU" format: html: toc: true toc-depth: 3 # number-sections: true code-tools: true execute: enabled: false --- This guide covers advanced training configurations for multi-GPU setups using Axolotl. ## Overview {#sec-overview} When training on multiple GPUs, Axolotl supports 3 sharding/parallelism strategies. Additionally, you can layer specific optimization features on top of that strategy. You generally cannot combine these strategies; they are mutually exclusive. 1. **DeepSpeed**: Powerful optimization library, supports ZeRO stages 1-3. 2. **FSDP (Fully Sharded Data Parallel)**: PyTorch's native sharding implementation (Recommended). 3. **DDP (Distributed Data Parallel)**: PyTorch's native parallelism implementation (Default if neither of the above are selected). These features can often be combined with the strategies above: * **Context Parallelism** (CP): Splits long sequences across GPUs (Ringmaster under FSDP2; DeepSpeed Ulysses under DeepSpeed). * **FSDP + QLoRA**: Combines 4-bit quantization with FSDP (Specific to FSDP). ## DeepSpeed {#sec-deepspeed} ### Configuration {#sec-deepspeed-config} Add to your YAML config: ```{.yaml} deepspeed: deepspeed_configs/zero1.json ``` ### Usage {#sec-deepspeed-usage} ```{.bash} # Fetch deepspeed configs (if not already present) axolotl fetch deepspeed_configs # Passing arg via config axolotl train config.yml # Passing arg via cli axolotl train config.yml --deepspeed deepspeed_configs/zero1.json ``` ### ZeRO Stages {#sec-zero-stages} We provide default configurations for: - ZeRO Stage 1 (`zero1.json`) - ZeRO Stage 1 with torch compile (`zero1_torch_compile.json`) - ZeRO Stage 2 (`zero2.json`) - ZeRO Stage 3 (`zero3.json`) - ZeRO Stage 3 with bf16 (`zero3_bf16.json`) - ZeRO Stage 3 with bf16 and CPU offload params(`zero3_bf16_cpuoffload_params.json`) - ZeRO Stage 3 with bf16 and CPU offload params and optimizer (`zero3_bf16_cpuoffload_all.json`) ::: {.callout-tip} Choose the configuration that offloads the least amount to memory while still being able to fit on VRAM for best performance. Start from Stage 1 -> Stage 2 -> Stage 3. ::: ## Fully Sharded Data Parallel (FSDP) {#sec-fsdp} FSDP allows you to shard model parameters, gradients, and optimizer states across data parallel workers. ::: {.callout-note} Only FSDP2 is supported. FSDP1 has been removed, and setting `fsdp_version: 1` raises an error. ::: ### Configuration {#sec-fsdp-config} ```{.yaml} fsdp_version: 2 fsdp_config: offload_params: false cpu_ram_efficient_loading: true auto_wrap_policy: TRANSFORMER_BASED_WRAP transformer_layer_cls_to_wrap: Qwen3DecoderLayer state_dict_type: FULL_STATE_DICT reshard_after_forward: true ``` When using `auto_wrap_policy: SIZE_BASED_WRAP`, set `min_num_params` to the minimum parameter count a module must have to be wrapped (e.g. `100000000`). With `peft_ensure_weight_tying: true`, tied embedding LoRA adapters and their output-head consumers stay in the root FSDP group under either wrapping policy. This keeps the shared parameters available for both uses and combines their gradients before FSDP/HSDP reduction. These weights remain gathered across the root forward, which can use more memory than separately wrapping the embedding and head. The root group follows the configured distributed mesh, including CP/EP data axes and TP on other layers. Keep the tied embedding and output head out of the tensor-parallel plan; TP-sharded adapters on those modules are rejected. ### Full checkpoints and resume `fsdp_config.state_dict_type: FULL_STATE_DICT` saves a single gathered model checkpoint and `optimizer.bin`. With expert parallelism, both checkpoints include every EP group's experts. Restore selects each rank's expert ownership before applying its FSDP/HSDP shard, including the different row and column packing of expert LoRA factors. LoRA/QLoRA Trainer checkpoints save only the adapter tensors; the frozen base model is loaded separately. Model-only intermediate checkpoints also write `pytorch_model_fsdp.bin` so Trainer uses the ownership-aware FSDP restore path. For `adamw_torch_8bit`, the full optimizer file records quantized codes, block scales, lookup tables, and shard ownership. Resuming with the same layout preserves these exactly. Changing the mesh is supported when the saved blocks remain compatible with the new optimizer shards. A change that regroups blocks with different scales raises a restore-time error; it does not silently requantize optimizer state or disallow the training configuration. Full checkpoints assemble tensors on rank 0's CPU and use bounded GPU transfer buffers. Rank 0 therefore needs enough host memory for the complete checkpoint state. Adapter-only runs gather PEFT export state, including frozen adapter factors, saved modules, and required embeddings, without materializing other frozen base weights. Native full bitsandbytes checkpoints preserve each EP owner's packed weights and quantization metadata in a versioned `pytorch_model_fsdp.bin` payload. Resume supports rank rearrangements and compatible data-parallel resharding, but requires the same EP ownership ranges. Changing EP grouping requires an explicit checkpoint conversion. These files are training checkpoints; Hugging Face adapter exports retain their usual format. Generic `Accelerator.get_state_dict()` cannot represent packed EP weights and their owner-specific quantization metadata as an ordinary state dict, so it raises a collective error. Use the native `save_fsdp_model()` checkpoint path or Trainer adapter export instead. Ownership-aware full checkpoints support 4-bit dense weights and 4-bit or 8-bit parametrized experts; dense `Int8Params` in this path are rejected. Final adapter exports use ordinary PEFT keys and gather every expert into the base model's global layout. Models with fused expert parameters retain that adapter layout; this is separate from bitsandbytes quantization packing. Full-model Hugging Face exports use the original checkpoint conversion mapping, which unpacks fused Mixtral experts into separate `w1`, `w2`, and `w3` weights. Final exports do not include the native resume file. ### FSDP + QLoRA {#sec-fsdp-qlora} For combining FSDP with QLoRA, see our [dedicated guide](fsdp_qlora.qmd). With `quantize_moe_experts: true` and expert parallelism, EP slices the packed bitsandbytes expert weights and their block scales together. Double-quantized scales remain compressed when shard boundaries align. If only a nested scale block crosses a boundary, EP unpacks those scales without dequantizing or requantizing the weights. Expert slices must align with the weight quantization blocks. The packed parameters and quantization metadata stay on the same device, and FSDP wraps them on the expert mesh. ### Migrating from FSDP1 to FSDP2 {#sec-migrate-fsdp1-fsdp2} FSDP1 has been removed. `fsdp_version: 1` and the bare top-level `fsdp:` list are rejected, and `fsdp_version` defaults to `2`. To migrate an FSDP1 config, replace the `fsdp:` list with `fsdp_config:` and follow the config field mapping below to update field names. #### Config mapping FSDP1 | FSDP2 -------- | -------- fsdp_sharding_strategy | reshard_after_forward fsdp_backward_prefetch_policy | **REMOVED** (dropped with a warning) fsdp_backward_prefetch | **REMOVED** (dropped with a warning) fsdp_forward_prefetch | **REMOVED** (rejected with an error) fsdp_sync_module_states | **REMOVED** (dropped with a warning) fsdp_limit_all_gathers | **REMOVED** (dropped with a warning) fsdp_cpu_ram_efficient_loading | cpu_ram_efficient_loading fsdp_state_dict_type | state_dict_type fsdp_use_orig_params | **REMOVED** (dropped with a warning; FSDP2 always trains original parameters) fsdp_activation_checkpointing | activation_checkpointing For more details, please see the migration guide in the [torchtitan repo](https://github.com/pytorch/torchtitan/blob/main/docs/fsdp.md). In Axolotl, if you were using the following FSDP1 config: ```{.yaml} fsdp_version: 1 fsdp_config: fsdp_offload_params: false fsdp_cpu_ram_efficient_loading: true fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP fsdp_transformer_layer_cls_to_wrap: Qwen3DecoderLayer fsdp_state_dict_type: FULL_STATE_DICT fsdp_sharding_strategy: FULL_SHARD ``` You can migrate to the following FSDP2 config: ```{.yaml} fsdp_version: 2 fsdp_config: offload_params: false cpu_ram_efficient_loading: true auto_wrap_policy: TRANSFORMER_BASED_WRAP transformer_layer_cls_to_wrap: Qwen3DecoderLayer state_dict_type: FULL_STATE_DICT reshard_after_forward: true ``` ## Context parallelism {#sec-sequence-parallelism} We support context parallelism (CP; other frameworks call the same thing sequence parallelism) via [axolotl-ringmaster](https://github.com/axolotl-ai-cloud/ringmaster) (`context_parallel_size`), which prefers Ulysses head-parallel attention and falls back to Ring attention for the remainder of the CP degree. This allows one to split up sequences across GPUs, which is useful in the event that a single sequence causes OOM errors during model training. Under DeepSpeed the same key maps onto DeepSpeed's Ulysses (ALST) implementation through accelerate. See our [dedicated guide](sequence_parallelism.qmd) for more information. ## Performance Optimization {#sec-performance} ### Liger Kernel Integration {#sec-liger} Please see [docs](custom_integrations.qmd#liger) for more info. ## Troubleshooting {#sec-troubleshooting} ### NCCL Issues {#sec-nccl} For NCCL-related problems, see our [NCCL troubleshooting guide](nccl.qmd). ### Common Problems {#sec-common-problems} ::: {.panel-tabset} ## Memory Issues - Reduce `micro_batch_size` - Reduce `eval_batch_size` - Adjust `gradient_accumulation_steps` - Consider using a higher ZeRO stage ## Training Instability - Start with DeepSpeed ZeRO-2 - Monitor loss values - Check learning rates ::: For more detailed troubleshooting, see our [debugging guide](debugging.qmd).