Llama-3 接入 verl 框架的强化学习训练实践

同时，在数据处理部分启用 chat template 自动填充：

data:
  chat_template: "llama-3"
  prompt_key: "prompt"
  max_prompt_length: 512
  max_response_length: 1024

若 verl 不原生支持 "llama-3" 模板，可在 verl/data/utils.py 中添加如下注册逻辑：

from transformers import AutoTokenizer

def get_chat_template(tokenizer_name_or_path):
    if "llama-3" in tokenizer_name_or_path.lower():
        return {
            "chat_template": "{% for message in messages %}{{'<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n' + message['content'] + '<|eot_id|>'}}{% endfor %}{% if add_generation_prompt %}{{ '<|start_header_id|>assistant<|end_header_id|>\n\n' }}{% endif %}"
        }
    return None

然后在数据加载器中动态注入该模板。

3.2 配置 GRPO 训练脚本（以 Llama-3-8B 为例）

创建 grpo_llama3.yaml 配置文件，内容如下：

data:
  train_files: ~/data/gsm8k/train.parquet
  val_files: ~/data/gsm8k/test.parquet
  prompt_key: prompt
  max_prompt_length: 512
  max_response_length: 1024
  train_batch_size: 1024
  val_batch_size: 1312
  shuffle: true
actor_rollout_ref:
  hybrid_engine: true
  model:
    path: meta-llama/Meta-Llama-3-8B-Instruct
    trust_remote_code: true
    enable_gradient_checkpointing: true
    use_remove_padding: false
  actor:
    strategy: fsdp
    ppo_mini_batch_size: 256
    ppo_micro_batch_size_per_gpu: 1
    ppo_max_token_len_per_gpu: 16384
    grad_clip: 1.0
    clip_ratio: 0.2
    entropy_coeff: 0.001
    use_kl_loss: true
    kl_loss_coef: 0.001
    kl_loss_type: low_var_kl
    ppo_epochs: 1
    optim:
      lr: 5e-7
      lr_warmup_steps_ratio: 0.1
      warmup_style: cosine
    fsdp_config:
      wrap_policy:
        min_num_params: 0
      param_offload: false
      optimizer_offload: false
  ref:
    fsdp_config:
      param_offload: false
      wrap_policy:
        min_num_params: 0
  rollout:
    name: vllm
    temperature: 0.7
    top_p: 0.95
    dtype: bfloat16
    gpu_memory_utilization: 0.8
    enforce_eager: true
    free_cache_engine: true
    load_format: dummy_dtensor
    tensor_model_parallel_size: 2
    max_num_batched_tokens: 8192
    max_num_seqs: 1024
    n: 8 # GRPO requires multiple samples per prompt
    enable_chunked_prefill: true
  critic: null # GRPO does not require critic
  reward_model:
    enable: false
algorithm:
  gamma: 1.0
  lam: 1.0
  adv_estimator: grpo
  kl_penalty: kl
  kl_ctrl:
    type: fixed
    kl_coef: 0.001
trainer:
  total_epochs: 3
  project_name: llama3-grpo-finetune
  experiment_name: gsm8k-v1
  logger: ['console']
  nnodes: 1
  n_gpus_per_node: 8
  default_local_dir: ./checkpoints/llama3-grpo
  save_freq: 100
  resume_mode: auto

4. 自定义 Reward 函数：提升训练可控性

4.1 实现 CustomRewardManager

由于 GRPO 不依赖外部 reward model，我们可以通过编写自定义奖励管理器来控制生成质量。

在 verl/workers/reward_manager/custom_reward.py 中新增：

from verl import DataProto
import torch

class LengthBonusRewardManager:
    """鼓励更长、结构完整的回答"""
    def __init__(self, tokenizer, num_examine=5) -> None:
        self.tokenizer = tokenizer
        self.num_examine = num_examine
        self.print_counter = 0

    def __call__(self, data: DataProto):
        reward_tensor = torch.zeros_like(data.batch['responses'], dtype=torch.float32)
        for i in range(len(data)):
            if self.print_counter < self.num_examine:
                print(f"Sample {i}: Prompt and Response:")
                data_item = data[i]
                prompt_ids = data_item.batch['prompts']
                response_ids = data_item.batch['responses']
                valid_prompt_len = data_item.batch['attention_mask'][:prompt_ids.shape[-1]].sum().item()
                valid_response_len = data_item.batch['attention_mask'][prompt_ids.shape[-1]:].sum().item()
                prompt_str = self.tokenizer.decode(prompt_ids[-valid_prompt_len:], skip_special_tokens=True)
                response_str = self.tokenizer.decode(response_ids[:valid_response_len], skip_special_tokens=True)
                if self.print_counter < self.num_examine:
                    print("Prompt:", prompt_str)
                    print("Response:", response_str)
                self.print_counter += 1
            # 奖励长度 + 结尾标点存在性
            length_score = len(response_str.split()) * 0.1 # 按词数打分
            ends_with_punct = 1.0 if response_str.strip()[-1] in '.?!' else 0.0
            final_score = length_score + ends_with_punct
            reward_tensor[i, valid_response_len - 1] = final_score
        return reward_tensor

