从零复现DeepSpeed Chat三阶段RLHF训练全流程：实战代码解析与避坑指南

如果你已经对RLHF（基于人类反馈的强化学习）的理论框架有所了解，却苦于找不到清晰可操作的代码实现方案，那么这篇手把手教程正是为你准备的。我们将基于微软开源的DeepSpeed Chat项目，完整复现从监督微调（SFT）到奖励模型训练（RM），再到PPO强化学习的三阶段训练流程。不同于理论讲解，本文聚焦于 可落地的工程实践 ，包含完整的代码解析、参数配置技巧和常见报错解决方案。

1. 环境准备与项目配置

在开始之前，我们需要搭建适合大规模模型训练的基础环境。推荐使用Linux系统（如Ubuntu 20.04+）并配备NVIDIA显卡（建议A100 40GB以上）。

# 创建Python虚拟环境（推荐3.8-3.10版本）
conda create -n ds_chat python=3.9 -y
conda activate ds_chat

# 安装PyTorch与CUDA工具包
pip install torch==2.0.1+cu118 torchvision==0.15.2+cu118 --extra-index-url https://download.pytorch.org/whl/cu118

# 安装DeepSpeed及其依赖
pip install deepspeed==0.10.0
pip install transformers==4.31.0 datasets==2.13.1 accelerate==0.21.0

克隆DeepSpeed Chat仓库并准备示例数据集：

git clone https://github.com/microsoft/DeepSpeedExamples.git
cd DeepSpeedExamples/applications/DeepSpeed-Chat

# 下载示例数据集（约5GB）
wget https://huggingface.co/datasets/Dahoas/rm-static/resolve/main/data/train.jsonl
wget https://huggingface.co/datasets/Dahoas/rm-static/resolve/main/data/test.jsonl
mkdir -p data/rm-static
mv train.jsonl test.jsonl data/rm-static/

关键目录结构说明：

DeepSpeed-Chat/
├── training/
│   ├── step1_supervised_finetuning/    # SFT阶段代码
│   ├── step2_reward_model_finetuning/  # RM阶段代码
│   └── step3_rlhf_finetuning/          # PPO阶段代码
├── data/                               # 数据集存放位置
└── output/                             # 模型输出目录

2. 第一阶段：监督微调（SFT）

SFT阶段的目标是让预训练语言模型初步掌握指令遵循能力。我们使用高质量的人类标注对话数据进行微调。

2.1 数据准备与格式处理

DeepSpeed Chat要求训练数据为JSONL格式，每条数据包含"prompt"和"chosen"字段：

{
  "prompt": "解释量子计算的基本概念",
  "chosen": "量子计算是利用量子力学原理（如叠加和纠缠）进行信息处理的新型计算模式..."
}

创建配置文件 training/step1_supervised_finetuning/params.json ：

{
  "model_name_or_path": "facebook/opt-1.3b",
  "data_path": "Dahoas/rm-static",
  "per_device_train_batch_size": 8,
  "per_device_eval_batch_size": 8,
  "max_seq_len": 512,
  "learning_rate": 1e-5,
  "weight_decay": 0.1,
  "num_train_epochs": 1,
  "gradient_accumulation_steps": 8,
  "lr_scheduler_type": "cosine",
  "num_warmup_steps": 100,
  "seed": 42,
  "gradient_checkpointing": true,
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": {
      "device": "cpu"
    }
  },
  "fp16": {
    "enabled": true
  }
}

2.2 启动SFT训练

运行以下命令开始训练：

deepspeed --num_gpus=4 training/step1_supervised_finetuning/main.py \
  --data_path Dahoas/rm-static \
  --model_name_or_path facebook/opt-1.3b \
  --per_device_train_batch_size 8 \
  --per_device_eval_batch_size 8 \
  --max_seq_len 512 \
  --learning_rate 1e-5 \
  --weight_decay 0.1 \
  --num_train_epochs 1 \
  --gradient_accumulation_steps 8 \
  --lr_scheduler_type cosine \
  --num_warmup_steps 100 \
  --seed 42 \
  --gradient_checkpointing \
  --output_dir output/sft \
  --deepspeed training/step1_supervised_finetuning/ds_config.json

