总目录›Axolotl›第 8 章
Chapter 08

多 GPU：DeepSpeed / FSDP 配置

📌 commit ab1a0d8两种主流多卡方案在 Axolotl 里的写法

Axolotl 通过 accelerate 支持 DeepSpeed 和 FSDP 两条路径。本章给推荐配方和何时选哪个的判断。

8.1DeepSpeed

deepspeed: deepspeed_configs/zero2.json
# 或 inline
deepspeed:
  zero_optimization:
    stage: 2
    offload_optimizer:
      device: cpu
      pin_memory: true
  bf16:
    enabled: true
  train_batch_size: auto
  train_micro_batch_size_per_gpu: auto
  gradient_accumulation_steps: auto
  gradient_clipping: 1.0

Axolotl 自带 4 个 ds 配置：

文件用途
deepspeed_configs/zero1.jsonZeRO-1（切 optim）
zero2.jsonZeRO-2（+grad）★ 常用
zero3.jsonZeRO-3（+param）
zero3_bf16.jsonZeRO-3 + bf16 配方

8.2FSDP

fsdp:
  - full_shard
  - auto_wrap
fsdp_config:
  fsdp_offload_params: false
  fsdp_sync_module_states: true
  fsdp_use_orig_params: true             # ★ PEFT 必须 true
  fsdp_cpu_ram_efficient_loading: true
  fsdp_state_dict_type: SHARDED_STATE_DICT
  fsdp_transformer_layer_cls_to_wrap: LlamaDecoderLayer
  fsdp_activation_checkpointing: true

8.3DeepSpeed vs FSDP 选哪个

场景选原因
13B 8 卡 LoRADeepSpeed ZeRO-2最简单稳定
70B 8 卡 QLoRADeepSpeed ZeRO-3切 param 必备
多机 70B 全参DeepSpeed ZeRO-3成熟
PyTorch 原生 + 长上下文FSDPHSDP（hybrid shard）有优势
跟 torch.compile 配FSDPDeepSpeed 暂不兼容
非 NVIDIA GPU (AMD/Intel)FSDPDeepSpeed CUDA-only

8.4多机启动

# 节点 0
accelerate launch \
    --num_processes 16 \
    --num_machines 2 \
    --machine_rank 0 \
    --main_process_ip 10.0.0.5 \
    --main_process_port 12345 \
    -m axolotl.cli.train my_qlora.yml

# 节点 1（同样命令，machine_rank=1）

8.513B 8 卡 SFT 配方（DeepSpeed ZeRO-2）

base_model: meta-llama/Meta-Llama-3-8B
adapter: lora
lora_r: 32
lora_alpha: 64
lora_target_modules: [q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj]

deepspeed: deepspeed_configs/zero2.json

sequence_len: 4096
sample_packing: true
micro_batch_size: 4
gradient_accumulation_steps: 4
num_epochs: 1
learning_rate: 2e-4
optimizer: adamw_torch_fused
bf16: auto
flash_attention: true
gradient_checkpointing: true

8.670B QLoRA 配方（DeepSpeed ZeRO-3）

base_model: meta-llama/Meta-Llama-3-70B
load_in_4bit: true
adapter: qlora
lora_r: 16

deepspeed: deepspeed_configs/zero3_bf16.json

sequence_len: 2048
sample_packing: true
micro_batch_size: 1
gradient_accumulation_steps: 16
num_epochs: 1
learning_rate: 1e-4
optimizer: paged_adamw_8bit
bf16: auto
flash_attention: true
gradient_checkpointing: true

8.7常见踩坑

现象处理
FSDP + LoRA "param has no attribute _fsdp_wrapped_module"开 fsdp_use_orig_params: true
ZeRO-3 + LoRA save 时 OOM用 save_safetensors: true + save_only_model: true
多机 NCCL hang用 IB；查 NCCL_SOCKET_IFNAME
FSDP + torch.compile 慢启动warmup compile 几步丢掉

8.8仓库自带的 8 份 DeepSpeed config

不需要自己写 ds_config.json。deepspeed_configs/ 下 8 份 preset，按需选：

文件ZeRO stage典型场景
zero1.json 1（optimizer state 切） 小模型；速度优先
zero1_torch_compile.json 1 + torch.compile 小模型 + 编译加速
zero2.json 2（+ gradient 切） ★ 中等模型最常用
zero2_torch_compile.json 2 + torch.compile 同上配 fused 算子
zero3.json 3（+ param 切） 70B 全参 SFT
zero3_bf16.json 3 + bf16 优化 同上但 H100/A100
zero3_bf16_cpuoffload_params.json 3 + param CPU offload 显存极紧；速度降 30-50%
zero3_bf16_cpuoffload_all.json 3 + 全部 offload 最省显存；速度降 60%+

用法：deepspeed: deepspeed_configs/zero3_bf16.json 一行字段，axolotl 自动用 accelerate 接入。

8.9FSDP1 vs FSDP2

v2025-Q3 之后 Axolotl 同时支持 FSDP1 和 FSDP2。fsdp_version 字段控制：

维度FSDP1FSDP2（推荐）
YAML 触发 fsdp_version: 1 fsdp_version: 2
底层 API FullyShardedDataParallel fully_shard()（DTensor）
支持 ND 并行 否 ✅（DP × CP × TP）
QLoRA 兼容 需要 patch ✅ 原生
LoRA 与 trainable / frozen 混合需 fsdp_use_orig_params: true✅ 自动
checkpoint save SHARDED_STATE_DICT DTensor checkpoint

新项目首选 FSDP2。老 YAML 升级时把 fsdp_version: 1 改成 2，fsdp_config 子字段大部分兼容。

8.10ND Parallelism（2025-Q3 新）

FSDP2 + Context Parallel + Tensor Parallel 三维并行。配置示例：

fsdp_version: 2
fsdp_config:
  fsdp_offload_params: false

# 三个维度
context_parallel_size: 2     # CP 维（沿 seq 切）
tensor_parallel_size: 2      # TP 维（沿权重切）
# DP 维 = nproc / (CP * TP)，自动算

# 长上下文必备
sequence_len: 32768
sample_packing: true

实战：8 卡机配 CP=2 TP=2 DP=2，能训 32K 上下文的 70B 模型 LoRA。这是 Axolotl 跟 verl / Megatron-LM 路径拉平的关键能力。

8.11Sequence Parallelism / Ring Attention

2025-Q2 引入的 SP 走 src/axolotl/monkeypatch/ring_attn/。启用：

sequence_parallel_size: 4
sequence_parallel_mode: ring_attn   # 或 ulysses
flash_attention: true               # 必须配 fa2

跟 ND Parallelism 互斥（都管 seq 维）；想要极长 CoT 训练时选 SP；想三维并行 / TP 时选 ND。

8.12这章你需要带走的