多 GPU 微调技巧：DeepSpeed 与 FSDP 实践

分布式微调的核心挑战

当我们从单卡微调转向多卡并行时，显存瓶颈与通信效率成为两大拦路虎。单个 GPU 无法装下完整模型、优化器状态和批次数据，必须将模型切分到多个设备上。此时，你需要回答三个关键问题：如何切分模型？如何同步梯度？如何最小化冗余显存占用？DeepSpeed 和 FSDP 正是为此而生，它们分别代表了微软与 Meta 给出的工业级方案。

本教程将带你从零开始，理解两种技术的核心原理，并分别完成一次 7B 参数模型的多 GPU 微调实践。你不需要成为分布式系统专家，但需要具备 PyTorch 基础以及 Hugging Face Transformers 的简单使用经验。

ZeRO 策略速览：DeepSpeed 的分层哲学

DeepSpeed 是微软推出的分布式训练引擎，其中最核心的内存优化技术称为 ZeRO（零冗余优化器）。ZeRO 将训练状态划分为三个阶段，从保守到激进逐步削减冗余：

• ZeRO-1：优化器状态（optimizer states）分片。每个 GPU 只保存一部分优化器状态，更新时通过 AllGather 还原。
• ZeRO-2：在 ZeRO-1 基础上，对梯度（gradients）也进行分片。反向传播后，梯度先分片再 reduce，进一步降低内存。
• ZeRO-3：参数量（parameters）也分片。模型参数不再在每个 GPU 上完整保存，前向/反向传播时按需通过 AllGather 获取。存储需求随 GPU 数量近线性下降。

几乎所有大模型微调场景都建议直接使用 ZeRO-2（兼顾速度与内存）或 ZeRO-3（极致省内存，适合超大模型）。此外，DeepSpeed 还提供 CPU Offload 和 NVMe Offload 技术，允许将部分状态交换到内存或硬盘，以通信换内存。

FSDP 设计理念：PyTorch 原生分片

FSDP（Fully Sharded Data Parallel）是 PyTorch 官方在 1.11 版本后加入的分布式训练 API。它的核心思想与 DeepSpeed ZeRO-3 高度相似：将模型参数、梯度和优化器状态均分到所有 GPU 上，只有计算时需要时才通过通信还原。但在实现上更加原生，与 PyTorch 生态无缝结合，避免了额外依赖。

需要区分的是，简单的 DistributedDataParallel (DDP) 只是将梯度同步，模型副本在每张卡上完全一致；而 FSDP 把参数也分开了。下表可以帮你快速对比：

环境准备：安装与集群基础

本教程假设你有一个配备至少 2 张 GPU（NVIDIA A10/24G 以上）的节点，系统为 Linux，Python 3.10+，已安装 CUDA 驱动。

1. 安装 PyTorch 与 FSDP

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

（FSDP 已集成在 PyTorch 中，无需额外安装。）

2. 安装 DeepSpeed

pip install deepspeed

3. 安装 Hugging Face 生态

pip install transformers accelerate datasets peft bitsandbytes

4. 验证 GPU 可见性

import torch
print(torch.cuda.device_count())  # 应输出 >=2

实践一：使用 DeepSpeed ZeRO-3 微调 Llama-2-7B

我们将基于 Hugging Face transformers 和 accelerate 进行 LoRA 微调，以减少可训练参数，聚焦于分布式内存优化。

步骤 1：准备 DeepSpeed 配置文件

创建 ds_config.json：

{
  "train_batch_size": "auto",
  "train_micro_batch_size_per_gpu": 1,
  "gradient_accumulation_steps": 8,
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    },
    "offload_param": {
      "device": "cpu",
      "pin_memory": true
    },
    "overlap_comm": true,
    "contiguous_gradients": true,
    "reduce_bucket_size": 5e8,
    "stage3_prefetch_bucket_size": 5e8,
    "stage3_param_persistence_threshold": 1e6
  },
  "fp16": {
    "enabled": true
  },
  "optimizer": {
    "type": "AdamW",
    "params": {
      "lr": 2e-4,
      "weight_decay": 0.01
    }
  },
  "scheduler": {
    "type": "WarmupDecayLR",
    "params": {
      "warmup_min_lr": 0,
      "warmup_max_lr": 2e-4,
      "warmup_num_steps": 100,
      "total_num_steps": 1000
    }
  }
}