常见问题排查：

1. CUDA内存不足 ：减小 per_device_batch_size 或增加 gradient_accumulation_steps
2. NaN损失值 ：尝试降低学习率或启用梯度裁剪
3. 数据集加载失败 ：检查 data_path 是否指向正确的数据集目录

2.3 关键代码解析

SFT阶段的核心训练逻辑在 main.py 中实现：

def main():
    # 初始化tokenizer和模型
    tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path)
    model = create_hf_model(AutoModelForCausalLM, args.model_name_or_path, tokenizer)
    
    # 准备数据集
    train_dataset, eval_dataset = create_prompt_dataset(
        args.data_path, args.train_split, args.valid_split, ...)
    
    # DeepSpeed引擎初始化
    model_engine, optimizer, _, _ = deepspeed.initialize(
        model=model,
        config=ds_config,
        ...)
    
    # 训练循环
    for epoch in range(args.num_train_epochs):
        for step, batch in enumerate(train_dataloader):
            outputs = model_engine(**batch)
            loss = outputs.loss
            model_engine.backward(loss)
            model_engine.step()
            
            # 定期评估
            if step % args.eval_steps == 0:
                eval_loss = evaluate(model_engine, eval_dataloader)

3. 第二阶段：奖励模型训练（RM）

RM阶段的目标是训练一个能够评估回复质量的评分模型，为后续RLHF提供反馈信号。

3.1 数据格式要求

RM训练需要成对的偏好数据，每条数据包含：

• prompt：输入的提示
• chosen：人类偏好的回复
• rejected：人类不喜欢的回复

示例数据格式：

{
  "prompt": "写一首关于春天的诗",
  "chosen": "春风拂面百花开，燕子归来筑巢忙...",
  "rejected": "我不知道怎么写诗，你可以自己试试"
}

3.2 启动RM训练

配置文件 training/step2_reward_model_finetuning/params.json ：

{
  "model_name_or_path": "facebook/opt-350m",
  "data_path": "Dahoas/rm-static",
  "num_padding_at_beginning": 0,
  "per_device_train_batch_size": 8,
  "per_device_eval_batch_size": 8,
  "max_seq_len": 512,
  "learning_rate": 5e-6,
  "weight_decay": 0.01,
  "num_train_epochs": 1,
  "gradient_accumulation_steps": 8,
  "lr_scheduler_type": "cosine",
  "num_warmup_steps": 100,
  "seed": 42,
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": {
      "device": "cpu"
    }
  }
}

启动训练命令：

deepspeed --num_gpus=4 training/step2_reward_model_finetuning/main.py \
  --data_path Dahoas/rm-static \
  --model_name_or_path facebook/opt-350m \
  --num_padding_at_beginning 0 \
  --per_device_train_batch_size 8 \
  --per_device_eval_batch_size 8 \
  --max_seq_len 512 \
  --learning_rate 5e-6 \
  --weight_decay 0.01 \
  --num_train_epochs 1 \
  --gradient_accumulation_steps 8 \
  --lr_scheduler_type cosine \
  --num_warmup_steps 100 \
  --seed 42 \
  --output_dir output/rm \
  --deepspeed training/step2_reward_model_finetuning/ds_config.json

3.3 核心实现解析

RM模型的关键在于成对排序损失（Pairwise Ranking Loss）的实现：

