Chapter 10

加速集成：DeepSpeed / FSDP / Unsloth / FlashAttention

📌 本章配置示例来自 examples/deepspeed/ 、 examples/accelerate/ · commit 01398eb18dd4 。

10.1 LLaMA-Factory 的"加速三选一"

第 1 章说过 LLaMA-Factory 自己不写底层 kernel。它的加速能力全部来自下游 plug-in：

加速后端	主战场	YAML 开关
DeepSpeed	多卡 ZeRO，多机扩展	`deepspeed: examples/deepspeed/ds_zN_*.json`
FSDP	PyTorch 原生 ZeRO 等价物	用 `accelerate config` 配 + `accelerate launch`
Unsloth	单卡 LoRA / QLoRA 2-5× 加速	`use_unsloth: true`
FlashAttention	降 attention 显存与 latency	`flash_attn: fa2`
Liger Kernel	fused RMSNorm / cross-entropy	`enable_liger_kernel: true`

它们不互斥，但同时开启时要清楚约束，本章逐个讲。

10.2 DeepSpeed：仓库自带的 7 份配置

examples/deepspeed/ 已经写好：

文件	用途
`ds_z0_config.json`	仅做混精，没有切分（基本不用）
`ds_z2_config.json`	ZeRO-2：切优化器 + 梯度
`ds_z2_offload_config.json`	ZeRO-2 + CPU 卸载优化器
`ds_z2_autotp_config.json`	ZeRO-2 + 自动 TP（实验性）
`ds_z3_config.json`	ZeRO-3：参数也切
`ds_z3_offload_config.json`	ZeRO-3 + CPU 卸载（70B 神器）
`ds_z3_fp8_config.json`	ZeRO-3 + FP8（H100 专享）

典型 `ds_z3_config.json` 解读

{
  "train_batch_size":                  "auto",   // 让框架按 (micro × accum × world) 算
  "train_micro_batch_size_per_gpu":    "auto",   // 同上，由 YAML 推
  "gradient_accumulation_steps":       "auto",
  "gradient_clipping":                 "auto",
  "zero_allow_untested_optimizer":     true,
  "bf16": { "enabled": "auto" },
  "fp16": { "enabled": "auto" },

  "zero_optimization": {
    "stage":                                   3,
    "overlap_comm":                            false,
    "contiguous_gradients":                    true,
    "sub_group_size":                          1e9,
    "reduce_bucket_size":                      "auto",
    "stage3_prefetch_bucket_size":             "auto",
    "stage3_param_persistence_threshold":      "auto",
    "stage3_max_live_parameters":              1e9,
    "stage3_max_reuse_distance":               1e9,
    "stage3_gather_16bit_weights_on_model_save": true
  }
}

大部分字段是 "auto"，意思是由 LLaMA-Factory 在启动时按 YAML 字段填上去。所以你几乎不用碰 DS 配置本身，只动 YAML 里的 batch / accum / bf16。

怎么选 ZeRO stage

模型 / 微调方式	显存 / 卡	推荐 stage
7B LoRA	≥ 24 GB	不需要 DS（单卡或 DDP）
13B LoRA	≥ 24 GB	z2
13B 全参	≥ 40 GB × 8	z3
70B QLoRA	24 GB	z2 / z3
70B 全参	80 GB × 8	z3
70B 全参在 24 GB 上	极限	z3 + offload

启动命令

# 单机多卡（4 卡 A100）
FORCE_TORCHRUN=1 llamafactory-cli train examples/train_full/qwen3_full_sft.yaml

# 多机：仍然走 torchrun（DS 内部用 NCCL）
NNODES=2 NODE_RANK=0 MASTER_ADDR=10.0.0.1 MASTER_PORT=29500 \
  FORCE_TORCHRUN=1 llamafactory-cli train examples/train_full/qwen3_full_sft.yaml

YAML 里只要写 deepspeed: examples/deepspeed/ds_z3_config.json，框架会自动把这份 JSON 注入 DS 引擎。

10.3 FSDP（通过 accelerate）

FSDP 走 HF accelerate 配置文件而不是 DS JSON。仓库 examples/accelerate/ 提供：

examples/accelerate/
├── fsdp_config.yaml
├── fsdp_config_offload.yaml
└── single_config.yaml

启动方式：

accelerate launch --config_file examples/accelerate/fsdp_config.yaml \
    -m llamafactory.cli train examples/train_full/qwen3_full_sft.yaml

YAML 里不要写 deepspeed:，否则会冲突；FSDP 的并行策略由 accelerate config 决定。

场景	建议
已经熟 DS、不在乎换	用 DeepSpeed
想完全 PyTorch 原生 / 跟 PyTorch 新版	用 FSDP
未来上 PyTorch DTensor / FSDP2	用 FSDP
多机扩展性 / 上 NVMe offload	用 DeepSpeed（ZeRO-Infinity）

10.4 Unsloth 集成

Unsloth 用手写 Triton kernel 给 LoRA / QLoRA 提速 2-5×，详见本目录 Unsloth 那本书。在 LLaMA-Factory 里启用，只要 YAML 加一行：

use_unsloth: true

注意约束：

当前 Unsloth OSS 只支持单卡。多卡同时开 Unsloth + DeepSpeed 会冲突，框架会自动回退；
支持的模型有白名单（Llama、Mistral、Qwen、Gemma、Phi 等），新架构可能不支持；
已经搭配 use_unsloth: true 时，flash_attn 由 Unsloth 内部控制，不必再设；
QLoRA + Unsloth 是经典组合，quantization_bit: 4 + use_unsloth: true 在 4090 上能训 70B 长上下文。

10.5 FlashAttention：默认就该打开

flash_attn: fa2                # 推荐
# 可选：auto / disabled / sdpa

`flash_attn`	含义
`auto`	自动检测，能 fa2 用 fa2，否则 sdpa
`fa2`	显式用 FlashAttention 2/3（需要装 `flash-attn`）
`sdpa`	用 PyTorch 自带 `scaled_dot_product_attention`
`disabled`	关闭，回到原 attention 实现（debug 时）

训练长上下文（≥ 8K）必须开 fa2，否则 attention 显存按 $O(L^2)$ 爆炸。

10.6 Liger Kernel：fused 算子库

Liger 是 LinkedIn 出的一组 Triton fused kernel（RMSNorm、RoPE、cross-entropy、SwiGLU）。LLaMA-Factory 通过：

enable_liger_kernel: true

能在不开 Unsloth 时也获得 10-30% 加速。和 FSDP / DeepSpeed 兼容；和 Unsloth 同开则两者会相互覆盖 patch，建议二选一。

10.7 同时开多种加速时的"优先级"

当一份 YAML 同时打开 DeepSpeed、FlashAttention、Liger、Unsloth 时，框架行为：

组合	结果
DeepSpeed + flash_attn fa2	正常协同
DeepSpeed + Liger	正常协同
DeepSpeed + Unsloth	冲突：当前 Unsloth OSS 不支持多卡，自动忽略 use_unsloth
FSDP + flash_attn fa2	正常协同
Unsloth + flash_attn fa2	Unsloth 自带 attention 优化，`flash_attn` 由其接管
Unsloth + Liger	不要同开，patch 路径冲突
QLoRA + DeepSpeed ZeRO-3	需要 bnb >= 0.43，`quantization_device_map: auto`

10.8 把"加速"做成思维模型

所有加速都在动这三笔账：

$$\text{总显存} \approx \underbrace{P_\text{参数}}_{\text{切?压?}} + \underbrace{G_\text{梯度}}_{\text{切?}} + \underbrace{O_\text{优化器}}_{\text{切?卸?}} + \underbrace{A_\text{激活}}_{\text{recompute?fa2?}}$$

手段	动哪一项
DeepSpeed ZeRO-1	切 $O$
ZeRO-2	切 $O + G$
ZeRO-3 / FSDP	切 $O + G + P$
Offload	把 $O$（或 $P$）卸到 CPU / NVMe
QLoRA	用 4-bit 表示 $P$
gradient_checkpointing	少存 $A$，反向时重算
FlashAttention	降 $A$ 的常数（不存 attention map）
Unsloth / Liger	kernel-level 降 $A$ 与 latency

遇到 OOM，从右往左加：先开 fa2 + gradient_checkpointing；不够开 ZeRO-2；再不够 ZeRO-3；再不够 offload；再不够 QLoRA。

10.9 GaLore / APOLLO：低秩优化器的真实字段

GaLore（2024-03）和 APOLLO（2025-01）是两个把"梯度投影到低秩子空间"作为加速手段的优化器。LLaMA-Factory 通过 FinetuningArguments → GaloreArguments / ApolloArguments mixin 暴露。落到代码层在 train/trainer_utils.py:200–349。

YAML 字段	含义	典型值
`use_galore` / `use_apollo`	开关	二选一
`galore_target` / `apollo_target`	目标模块，`all` 自动找 nn.Linear	`all`
`galore_rank` / `apollo_rank`	低秩维度	16 / 256
`galore_update_interval`	多少步重投一次	200
`galore_scale`	梯度缩放	0.25–4.0
`galore_proj_type`	`std / reverse_std / right / left / full`	std
`galore_layerwise` / `apollo_layerwise`	每层独立 optimizer dict；要求 grad_accum=1	false
`apollo_proj` / `apollo_proj_type`	投影变种（APOLLO 特有）	—
`apollo_scale_type` / `apollo_scale_front`	scale 应用方式	—

"GaLore 还是 APOLLO？"经验：

纯节省优化器显存 → GaLore（成熟，社区验证最多）；
追求收敛速度提升 → APOLLO（Hessian-aware）；
都可以和 LoRA 共存，但跟 BAdam / Muon 不能同时开。

10.10 Muon：2025-04 新优化器

Muon 不通过 use_muon flag，而是写在 HF Trainer 的 optim 字段（同 adamw_torch 一类）：

optim: muon                  # 或 muon_8bit
optim_args: "..."            # Muon 额外参数

Muon 是 newton-schulz 风格的二阶优化器，对长训练比 AdamW 省一半显存（无一阶矩），收敛速度也更稳。但前期 step 比 AdamW 慢，小数据 / 短训练不要换 Muon。

10.11 BAdam：块坐标优化器

BAdam 把模型参数分块轮流训，每次只 update 一小块 → 显存峰值降到全参的 1/N。LLaMA-Factory 通过：

use_badam: true
badam_mode: layer            # layer / ratio
badam_switch_mode: ascending # ascending / descending / random / fixed
badam_switch_interval: 50    # 多少 step 换一块
badam_update_ratio: 0.05     # ratio 模式下每次训占比

BAdam 的特别处：trainable 参数始终保持半精度（adapter.py:319），不做 fp32 upcast。所以 BAdam 全参训能在不上 LoRA 的情况下跑出接近 LoRA 的显存占用。

10.12 KTransformers：CPU-GPU 混合推理 / 训练

2025 引入的 KTransformers 是 CPU-GPU 混合内核库，特别适合显存远不够但内存够的场景（消费级 4090 + 大内存训 70B 等）。LLaMA-Factory 在 ModelArguments → KTransformersArguments mixin 给入口。开关：

use_ktransformers: true
ktransformers_*: ...

启用后 init_adapter 会走特殊的 KT-PEFT 路径（adapter.py:252–257）。约束：

只支持 LoRA，不支持 OFT / full / freeze；
需要 KTransformers 这个独立包装好；
和 DeepSpeed 互斥（launcher.py 检测后跳过 torchrun）。

10.13 Megatron-core 后端（2025-10）

千亿模型 LoRA / 全参 SFT 走 Megatron 路径，开关 use_mca: true（由 mcore_adapter 包驱动）。本质：

YAML 入口和原 LLaMA-Factory 一样（同一份 train.yaml）；
底层把 transformer 跑到 Megatron-core 的 TP / PP / VPP 切分上；
Trainer 用 McaSeq2SeqTrainingArguments（training_args.py:26–36）替代 HF Seq2SeqTrainingArguments。

触发条件：环境变量 USE_MCA=1 + YAML 写 strategy: megatron（在 actor/critic 配置里）。本书基线 commit 上还在 beta，生产前一定先在小模型跑通再上千亿。

← 上一章

多模态：LLaVA / Qwen2-VL / Qwen3-VL 微调

常见错误与调试