ChatGLM2微调实战：基于DeepSpeed的高效训练方案解析

1次阅读

没有评论

共计 2666 个字符，预计需要花费 7 分钟才能阅读完成。

最近在尝试微调 ChatGLM2 模型时，遇到了显存爆炸、训练龟速的问题。经过一番折腾，发现 DeepSpeed 简直是救星！今天就把这套实战经验整理出来，包含从环境配置到性能调优的全流程，希望能帮到同样被大模型微调折磨的小伙伴们。

微调大模型时，我们常遇到三个头疼问题：

显存黑洞：加载 13B 参数的 ChatGLM2，光模型权重就吃掉 24GB 显存，加上梯度、优化器状态，消费级显卡直接 OOM
训练龟速：单卡跑满 batch_size= 1 都要 10 秒 /step，百万级数据等到天荒地老
多卡闲置：传统 DataParallel 下，多卡显存利用率不到 50%

对比原生 PyTorch，DeepSpeed 的三大杀手锏：

ZeRO 优化：通过分片优化器状态(Stage1)、梯度(Stage2)、参数(Stage3)，显存占用呈数量级下降
CPU Offload：将暂时不用的数据卸载到内存，显卡专注计算
梯度检查点：用时间换空间，仅保留关键节点的激活值

实测在 RTX 3090 上：

方案	最大 batch_size	显存占用	速度(iter/s)
PyTorch 原生	1	38GB	0.1
DeepSpeed-ZeRO2	8	22GB	1.8

pip install deepspeed transformers==4.33.0 torch==2.0.1

{
  "train_batch_size": 32,
  "gradient_accumulation_steps": 4,
  "optimizer": {
    "type": "AdamW",
    "params": {
      "lr": 2e-5,
      "weight_decay": 0.01
    }
  },
  "fp16": {
    "enabled": true,
    "loss_scale_window": 100
  },
  "zero_optimization": {
    "stage": 2,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    },
    "allgather_partitions": true,
    "allgather_bucket_size": 5e8,
    "overlap_comm": true,
    "reduce_scatter": true,
    "reduce_bucket_size": 5e8
  },
  "gradient_clipping": 1.0,
  "steps_per_print": 50
}

Stage 选择：个人推荐 Stage2 平衡显存与速度，13B 模型显存下降 60%
Offload 策略：内存足够时建议开启 CPU 卸载，可再省 20% 显存
通信优化：allgather_bucket_size 调大有助于提升多卡效率

import deepspeed
from transformers import AutoTokenizer, AutoModel

# 初始化模型
tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True)
model = AutoModel.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True)

# DeepSpeed 初始化
model_engine, optimizer, _, _ = deepspeed.initialize(
    model=model,
    model_parameters=model.parameters(),
    config_params="ds_config.json"
)

# 训练循环
for batch in dataloader:
    inputs = tokenizer(batch["text"], padding=True, return_tensors="pt").to(model_engine.device)
    outputs = model_engine(**inputs, labels=inputs["input_ids"])
    model_engine.backward(outputs.loss)
    model_engine.step()

# ds_config.json 中设置
"gradient_accumulation_steps": 4,
"train_micro_batch_size_per_gpu": 8

# 等效 batch_size = 8 * 4 * GPU 数量

当显存不足时，增大 gradient_accumulation_steps
通讯开销大时，适当增加 micro_batch_size

{
  "fp16": {
    "enabled": true,
    "loss_scale": 0,
    "initial_scale_power": 16,
    "loss_scale_window": 1000,
    "hysteresis": 2,
    "min_loss_scale": 1
  }
}