并在 verl/workers/reward_manager/__init__.py 中注册：

from .custom_reward import LengthBonusRewardManager

__all__ = ['NaiveRewardManager', 'LengthBonusRewardManager']

4.2 在配置中启用自定义 Reward

修改主配置文件中的 reward_manager 字段：

reward_manager: custom
  custom_reward_class: LengthBonusRewardManager

并在 main_ppo.py 中增加导入逻辑：

if config.reward_manager == 'custom':
    from verl.workers.reward_manager.custom_reward import LengthBonusRewardManager
    reward_manager = LengthBonusRewardManager(tokenizer=tokenizer, num_examine=5)
else:
    reward_manager = NaiveRewardManager(tokenizer=tokenizer, rm_model=None)

5. 启动训练与常见问题解决

5.1 启动命令

set -x
export VLLM_ATTENTION_BACKEND=XFORMERS
CONFIG_PATH="./configs/grpo_llama3.yaml"
python3 -m verl.trainer.main_ppo --config_path=$CONFIG_PATH

5.2 常见问题与解决方案

❌ 问题 1：vLLM 加载 Llama-3 报错'unknown model type'

原因：vLLM 版本过低不支持 Llama-3 架构。

解决方案：升级至 vLLM >= 0.5.4，并确认 HF 模型路径正确。

❌ 问题 2：FSDP 训练时 OOM

原因：Llama-3-8B 显存占用高，batch size 设置过大。

解决方案：

• 降低 ppo_micro_batch_size_per_gpu 至 1
• 开启 enable_gradient_checkpointing
• 使用 bfloat16 精度
• 调整 gpu_memory_utilization: 0.7

❌ 问题 3：生成结果重复或陷入循环

原因：temperature 设置过低或 top_p 过小。

建议参数：

rollout:
  temperature: 0.7~1.0
  top_p: 0.9~0.95
  n: 8 # 提高多样性采样

6. 模型保存与 HuggingFace 格式转换

6.1 检查点结构说明

verl 默认保存的是包含优化器状态的完整 checkpoint，路径结构如下：

checkpoints/
└── llama3-grpo/
    └── global_step_100/
        └── actor/
            ├── model_world_size_8_rank_0.pt
            ├── model_world_size_8_rank_1.pt
            ...

6.2 转换为 HuggingFace 可加载格式

使用以下脚本将 FSDP 分片权重合并为标准 HF 格式：

#!/usr/bin/env python
import torch
from collections import defaultdict
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer

def convert_fsdp_to_hf(fsdp_ckpt_dir, hf_model_path, output_dir, step=100, world_size=8):
    state_dict = defaultdict(list)
    ckpt_path = f"{fsdp_ckpt_dir}/global_step_{step}/actor"
    for rank in range(world_size):
        file_path = f"{ckpt_path}/model_world_size_{world_size}_rank_{rank}.pt"
        print(f"Loading rank {rank}...")
        shard = torch.load(file_path)
        for k, v in shard.items():
            state_dict[k].append(v.to_local())
    merged_state_dict = {}
    for k, v_list in state_dict.items():
        merged_state_dict[k] = torch.cat(v_list, dim=0)
    config = AutoConfig.from_pretrained(hf_model_path)
    model = AutoModelForCausalLM.from_config(config)
    model.load_state_dict(merged_state_dict)
    model.save_pretrained(output_dir, max_shard_size="10GB")
    tokenizer = AutoTokenizer.from_pretrained(hf_model_path)
    tokenizer.save_pretrained(output_dir)
    print(f"Model saved to {output_dir}")

if __name__ == "__main__":
    convert_fsdp_to_hf(
        fsdp_ckpt_dir="./checkpoints",
        hf_model_path="meta-llama/Meta-Llama-3-8B-Instruct",
        output_dir="./hf_checkpoints/llama3-grpo-step100",
        step=100,
        world_size=8
    )

转换完成后即可用标准方式加载：

from transformers import pipeline
pipe = pipeline(
    "text-generation",
    model="./hf_checkpoints/llama3-grpo-step100",
    device_map="auto",
    torch_dtype=torch.bfloat16
)

7. 总结

本文系统介绍了如何将 Llama-3 系列模型成功接入 verl 框架，完成基于 GRPO 的强化学习训练全流程。主要内容包括：

1. 环境搭建：正确安装 verl 及其依赖，特别是 vLLM 和 FlashAttention；
2. 模型接入：通过配置文件指定 Llama-3 路径，并处理其特有的 chat template；
3. 训练配置：合理设置 GRPO 参数，尤其是 n 采样数、KL 正则系数等；
4. 自定义奖励：实现 CustomRewardManager 来引导模型生成更高质量响应；
5. 问题排查：针对 OOM、生成异常等问题提供实用调参建议；
6. 模型导出：将 verl 的分布式 checkpoint 转换为 HuggingFace 标准格式，便于部署。

通过上述实践，开发者可以充分利用 verl 的高性能与灵活性，高效完成 Llama-3 等先进大模型的对齐训练任务，为实际应用场景提供强有力的支持。