• stage: 3 启用参数/梯度/优化器完整分片。
• offload_optimizer 和 offload_param 将优化器与参数卸载到 CPU 内存，可进一步降低 GPU 显存。
• 批次大小由 auto 根据 micro batch 大小和梯度累积步数计算。

步骤 2：编写微调脚本

创建一个名为 train_deepspeed.py 的文件：

import torch
import deepspeed
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from datasets import load_dataset

model_name = "meta-llama/Llama-2-7b-hf"
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token

# 4-bit 量化基础模型以适配显存（可选）
from transformers import BitsAndBytesConfig
quant_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True,
)

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=quant_config,
    device_map={"": 0},  # 临时指定主设备
    torch_dtype=torch.float16,
)
model = prepare_model_for_kbit_training(model)

# 配置 LoRA
lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    target_modules=["q_proj", "v_proj"],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)
model = get_peft_model(model, lora_config)

# 加载数据集
dataset = load_dataset("timdettmers/openassistant-guanaco", split="train")
def tokenize(example):
    return tokenizer(example["text"], truncation=True, max_length=512)
dataset = dataset.map(tokenize, batched=True)

# 训练参数
training_args = TrainingArguments(
    output_dir="./llama2-7b-lora-ds",
    per_device_train_batch_size=1,
    gradient_accumulation_steps=8,
    logging_steps=10,
    save_steps=200,
    num_train_epochs=1,
    learning_rate=2e-4,
    fp16=True,
    deepspeed="./ds_config.json",  # 指向配置文件
    report_to="none",
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    tokenizer=tokenizer,
)
trainer.train()

步骤 3：启动多 GPU 训练

不需要使用 torchrun，直接通过 deepspeed 命令启动：

deepspeed --num_gpus=2 train_deepspeed.py

训练过程中可通过 nvidia-smi 观察，显存占用将显著低于普通 DDP。注意：4-bit 量化 + ZeRO-3 的组合可能引起冲突，如果报错可以将 BitsAndBytesConfig 注释掉，直接使用半精度模型。

实践二：使用 FSDP 微调 Mistral-7B

FSDP 通过 torch.distributed.fsdp API 直接控制模型分片。我们将结合 Hugging Face accelerate 库简化配置，同样使用 LoRA。

步骤 1：配置 accelerate

运行 accelerate config 交互式生成配置文件，关键选择：

• 分布式类型: FSDP
• 分片策略: FULL_SHARD（对应 ZeRO-3）
• 自动封装策略: TRANSFORMER_BASED_WRAP
• 精度: fp16 或 bf16

这会在当前目录生成 default_config.yaml。

步骤 2：编写 FSDP 训练脚本 train_fsdp.py

import torch
from accelerate import Accelerator
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import LoraConfig, get_peft_model
from datasets import load_dataset
from torch.utils.data import DataLoader

model_name = "mistralai/Mistral-7B-v0.1"
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token

# 加载模型（不指定 device_map，交由 FSDP 接管）
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    use_cache=False,  # 梯度检查点与 FSDP 搭配时建议关闭
)

lora_config = LoraConfig(
    r=8,
    lora_alpha=16,
    target_modules=["q_proj", "v_proj"],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)
model = get_peft_model(model, lora_config)

dataset = load_dataset("timdettmers/openassistant-guanaco", split="train")
def tokenize(example):
    return tokenizer(example["text"], truncation=True, max_length=512, padding="max_length")
dataset = dataset.map(tokenize, batched=True)
dataset.set_format(type="torch", columns=["input_ids", "attention_mask"])

dataloader = DataLoader(dataset, batch_size=1, shuffle=True)

optimizer = torch.optim.AdamW(model.parameters(), lr=2e-4)

