全参数微调：DeepSpeed & FSDP

何时需要全参数微调

大多数场景 LoRA/QLoRA 已经足够，但以下情况考虑全参数微调：

• 需要从根本上改变模型的"世界观"（语言、领域）
• 持续预训练（Continual Pretraining）在新语料上
• LoRA 效果达到瓶颈，评估指标不再提升
• 数据量 > 10 万条，充分利用大数据集
• 追求最高精度，不关心显存成本

DeepSpeed ZeRO：分割显存压力

DeepSpeed 配置

{
  "zero_optimization": {
    "stage": 2,
    "overlap_comm": true,
    "contiguous_gradients": true,
    "sub_group_size": 1e9,
    "reduce_bucket_size": "auto",
    "stage3_prefetch_bucket_size": "auto",
    "stage3_param_persistence_threshold": "auto",
    "gather_16bit_weights_on_model_save": true
  },
  "bf16": {"enabled": true},
  "gradient_clipping": 1.0,
  "train_micro_batch_size_per_gpu": 4,
  "gradient_accumulation_steps": 4
}

启动多 GPU 训练

# 单机 8 GPU 训练
deepspeed --num_gpus=8 train.py \
  --deepspeed ds_config.json \
  --model_name_or_path meta-llama/Meta-Llama-3-8B \
  --dataset_path data/train.jsonl \
  --output_dir ./output \
  --num_train_epochs 3

# 或使用 torchrun（FSDP）
torchrun --nproc_per_node=8 \
  --master_addr="localhost" \
  --master_port=29500 \
  train_fsdp.py

FSDP vs DeepSpeed 选择

全参数微调训练技巧

全参数微调与 LoRA 的显存对比

Flash Attention 2 的工作原理

# 安装 Flash Attention 2（需要 CUDA 11.6+，Ampere/Ada 架构 GPU）
pip install flash-attn --no-build-isolation

# 在 Transformers 中启用
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3-8B",
    attn_implementation="flash_attention_2",   # 关键参数
    torch_dtype=torch.bfloat16,
)

多机训练的环境配置

# 多机 DeepSpeed 训练（主节点 + 工作节点）
# 需要所有节点 SSH 免密登录，共享存储（NFS/NAS）

# 主节点上创建 hostfile（列出所有节点）
cat hostfile
worker-0 slots=8  # 8张A100
worker-1 slots=8

# 在主节点启动（DeepSpeed 自动 SSH 到工作节点）
deepspeed --hostfile hostfile \
  --num_gpus 8 \
  --num_nodes 2 \
  train.py \
  --deepspeed ds_z3_config.json \
  --model_name_or_path meta-llama/Meta-Llama-3-70B

# torchrun 多机版本
torchrun --nnodes=2 \
  --nproc_per_node=8 \
  --rdzv_id=100 \
  --rdzv_backend=c10d \
  --rdzv_endpoint="worker-0:29400" \
  train_fsdp.py

分布式训练的常见问题与调试

from transformers import AutoModelForCausalLM
import torch

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3-8B",
    attn_implementation="flash_attention_2",
    torch_dtype=torch.bfloat16,
)

# 开启 Gradient Checkpointing（大幅节省激活值显存）
model.gradient_checkpointing_enable(
    gradient_checkpointing_kwargs={"use_reentrant": False}  # 推荐 False，更稳定
)

# 注意：Gradient Checkpointing 与以下特性不兼容
# - 某些自定义注意力实现
# - model.generate() 时需要先 model.gradient_checkpointing_disable()
# - 速度下降约 20-30%，但激活值显存减少约 10 倍

# TrainingArguments 中的稳定性关键参数
from transformers import TrainingArguments

args = TrainingArguments(
    output_dir="./output",
    bf16=True,                        # 优先 BF16（数值范围比 FP16 更大）
    max_grad_norm=1.0,                 # 梯度裁剪阈值，防止梯度爆炸
    learning_rate=2e-5,                # 全参数微调用较小学习率
    warmup_ratio=0.03,                 # 前 3% 步 warmup，防止初始训练不稳定
    lr_scheduler_type="cosine",        # Cosine 衰减，优于线性
    weight_decay=0.01,                 # L2 正则化，防止过拟合
    gradient_checkpointing=True,
    gradient_accumulation_steps=4,
    per_device_train_batch_size=2,
    save_strategy="steps",
    save_steps=500,                    # 每 500 步保存 checkpoint
    save_total_limit=3,                # 只保留最近 3 个 checkpoint
)