class RewardModel(nn.Module):
    def __init__(self, base_model):
        super().__init__()
        self.rwtransformer = base_model
        self.v_head = nn.Linear(base_model.config.hidden_size, 1)
    
    def forward(self, input_ids, attention_mask):
        # 获取transformer输出
        transformer_outputs = self.rwtransformer(
            input_ids=input_ids,
            attention_mask=attention_mask)
        
        # 计算奖励分数
        hidden_states = transformer_outputs.last_hidden_state
        rewards = self.v_head(hidden_states).squeeze(-1)
        
        # 分割chosen和rejected部分
        bs = input_ids.shape[0] // 2
        chosen_rewards = rewards[:bs]
        rejected_rewards = rewards[bs:]
        
        # 计算成对排序损失
        loss = -torch.log(
            torch.sigmoid(chosen_rewards - rejected_rewards)
        ).mean()
        
        return {
            "loss": loss,
            "chosen_rewards": chosen_rewards,
            "rejected_rewards": rejected_rewards
        }

4. 第三阶段：RLHF与PPO训练

这是最复杂的阶段，我们将使用PPO算法结合前面训练的SFT和RM模型进行强化学习微调。

4.1 配置文件准备

创建 training/step3_rlhf_finetuning/params.json ：

{
  "actor_model_name_or_path": "output/sft",
  "critic_model_name_or_path": "output/rm",
  "num_train_epochs": 1,
  "per_device_train_batch_size": 8,
  "per_device_mini_batch_size": 8,
  "generation_batch_numbers": 1,
  "ppo_epochs": 1,
  "max_answer_seq_len": 256,
  "max_prompt_seq_len": 256,
  "actor_learning_rate": 9.65e-6,
  "critic_learning_rate": 5e-6,
  "actor_weight_decay": 0.1,
  "critic_weight_decay": 0.1,
  "enable_hybrid_engine": true,
  "inference_tp_size": 1,
  "tp_gather_partition_size": 8,
  "zero_stage": 3,
  "actor_gradient_checkpointing": true,
  "critic_gradient_checkpointing": true
}

4.2 启动PPO训练

deepspeed --num_gpus=8 training/step3_rlhf_finetuning/main.py \
  --actor_model_name_or_path output/sft \
  --critic_model_name_or_path output/rm \
  --num_train_epochs 1 \
  --per_device_train_batch_size 8 \
  --per_device_mini_batch_size 8 \
  --generation_batch_numbers 1 \
  --ppo_epochs 1 \
  --max_answer_seq_len 256 \
  --max_prompt_seq_len 256 \
  --actor_learning_rate 9.65e-6 \
  --critic_learning_rate 5e-6 \
  --actor_weight_decay 0.1 \
  --critic_weight_decay 0.1 \
  --enable_hybrid_engine \
  --inference_tp_size 1 \
  --tp_gather_partition_size 8 \
  --zero_stage 3 \
  --actor_gradient_checkpointing \
  --critic_gradient_checkpointing \
  --output_dir output/ppo \
  --deepspeed training/step3_rlhf_finetuning/ds_config.json

4.3 PPO核心算法实现

DeepSpeed Chat中的PPO实现主要包含以下几个关键组件：

class PPOTrainer:
    def __init__(self, actor, critic, ref_model, reward_model, ...):
        self.actor = actor      # 待优化的策略模型
        self.critic = critic    # 价值函数模型
        self.ref_model = ref_model  # SFT模型（用于KL散度计算）
        self.reward_model = reward_model  # RM模型
    
    def generate_experience(self, prompts):
        # 使用当前策略生成回复
        seq = self.actor.generate(prompts)
        
        # 计算各组件输出
        with torch.no_grad():
            actor_logits = self.actor(seq).logits
            ref_logits = self.ref_model(seq).logits
            rewards = self.reward_model.forward_value(seq)['chosen_end_scores']
            values = self.critic.forward_value(seq, return_value_only=True)
        
        return {
            'prompts': prompts,
            'logprobs': gather_log_probs(actor_logits, seq),
            'ref_logprobs': gather_log_probs(ref_logits, seq),
            'values': values,
            'rewards': rewards,
            'input_ids': seq
        }
    
    def train_rlhf(self, experience):
        # 计算KL惩罚后的奖励
        rewards = self.compute_rewards(
            experience['logprobs'],
            experience['ref_logprobs'],
            experience['rewards'])
        
        # 计算优势函数和回报
        advantages, returns = self.get_advantages_and_returns(
            experience['values'],
            rewards)
        
        # 策略损失（带裁剪）
        actor_loss = self.actor_loss_fn(
            new_logprobs,
            experience['logprobs'],
            advantages)
        
        # 价值函数损失
        critic_loss = self.critic_loss_fn(
            new_values,
            experience['values'],
            returns)
        
        return actor_loss, critic_loss

5. 实战技巧与性能优化

5.1 混合引擎加速

DeepSpeed的Hybrid Engine可以显著提升RLHF训练效率：

# 在DeepSpeed配置中启用混合引擎
{
  "train_micro_batch_size_per_gpu": 8,
  "hybrid_engine": {
    "enabled": true,
    "max_out_tokens": 512,
    "inference_tp_size": 1,
    "release_inference_cache": false,
    "pin_parameters": true,
    "tp_gather_partition_size": 8
  }
}

5.2 内存优化策略

技术	配置方法	内存节省	计算开销
梯度检查点	--gradient_checkpointing	30-50%	增加25%
ZeRO-3	"zero_optimization": {"stage": 3}	线性减少	通信开销
CPU Offload	"offload_optimizer": {"device": "cpu"}	额外20%	数据传输开销
FP16混合精度	"fp16": {"enabled": true}	50%	可忽略

5.3 自定义数据集支持

要实现自定义数据集，需要继承 PromptRawDataset 类：

class CustomDataset(PromptRawDataset):
    def __init__(self, output_path, seed, local_rank, dataset_name):
        super().__init__(output_path, seed, local_rank, dataset_name)
        self.dataset = load_dataset("json", data_files="custom_data.jsonl")["train"]
    
    def get_prompt(self, sample):
        return f"Human: {sample['instruction']}\nAssistant:"
    
    def get_chosen(self, sample):
        return sample["response"]
    
    def get_rejected(self, sample):
        return sample["bad_response"]

然后在 data_utils.py 中注册数据集：

def get_raw_dataset(dataset_name, output_path, seed, local_rank):
    if "custom" in dataset_name:
        return CustomDataset(output_path, seed, local_rank, dataset_name)
    ...

6. 模型评估与部署

训练完成后，可以使用以下脚本测试模型效果：

from transformers import pipeline

chatbot = pipeline("text-generation", 
                  model="output/ppo",
                  device=0)

while True:
    prompt = input("User: ")
    if prompt.lower() == "exit":
        break
        
    response = chatbot(
        f"Human: {prompt}\nAssistant:",
        max_length=256,
        do_sample=True,
        temperature=0.7,
        top_p=0.9)
    
    print(f"Assistant: {response[0]['generated_text']}")

对于生产环境部署，建议使用DeepSpeed Inference：

import deepspeed

model = deepspeed.init_inference(
    model=AutoModelForCausalLM.from_pretrained("output/ppo"),
    mp_size=2,
    dtype=torch.float16,
    replace_method="auto"
)

7. 常见问题解决方案

问题1 ：训练过程中出现NaN损失

• 检查学习率是否过高
• 尝试添加梯度裁剪（ --max_grad_norm 1.0 ）
• 验证输入数据是否包含异常值

问题2 ：GPU内存不足

• 减小batch size
• 增加gradient_accumulation_steps
• 启用ZeRO-3和CPU offload

问题3 ：生成结果质量差

• 检查RM模型是否训练充分
• 调整PPO阶段的KL惩罚系数（默认0.02）
• 增加SFT阶段的数据量和训练时间

问题4 ：训练速度慢

• 启用Hybrid Engine
• 使用更高效的GPU（如A100/H100）
• 优化数据加载流程（预加载数据到内存）

通过本教程，你应该已经掌握了使用DeepSpeed Chat实现完整RLHF三阶段训练的全流程。实际应用中可能需要根据具体任务调整超参数和数据格式，但核心框架保持不变。建议从小规模实验开始，逐步扩大训练规模，同时密切关注训练动态和模型表现。