Skip to content

[Feat] Add Sampled-token OPD Support to RL Training - #1977

Open
YanhuiDua wants to merge 8 commits into
InternLM:mainfrom
YanhuiDua:dev-opd
Open

[Feat] Add Sampled-token OPD Support to RL Training #1977
YanhuiDua wants to merge 8 commits into
InternLM:mainfrom
YanhuiDua:dev-opd

Conversation

@YanhuiDua

@YanhuiDua YanhuiDua commented Jul 27, 2026

Copy link
Copy Markdown
Collaborator

背景

本 PR 为 XTuner RL 训练链路增加 PG-OPD(Policy Gradient On-Policy Distillation)支持。Rollout 样本通过外部 Teacher server 获取 response token log probabilities,并使用 old_logprobs 构造 reverse-KL token advantage。通过 task_adv_weightopd_adv_weight,同时支持 Pure OPD 以及任务优势与 OPD 优势混合训练。支持sglang和lmdeploy推理后端,lmdeploy推理依赖PR: InternLM/lmdeploy#4793

主要修改

1. 新增 OPD 配置和异步 Teacher Logprob Client

新增 OPDConfigOPDTeacherConfig,用于配置 PG-OPD 的优势权重、Teacher 服务以及不同数据源到 Teacher 的路由关系。
示例配置如下:

opd_config = OPDConfig(                                                                                                                                                                                                
    mode="pg-opd",                                                                                                                                                                                                     
    task_adv_weight=0.0,                                                                                                                                                                                               
    opd_adv_weight=1.0,                                                                                                                                                                                                
    teachers=[                                                                                                                                                                                                         
        OPDTeacherConfig(                                                                                                                                                                                              
            name="teacher",                                                                                                                                                                                            
            endpoint="http://127.0.0.1:13141",                                                                                                                                                                         
            api_key=None,                                                                                                                                                                                              
            request_timeout_s=1200.0,                                                                                                                                                                                  
            max_retry_per_sample=2,                                                                                                                                                                                    
            max_concurrency=128,                                                                                                                                                                                       
        )                                                                                                                                                                                                              
    ],                                                                                                                                                                                                                 
    data_source_teacher_map={                                                                                                                                                                                          
        "math_dapo": "teacher",                                                                                                                                                                                        
    },                                                                                                                                                                                                                 
)                                                                                                                                                                                                                      

其中:

  • task_adv_weight:任务 reward advantage 的权重,设置为 0.0 时为 Pure OPD。
  • opd_adv_weight:OPD reverse-KL advantage 的权重。
  • teachers:Teacher 服务列表,支持配置 endpoint、超时、重试次数和最大并发数。
  • data_source_teacher_map:根据样本的 data_source 将其路由到指定 Teacher。

新增异步 TeacherLogprobClient。每个 AgentLoop 会根据 OPDTeacherConfig 创建对应的 Client,通过推理引擎 /generate 接口对 prompt_ids + response_ids 进行 scoring,只保留 response token 对应的 Teacher logprobs。

2. 将 Teacher scoring 接入 AgentLoop

将 Teacher scoring 集成到 AgentLoop 的 rollout group 收集流程中。Teacher scoring 请求重试耗尽后,对应样本会被标记为 FAILED,错误原因记录在 RolloutState.error_msg 中。失败样本不会进入 replay buffer,也不会参与训练。

3. 将 rollout group filtering 从 ProduceStrategy 下沉到 AgentLoop

is_valid_sample_fn 从 ProduceStrategy 迁移到 TaskSpecConfig.filter_func,由 AgentLoop 在收集 rollout group 时完成过滤。

当同时存在 filter_funcOPDConfig 时:

rollout / reward                                                                                                                                                                                                       
    -> filter_func                                                                                                                                                                                                     
    -> teacher_logprobs                                                                                                                                                                                                

只有状态为 COMPLETED 且通过 filter_func 的 group 才会请求 Teacher;未通过的 group 会被标记为 FILTERED,从而避免无效的 Teacher scoring 请求。

当不存在 filter_func 时,AgentLoop 使用 OPD fast path。对于 batch judge,执行顺序为:

rollout                                                                                                                                                                                                                
    -> teacher_logprobs                                                                                                                                                                                                
    -> reward

4. 新增 DAPO Math Pure PG-OPD recipe 和 TensorBoard 指标

新增 DAPO Math Pure PG-OPD 配置,新增 TensorBoard 指标 opd_reverse_kl,记录所有训练 worker 上有效 response token 的平均 reverse KL。

run_pg_opd.sh 会按照以下顺序启动:

  1. 检查 Student 模型、Teacher 模型、训练数据和运行依赖。
  2. TEACHER_CUDA_VISIBLE_DEVICES 上启动 SGLang Teacher Server。
  3. 轮询 /health_generate,等待 Teacher Server 完全就绪。
  4. Teacher Server ready 后,才在 STUDENT_CUDA_VISIBLE_DEVICES 上启动 Student rollout 和训练。
  5. 训练退出时清理 Teacher 进程和本次启动的 Ray 集群。

实验结果

Pure OPD

student model: Qwen3-8B
teacher model: Qwen3-32B
results: TBD

@YanhuiDua YanhuiDua changed the title [Feat] Add PG-OPD Support to RL Training [Feat] Add Sampled-token OPD Support to RL Training Jul 27, 2026
@YanhuiDua YanhuiDua changed the title [Feat] Add Sampled-token OPD Support to RL Training [Feat] Add Sampled-token OPD Support with SGLang to RL Training Jul 27, 2026
@YanhuiDua YanhuiDua changed the title [Feat] Add Sampled-token OPD Support with SGLang to RL Training [Feat] Add Sampled-token OPD Support to RL Training Jul 28, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant