基于Contiue来阅读open-r1中的GRPO训练代码

原创快乐王子HP 快乐王子AI说 2025年04月03日 23:54 广东

前面安装了vscode[1]同时也安装了Coninue的相关插件[2]，现在想用它们来阅读一下open-r1项目的代码[3]。

首先，从启动训练开始(以GRPO为例子）

第一步，使用TRL的vLLM后端

CUDA_VISIBLE_DEVICES=0 trl vllm-serve --model deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B

第二步，启动GRPO

CUDA_VISIBLE_DEVICES=1,2,3,4,5,6,7 ACCELERATE_LOG_LEVEL=info \     accelerate launch --config_file recipes/accelerate_configs/zero2.yaml --num_processes 7 \     src/open_r1/grpo.py --config recipes/DeepSeek-R1-Distill-Qwen-1.5B/grpo/config_demo.yaml

查看vllm的服务启动帮助文档

usage: trl vllm-serve [-h] --model MODEL [--revision REVISION] [--tensor_parallel_size TENSOR_PARALLEL_SIZE] [--host HOST] [--port PORT] [--gpu_memory_utilization GPU_MEMORY_UTILIZATION] [--dtype DTYPE]                       [--max_model_len MAX_MODEL_LEN] [--enable_prefix_caching ENABLE_PREFIX_CACHING]

关于zero2.yaml文件

（https://github.com/huggingface/open-r1/blob/main/recipes/accelerate_configs/zero2.yaml）

1.核心配置:    - 使用 DeepSpeed 的 Zero Stage 2 优化 (zero_stage: 2)    - 混合精度训练采用 bf16 (mixed_precision: bf16)    - 单机 8 GPU 训练 (num_machines: 1, num_processes: 8)2.Zero Stage 2 特点:    - 优化器状态分区，减少内存占用    - 没有启用参数或优化器卸载 (offload_optimizer_device: none, offload_param_device: none)    - 比 Stage 3 内存效率稍低，但通信开销更小3.硬件配置:    - 纯 GPU 训练 (use_cpu: false)    - 不涉及 TPU (tpu_* 相关配置均为 false)    - 适合具有 8 个 GPU 的单个节点4.使用场景:    - 中等规模模型训练    - 当 GPU 内存足够容纳模型参数和激活值时    - 需要比 Zero Stage 1 更高的内存效率，但不想承受 Stage 3 的通信开销5.性能考虑:    - bf16 混合精度可以在支持它的硬件上提供良好的训练速度和内存效率    - 8 个 GPU 的配置适合大多数单节点服务器这个配置文件适合在单个多 GPU 节点上训练中等规模模型，在内存效率和通信开销之间取得平衡。

recipes/DeepSeek-R1-Distill-Qwen-1.5B/grpo/config_demo.yaml文件分析

（https://github.com/huggingface/open-r1/blob/main/recipes/DeepSeek-R1-Distill-Qwen-1.5B/grpo/config_demo.yaml）

1.模型架构:  - 基于1.5B参数的蒸馏版Qwen模型  - 使用Flash Attention 2优化注意力计算  - bfloat16混合精度训练2.训练策略:  - 采用GRPO(可能是一种强化学习优化算法)训练方法  - 结合三种奖励函数:准确性、格式正确性和标签计数  - 使用vLLM加速推理过程3.数据处理:  - 专门设计的复杂对话模板  - 数学领域专用数据集(OpenR1-Math-220k)  - 要求模型以和标签分步输出4.资源利用:  - 梯度检查点和梯度累积优化显存使用  - 适中的batch size(16)和上下文长度(512/2048)5.监控与部署:  - 完整的训练日志记录(W&B)  - 模型自动推送至HuggingFace Hub  - 严格的模型保存策略

grpo.py文件

（https://github.com/huggingface/open-r1/blob/main/src/open_r1/grpo.py）

```mermaidgraph TD    A[开始] --> B[设置随机种子]    B --> C[配置日志系统]    C --> D[检查检查点]    D --> E[初始化WandB]    E --> F[加载数据集]    F --> G[加载tokenizer]    G --> H[获取奖励函数]    H --> I[格式化对话数据]    I --> J[初始化模型参数]    J --> K[创建GRPOTrainer]    K --> L{是否有检查点?}    L -- 是 --> M[从检查点恢复训练]    L -- 否 --> N[开始新训练]    M --> O[训练模型]    N --> O    O --> P[保存模型和指标]    P --> Q{是否评估?}    Q -- 是 --> R[执行评估]    Q -- 否 --> S    R --> S[保存评估结果]    S --> T{是否推送至Hub?}    T -- 是 --> U[推送模型]    T -- 否 --> V[结束]    U --> V```

rewards.py

（https://github.com/huggingface/open-r1/blob/main/src/open_r1/rewards.py）

结合医学场景来探索

def medical_accuracy_reward(response: str, golden_answer: str) -> float:    """评估医学准确性，需要与标准医学答案对比"""    # 这里可以集成医学知识库或NLP模型进行专业评估    medical_terms_score = calculate_medical_terms_match(response, golden_answer)    treatment_score = evaluate_treatment_correctness(response, golden_answer)    return 0.6 * medical_terms_score + 0.4 * treatment_scoredef safety_reward(response: str) -> float:    """安全性评估：检查是否有危险建议"""    dangerous_keywords = ["自行停药", "未经医生", "高剂量", "随意服用"]    for keyword in dangerous_keywords:        if keyword in response:            return 0.0  # 发现危险建议直接0分    return 1.0def citation_reward(response: str) -> float:    """参考文献引用评估"""    citation_formats = ["[1]", "(Smith et al., 2020)", "根据最新指南"]    return 1.0 if any(fmt in response for fmt in citation_formats) else 0.5def patient_language_reward(response: str) -> float:    """患者友好语言评估"""    complex_terms = ["病理学", "分子机制", "流行病学"]    simplified_explanations = ["简单说", "通俗理解", "换句话说"]        complex_count = sum(term in response for term in complex_terms)    simple_count = sum(term in response for term in simplified_explanations)        if complex_count == 0:         return 1.0    return simple_count / (complex_count + 1)  # 确保至少解释了部分复杂术语def empathy_reward(response: str) -> float:    """同理心评估"""    empathy_keywords = ["理解您", "不用担心", "建议咨询", "我们会帮助"]    return min(1.0, 0.2 * sum(kw in response for kw in empathy_keywords))