Multi-GPU

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

1 Overview

Axolotl supports several methods for multi-GPU training:

DeepSpeed (recommended)
FSDP (Fully Sharded Data Parallel)
Sequence parallelism
FSDP + QLoRA

2 DeepSpeed

DeepSpeed is the recommended approach for multi-GPU training due to its stability and performance. It provides various optimization levels through ZeRO stages.

2.1 Configuration

Add to your YAML config:

deepspeed: deepspeed_configs/zero1.json

2.2 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

2.3 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.

3 FSDP

3.1 Basic FSDP Configuration

fsdp:
  - full_shard
  - auto_wrap
fsdp_config:
  fsdp_offload_params: true
  fsdp_state_dict_type: FULL_STATE_DICT
  fsdp_transformer_layer_cls_to_wrap: LlamaDecoderLayer

4 Sequence parallelism

We support sequence parallelism (SP) via the ring-flash-attention project. 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.

First, install ring-flash-attn, recommended via pip install axolotl[ring-flash-attn], or from source with pip install .[ring-flash-attn].

Your Axolotl YAML config should contain the following lines:

sequence_parallel_degree: 4  # Split each sequence into 4 parts, one per GPU
flash_attention: true  # Required with sequence parallelism

# Optional; strides across the key dimension. Larger values use more memory but will make training faster.
heads_k_stride: 1

See our dedicated guide for more details.

4.1 FSDP + QLoRA

For combining FSDP with QLoRA, see our dedicated guide.

5 Performance Optimization

5.1 Liger Kernel Integration

Please see docs for more info.

6 Troubleshooting

6.1 NCCL Issues

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

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} Axolotl supports several methods for multi-GPU training: - DeepSpeed (recommended) - FSDP (Fully Sharded Data Parallel) - Sequence parallelism - FSDP + QLoRA ## DeepSpeed {#sec-deepspeed} DeepSpeed is the recommended approach for multi-GPU training due to its stability and performance. It provides various optimization levels through ZeRO stages. ### 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. ::: ## FSDP {#sec-fsdp} ### Basic FSDP Configuration {#sec-fsdp-config} ```{.yaml} fsdp: - full_shard - auto_wrap fsdp_config: fsdp_offload_params: true fsdp_state_dict_type: FULL_STATE_DICT fsdp_transformer_layer_cls_to_wrap: LlamaDecoderLayer ``` ## Sequence parallelism {#sec-sequence-parallelism} We support sequence parallelism (SP) via the [ring-flash-attention](https://github.com/zhuzilin/ring-flash-attention) project. 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. First, install `ring-flash-attn`, recommended via `pip install axolotl[ring-flash-attn]`, or from source with `pip install .[ring-flash-attn]`. Your Axolotl YAML config should contain the following lines: ```{.yaml} sequence_parallel_degree: 4 # Split each sequence into 4 parts, one per GPU flash_attention: true # Required with sequence parallelism # Optional; strides across the key dimension. Larger values use more memory but will make training faster. heads_k_stride: 1 ``` See our [dedicated guide](sequence_parallelism.qmd) for more details. ### FSDP + QLoRA {#sec-fsdp-qlora} For combining FSDP with QLoRA, see our [dedicated guide](fsdp_qlora.qmd). ## 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).