# Accelerator 初始化 （会自动加载 FSDP 配置）
accelerator = Accelerator()
model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader)

model.train()
for step, batch in enumerate(dataloader):
    outputs = model(**batch)
    loss = outputs.loss
    accelerator.backward(loss)
    optimizer.step()
    optimizer.zero_grad()
    if step % 10 == 0:
        accelerator.print(f"Step {step} loss: {loss.item()}")

步骤 3：启动训练

accelerate launch --config_file default_config.yaml train_fsdp.py

默认使用所有可见 GPU。如果希望手动指定数量，在启动前设置 CUDA_VISIBLE_DEVICES=0,1。

调优技巧与常见陷阱

1. 批次大小与通信重叠

• 每张卡上的 per_device_train_batch_size 尽量设为 1，通过梯度累积达到等效大 batch。
• DeepSpeed 开启 overlap_comm 可将通信与反向传播重叠；FSDP 中设置 fsdp_config 的 forward_prefetch 和 backward_prefetch 有类似效果。

2. CPU Offload 的明智使用

• 当显存仍然吃紧时，优先 offload 优化器状态（ZeRO-2/3 均支持）。参数 offload 会频繁触发 GPU-CPU 拷贝，显著拖慢训练速度，仅当你真的无法容纳参数时使用。
• FSDP 通过 offload_params 策略实现类似功能。

3. 混合精度

• 始终启用 fp16 或 bf16。bf16 更稳定但需要 Ampere 以上架构。
• 配合梯度缩放（GradScaler）以避免下溢；在 DeepSpeed 中自动管理，FSDP 需通过 accelerate 或手动设置。

4. LoRA 与分片叠加

• 使用 PEFT 时，只有基础模型的参数分片，LoRA 适配器仍然在每个 GPU 上保存完整副本（参数量极少，影响不大）。注意不要将 LoRA 层也纳入分片范围。

5. 监控与 profiling

• nvidia-smi dmon -s pucvmet 可以实时查看显存与利用率。
• DeepSpeed 内置 ds_report 可生成内存分析报告；FSDP 可通过 PyTorch Profiler 配合 TensorBoard 分析通信开销。

6. 常见报错

• RuntimeError: Expected all tensors to be on the same device：可能某个模块未被 FSDP 包裹。确保 model 最初没有指定 device_map，由 accelerator 统一处理。
• OOM：减少 micro batch size，增加梯度累积，或启用 CPU offload；考虑使用 4-bit 量化基础模型。

选型指南：何时用 DeepSpeed，何时用 FSDP？

• 团队熟悉 PyTorch 原生工具，想避免额外依赖：选 FSDP。经过 PyTorch 2.x 的迭代，FSDP 性能已非常接近 DeepSpeed，且与 torch.compile 等新特性兼容更好。
• 需要极致的显存节省和灵活的 offload 方案：选 DeepSpeed。其 ZeRO-Infinity 可将状态 offload 到 NVMe，支持训练比 GPU 总显存大十倍的模型。
• 实验阶段快速跑通小规模模型：DDP 即可，无需引入复杂性。
• 混合使用不同厂商的 GPU（如部分节点用 AMD）：DeepSpeed 支持 ROCm，而 FSDP 在 PyTorch 层面的可移植性更好，需按实际驱动支持选择。

不论选择哪一种，核心思想都是用通信换内存，因此网络带宽（NVLink/NVSwitch/Infiniband）会极大影响训练吞吐。在单机多卡环境下，NVLink 连接的两张卡通常能获得近线性的加速比。

总结

多 GPU 微调不再是只能由大型实验室操作的“黑科技”。掌握 DeepSpeed 的 ZeRO 配置和 FSDP 的 accelerate 工作流后，你就能在消费级或企业级多卡服务器上从容加载并微调 7B-13B 的模型。建议先在单卡上调试通代码逻辑，再扩展到多卡，最后针对显存和吞吐调优配置。动手试试吧，你的第一份分布式微调 Checkpoint 就在前方。