DeepSpeed 微调 LLaMA-2完整步骤

用户2584859014845
3 阅读3分钟

以下是使用 DeepSpeed 微调 LLaMA-2 的完整步骤，涵盖环境配置、数据处理、训练脚本编写和部署优化，适用于单机多卡或多机分布式场景：

1. 环境准备

(1) 硬件要求

  • 最低配置：至少2张A100 80GB GPU（7B模型）或8张A100（70B模型）
  • 推荐配置：NVIDIA A100/H100集群 + NVLink互联

(2) 软件依赖

# 创建conda环境
conda create -n llama2 python=3.10 -y
conda activate llama2

# 安装核心库
pip install torch==2.0.1+cu117 --extra-index-url https://download.pytorch.org/whl/cu117
pip install transformers==4.31.0 datasets accelerate deepspeed sentencepiece

(3) 获取LLaMA-2访问权限

  1. 申请Meta官方许可：ai.meta.com/resources/m…
  2. 下载模型权重（选择7B/13B/70B版本）： 
    huggingface-cli login  # 使用HF账号认证
huggingface-cli download meta-llama/Llama-2-7b-hf --local-dir ./llama-2-7b

2. 数据准备

(1) 数据集格式

  • 使用JSON格式，每条数据包含instruction和output： 
    [
  {"instruction": "解释量子力学", "output": "量子力学是研究微观粒子运动规律的物理学分支..."},
  {"instruction": "写一首关于春天的诗", "output": "春风拂面百花开..."}
]
  • 保存为train.jsonl和eval.jsonl

(2) 数据加载脚本

from datasets import load_dataset
dataset = load_dataset("json", data_files={"train": "train.jsonl", "eval": "eval.jsonl"})

3. 配置DeepSpeed

(1) 创建DeepSpeed配置文件ds_config.json

{
  "train_micro_batch_size_per_gpu": 4,
  "gradient_accumulation_steps": 8,
  "optimizer": {
    "type": "AdamW",
    "params": {
      "lr": 2e-5,
      "weight_decay": 0.01
    }
  },
  "fp16": {
    "enabled": true,
    "loss_scale_window": 100
  },
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": {
      "device": "cpu"
    },
    "allgather_partitions": true,
    "allgather_bucket_size": 5e8
  },
  "steps_per_print": 50
}

(2) 关键参数说明

  • zero_optimization.stage=3：启用ZeRO-3优化，显存占用最小化
  • offload_optimizer：将优化器状态卸载到CPU，节省GPU显存
  • train_micro_batch_size_per_gpu：根据GPU显存调整（7B模型建议2-4）

4. 微调脚本train.py

import torch
from transformers import (
    LlamaForCausalLM,
    LlamaTokenizer,
    TrainingArguments
)
from trl import SFTTrainer

# 1. 加载模型和分词器
model = LlamaForCausalLM.from_pretrained(
    "./llama-2-7b",
    torch_dtype=torch.float16,
    device_map="auto"
)
tokenizer = LlamaTokenizer.from_pretrained("./llama-2-7b")
tokenizer.pad_token = tokenizer.eos_token

# 2. 定义训练参数
training_args = TrainingArguments(
    output_dir="./output",
    per_device_train_batch_size=4,
    num_train_epochs=3,
    logging_steps=10,
    save_steps=500,
    fp16=True,
    deepspeed="./ds_config.json",  # 关键：启用DeepSpeed
)

# 3. 创建Trainer
trainer = SFTTrainer(
    model=model,
    args=training_args,
    train_dataset=dataset["train"],
    dataset_text_field="instruction",  # 指定文本字段
    max_seq_length=1024,
    tokenizer=tokenizer
)

# 4. 开始训练
trainer.train()

5. 启动训练

(1) 单机多卡（4卡为例）

deepspeed --num_gpus=4 train.py

(2) 多机训练

# 主机（rank=0）
deepspeed --hostfile=hostfile --master_addr=主节点IP --master_port=29500 train.py

# hostfile内容示例
node1 slots=4  # 第一台机器4卡
node2 slots=4  # 第二台机器4卡

6. 监控与优化

(1) 显存监控

watch -n 1 nvidia-smi  # 实时查看GPU利用率

(2) 训练日志解读

  • 关注loss下降曲线和grad_norm（梯度裁剪阈值）
  • 如果出现OOM（内存不足）：
    • 减小per_device_train_batch_size
    • 增加gradient_accumulation_steps

7. 模型导出与部署

(1) 保存微调后的模型

trainer.save_model("./llama-2-7b-finetuned")

(2) 转换为推理优化格式

# 导出为vLLM兼容格式（如需高性能API）
python -m vllm.entrypoints.convert_model ./llama-2-7b-finetuned --output-dir ./llama-2-7b-vllm

(3) 启动推理服务

python -m vllm.entrypoints.api_server --model ./llama-2-7b-vllm --tensor-parallel-size 4

常见问题解决

  1. CUDA内存不足：

    • 在ds_config.json中启用"offload_param": {"device": "cpu"}
    • 使用梯度检查点：model.gradient_checkpointing_enable()

  2. 多机通信失败：

    • 检查防火墙设置sudo ufw allow 29500
    • 确保所有节点SSH免密登录

  3. 微调效果不佳：

    • 尝试调整学习率（1e-5到5e-5）
    • 增加数据量或使用数据增强

性能参考（7B模型）

硬件Batch Size显存占用Tokens/sec
1x A100 80GB472GB1200
4x A100 80GB164×42GB4800

如果需要针对特定场景（如医疗文本微调）的完整代码库，可以参考：

avatar
文章
阅读
粉丝