DAPO代码实现浅析

参考verl对dapo的实现，首先咱们看一下入口.sh和.py文件，在./recipe/dapo/文件夹中有以下目录

1. .
2. ├── config
3. │   ├── dapo_megatron_trainer.yaml
4. │   └── dapo_trainer.yaml
5. ├── dapo_ray_trainer.py
6. ├── main_dapo.py
7. ├── prepare_dapo_data.sh
8. ├── README.md
9. ├── run_dapo_qwen2.5_32b.sh

复制代码

整体的执行顺序：

• main_dapo.py：数据加载初始化、初始化actor_rollout model、rm model，加载reward_manager
  
• dapo_ray_trainer.py：RL训练流程
  
  
  
  • 对batch进行repeate，每个q采样n次
    
  • 记录每个采样的log，以及对应的reward_score 和 advantage
    
    
    
    • filter掉一个q的所有sample的score都是1或都是0，继续获取新的q进行采样，直到满足要求的batch的大小达到train_prompt_bsz。（值得注意的是，batch大小是gen_prompt_bsz=3*train_prompt_bsz，通过提高采样q的个数，避免满足要求的q不到train_prompt_bsz）。
      
    
    
  • 每mini_batch的data进行模型更新
    
    
    
    • 每micro_batch的data进行前向传播（token-mean loss）与梯度计算
      
    
    
  
  

具体代码实例:
main_dapo.py

1. # Copyright 2024 Bytedance Ltd. and/or its affiliates
2. #
3. # Licensed under the Apache License, Version 2.0 (the "License");
4. # you may not use this file except in compliance with the License.
5. # You may obtain a copy of the License at
6. #
7. #     http://www.apache.org/licenses/LICENSE-2.0
8. #
9. # Unless required by applicable law or agreed to in writing, software
10. # distributed under the License is distributed on an "AS IS" BASIS,
11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12. # See the License for the specific language governing permissions and
13. # limitations under the License.
14. """
15. Note that we don't combine the main with ray_trainer as ray_trainer is used by other main.
16. """
18. import os
19. import socket
21. import hydra
22. import ray
23. from omegaconf import OmegaConf
25. from verl.trainer.ppo.reward import load_reward_manager
26. from verl.utils.device import is_cuda_available
28. from .dapo_ray_trainer import RayDAPOTrainer
31. @hydra.main(config_path="config", config_name="dapo_trainer", version_base=None)
32. def main(config):
33.     run_ppo(config)
35. #################################################################
36. # RL训练入口
37. #################################################################
38. def run_ppo(config) -> None:
39.     if not ray.is_initialized():
40.         # this is for local ray cluster
41.         default_runtime_env = {
42.             "env_vars": {"TOKENIZERS_PARALLELISM": "true", "NCCL_DEBUG": "WARN", "VLLM_LOGGING_LEVEL": "WARN"}
43.         }
44.         ray_init_kwargs = config.ray_kwargs.get("ray_init", {})
45.         runtime_env_kwargs = ray_init_kwargs.get("runtime_env", {})
46.         runtime_env = OmegaConf.merge(default_runtime_env, runtime_env_kwargs)
47.         ray_init_kwargs = OmegaConf.create({**ray_init_kwargs, "runtime_env": runtime_env})
48.         print(f"ray init kwargs: {ray_init_kwargs}")
49.         ray.init(**OmegaConf.to_container(ray_init_kwargs))
51.     try:
52.         if (
53.             is_cuda_available
54.             and config.global_profiler.tool == "nsys"
55.             and OmegaConf.select(config.global_profiler, "steps") is not None
56.             and len(OmegaConf.select(config.global_profiler, "steps")) > 0
57.         ):
58.             nsight_options = OmegaConf.to_container(
59.                 config.global_profiler.global_tool_config.nsys.controller_nsight_options
60.             )
61.             runner = TaskRunner.options(runtime_env={"nsight": nsight_options}).remote()
62.         else:
63.             runner = TaskRunner.remote()
64.         ray.get(runner.run.remote(config))
65.     finally:
66.         if ray.is_initialized():
67.             ray.shutdown()
70. @ray.remote(num_cpus=1)  # please make sure main_task is not scheduled on head
71. class TaskRunner:
72.     def run(self, config):
73.         # print initial config
74.         from pprint import pprint
76.         from omegaconf import OmegaConf
78.         from verl.utils.fs import copy_to_local
80.         print(f"TaskRunner hostname: {socket.gethostname()}, PID: {os.getpid()}")
82.         pprint(OmegaConf.to_container(config, resolve=True))  # resolve=True will eval symbol values
83.         OmegaConf.resolve(config)
85.         # download the checkpoint from hdfs
86.         local_path = copy_to_local(config.actor_rollout_ref.model.path)
88.         # instantiate tokenizer
89.         from verl.utils import hf_processor, hf_tokenizer
91.         tokenizer = hf_tokenizer(local_path)
92.         processor = hf_processor(local_path, use_fast=True)  # used for multimodal LLM, could be none
94.         from verl.single_controller.ray import RayWorkerGroup
96.         #################################################################
97.         # 加载actor worker
98.         #################################################################
99.         # define worker classes
100.         if config.actor_rollout_ref.actor.strategy in {"fsdp", "fsdp2"}:
101.             assert config.critic.strategy in {"fsdp", "fsdp2"}
103.             from verl.workers.fsdp_workers import ActorRolloutRefWorker, CriticWorker
105.             ray_worker_group_cls = RayWorkerGroup
107.         elif config.actor_rollout_ref.actor.strategy == "megatron":
108.             assert config.actor_rollout_ref.actor.strategy == config.critic.strategy
109.             from verl.workers.megatron_workers import ActorRolloutRefWorker, CriticWorker
111.             ray_worker_group_cls = RayWorkerGroup
113.         else:
114.             raise NotImplementedError
116.         from verl.trainer.ppo.ray_trainer import ResourcePoolManager, Role
118.         role_worker_mapping = {
119.             Role.ActorRollout: ray.remote(ActorRolloutRefWorker),
120.             Role.Critic: ray.remote(CriticWorker),
121.         }
123.         global_pool_id = "global_pool"
124.         resource_pool_spec = {
125.             global_pool_id: [config.trainer.n_gpus_per_node] * config.trainer.nnodes,
126.         }
127.         mapping = {
128.             Role.ActorRollout: global_pool_id,
129.             Role.Critic: global_pool_id,
130.         }
132.         # we should adopt a multi-source reward function here
133.         # - for rule-based rm, we directly call a reward score
134.         # - for model-based rm, we call a model
135.         # - for code related prompt, we send to a sandbox if there are test cases
136.         # - finally, we combine all the rewards together
137.         # - The reward type depends on the tag of the data
138.         if config.reward_model.enable:
139.             if config.reward_model.strategy in {"fsdp", "fsdp2"}:
140.                 from verl.workers.fsdp_workers import RewardModelWorker
141.             elif config.reward_model.strategy == "megatron":
142.                 from verl.workers.megatron_workers import RewardModelWorker
143.             else:
144.                 raise NotImplementedError
145.             role_worker_mapping[Role.RewardModel] = ray.remote(RewardModelWorker)
146.             mapping[Role.RewardModel] = global_pool_id
148.         # reference model
149.         if config.algorithm.use_kl_in_reward or config.actor_rollout_ref.actor.use_kl_loss:
150.             role_worker_mapping[Role.RefPolicy] = ray.remote(ActorRolloutRefWorker)
151.             mapping[Role.RefPolicy] = global_pool_id
153.         #################################################################
154.         # 加载reward manager函数。用于根据data计算对应的reward score
155.         #################################################################
156.         reward_fn = load_reward_manager(
157.             config,
158.             tokenizer,
159.             0,
160.             max_resp_len=config.data.max_response_length,
161.             overlong_buffer_cfg=config.reward_model.overlong_buffer,
162.         )
164.         # Note that we always use function-based RM for validation
165.         val_reward_fn = load_reward_manager(
166.             config,
167.             tokenizer,
168.             1,
169.             max_resp_len=config.data.max_response_length,
170.             overlong_buffer_cfg=config.reward_model.overlong_buffer,
171.         )
172.         resource_pool_manager = ResourcePoolManager(resource_pool_spec=resource_pool_spec, mapping=mapping)
174.         #################################################################
175.         # 加载主要的DAPO RL训练类，并运行.fit()
176.         #################################################################
177.         trainer = RayDAPOTrainer(
178.             config=config,
179.             tokenizer=tokenizer,
180.             processor=processor,
181.             role_worker_mapping=role_worker_mapping,
182.             resource_pool_manager=resource_pool_manager,
183.             ray_worker_group_cls=ray_worker_group_cls,
184.             reward_fn=reward_fn,
185.             val_reward_fn=val_reward_fn,
186.         )
187.         trainer.init_workers()
188.         trainer.fit()
191. if __name__ == "__main__":
192.     main()

复制代码

我们紧接着来看一下from verl.trainer.ppo.reward import load_reward_manager。
配置文件中verl/recipe/dapo/run_dapo_qwen2.5_32b.sh给出了reward的类型

1. enable_overlong_buffer=True
2. overlong_buffer_len=$((1024 * 4)) # overlong soft
3. overlong_penalty_factor=1.0
5. reward_model.reward_manager=dapo \
6. reward_model.overlong_buffer.enable=${enable_overlong_buffer} \
7. reward_model.overlong_buffer.len=${overlong_buffer_len} \
8. reward_model.overlong_buffer.penalty_factor=${overlong_penalty_factor} \

复制代码

verl.trainer.ppo.reward.py

1. def load_reward_manager(
2.     config: DictConfig, tokenizer: Any, num_examine: int, **reward_kwargs: Any
3. ) -> AbstractRewardManager:
4.     """
5.     Load and initialize a reward manager based on the configuration.
7.     Args:
8.         config: PPO trainer configuration object containing reward_model fields.
9.         tokenizer: Tokenizer object used for processing text.
10.         num_examine: Number of samples to examine.
11.         **reward_kwargs: Additional keyword arguments for the reward manager.
13.     Returns:
14.         An instance of the specified reward manager class.
15.     """
17.     # Try to get a custom reward function based on the configuration
18.     # user defined reward manager can be registered in custom_reward_fn
19.     compute_score = get_custom_reward_fn(config)
20.     final_compute_score = compute_score
22.     # The list of pre-defined reward managers are defined in `verl/workers/reward_manager/`:
23.     # naive: NaiveRewardManager
24.     # prime: PrimeRewardManager
25.     # batch: BatchRewardManager
26.     # dapo: DAPORewardManager
27.     # Note(haibin.lin): For custom reward managers, please make sure they are imported and
28.     # registered via `verl.workers.reward_manager.register`
29.     # By default reward_manager is set to naive (NaiveRewardManager)
30.     #################################################################
31.     # 在这里加载具体的reward_manager
32.     #################################################################
33.     reward_manager_name = config.reward_model.get("reward_manager", "naive")
34.     reward_manager_cls = get_reward_manager_cls(reward_manager_name)
36.     if compute_score is None:
37.         sandbox_config = config.reward_model.get("sandbox_fusion")
38.         sandbox_url = sandbox_config.get("url") if sandbox_config else None
39.         memory_limit_mb = sandbox_config.get("memory_limit_mb", 1024)
40.         if sandbox_url:
41.             sandbox_manager = multiprocessing.Manager()
42.             # Create a semaphore to control concurrent access to the sandbox
43.             _concurrent_semaphore = sandbox_manager.Semaphore(sandbox_config.get("max_concurrent", 64))
44.             final_compute_score = partial(
45.                 default_compute_score,
46.                 sandbox_fusion_url=sandbox_url,
47.                 concurrent_semaphore=_concurrent_semaphore,
48.                 memory_limit_mb=memory_limit_mb,
49.             )
50.         else:
51.             final_compute_score = default_compute_score
53.     #################################################################
54.     # 这里的reward_manager_cls 其实是DAPO，
55.     #################################################################
56.     # Instantiate and return the reward manager with the specified parameters
57.     return reward_manager_cls(
58.         tokenizer=tokenizer,
59.         num_examine=num_examine,
60.         compute_score=final_compute_score,
61.         reward_fn_key=config.data.reward_fn_key,
62.         **reward_kwargs,
63.     )

复制代码

这里需要知道dapo的reward_manager_cls 具体是什么，因为reward需要batch数据才能计算，因此对于reward manager咱们先按下不表（其实dapo对应的reward_manager_cls是在verl/verl/workers/reward_manager/dapo.py），先去dapo_ray_trainer.py看一下batch是怎么采样的，再回来仔细阅读reward的具体计算方法。
dapo_ray_trainer.py

1. #################################################################
2. # RayDAPOTrainer继承于RayPPOTrainer
3. # fit()函数：执行dapo的训练，包括（1）动态采样（2）overlong soft reward计算（3）token-level loss
4. #################################################################
5. class RayDAPOTrainer(RayPPOTrainer):
6.     """
7.     Note that this trainer runs on the driver process on a single CPU/GPU node.
8.     """
10.     def fit(self):
11.         """
12.         The training loop of PPO.
13.         The driver process only need to call the compute functions of the worker group through RPC
14.         to construct the PPO dataflow.
15.         The light-weight advantage computation is done on the driver process.
16.         """
17.         from omegaconf import OmegaConf
19.         from verl.utils.tracking import Tracking
21.         logger = Tracking(
22.             project_name=self.config.trainer.project_name,
23.             experiment_name=self.config.trainer.experiment_name,
24.             default_backend=self.config.trainer.logger,
25.             config=OmegaConf.to_container(self.config, resolve=True),
26.         )
28.         self.global_steps = 0
29.         self.gen_steps = 0
31.         # load checkpoint before doing anything
32.         self._load_checkpoint()
34.         # perform validation before training
35.         # currently, we only support validation using the reward_function.
36.         if self.val_reward_fn is not None and self.config.trainer.get("val_before_train", True):
37.             val_metrics = self._validate()
38.             assert val_metrics, f"{val_metrics=}"
39.             pprint(f"Initial validation metrics: {val_metrics}")
40.             logger.log(data=val_metrics, step=self.global_steps)
41.             if self.config.trainer.get("val_only", False):
42.                 return
44.         if self.config.actor_rollout_ref.rollout.get("skip_rollout", False):
45.             rollout_skip = RolloutSkip(self.config, self.actor_rollout_wg)
46.             rollout_skip.wrap_generate_sequences()
48.         # add tqdm
49.         progress_bar = tqdm(total=self.total_training_steps, initial=self.global_steps, desc="Training Progress")
51.         # we start from step 1
52.         self.global_steps += 1
53.         self.gen_steps += 1
54.         last_val_metrics = None
56.         prev_step_profile = False
57.         curr_step_profile = (
58.             self.global_steps in self.config.global_profiler.steps
59.             if self.config.global_profiler.steps is not None
60.             else False
61.         )
62.         next_step_profile = False
64.         timing_raw = defaultdict(float)
65.         batch = None
66.         #################################################################
67.         # num_prompt_in_batch：记录filter后，std不等于0的q的个数，当模型更新后重新赋值为0
68.         # num_gen_batches: 记录当前使用了多少个gen_batch，当模型更新后重新赋值为0
69.         #################################################################
70.         num_prompt_in_batch = 0
71.         num_gen_batches = 0
72.         #################################################################
73.         # 正式开始训练，循环每个epoch后，循环每个gen_batch
74.         #################################################################
75.         for epoch in range(self.config.trainer.total_epochs):
76.             for batch_dict in self.train_dataloader:
77.                 metrics = {}
79.                 with marked_timer("start_profile", timing_raw):
80.                     self._start_profiling(
81.                         not prev_step_profile and curr_step_profile
82.                         if self.config.global_profiler.profile_continuous_steps
83.                         else curr_step_profile
84.                     )
86.                 #################################################################
87.                 # new_batch 是DataProto类型（具体见verl/verl/protocol.py），
88.                 # new_batch.batch是TensorDict类型
89.                 # new_batch中q的数量是可训练batch大小的3倍（增加采样的batch的q的个数）
90.                 #################################################################
91.                 new_batch: DataProto = DataProto.from_single_dict(batch_dict)
92.                 num_gen_batches += 1
93.                 # pop those keys for generation
94.                 if "multi_modal_data" in new_batch.non_tensor_batch.keys():
95.                     gen_batch = new_batch.pop(
96.                         batch_keys=["input_ids", "attention_mask", "position_ids"],
97.                         non_tensor_batch_keys=["raw_prompt_ids", "multi_modal_data"],
98.                     )
99.                 else:
100.                     # 从new_batch中提取对应的key，构建gen_batch
101.                     gen_batch = new_batch.pop(
102.                         batch_keys=["input_ids", "attention_mask", "position_ids"],
103.                         non_tensor_batch_keys=["raw_prompt_ids"],
104.                     )
105.                 # 这里为什么要repeate呢，因为每个prompt要采样n次，所以repeat n次。这里的interleave=True
106.                 # gen_batch: (bsz, response_length),
107.                                                                 # gen_batch_output: (bsz*n, response_length)
108.                 gen_batch_output = gen_batch.repeat(
109.                     repeat_times=self.config.actor_rollout_ref.rollout.n, interleave=True
110.                 )
112.                 is_last_step = self.global_steps >= self.total_training_steps
114.                 with marked_timer("step", timing_raw):
115.                     # generate a batch
116.                     with marked_timer("gen", timing_raw, "red"):
117.                         gen_batch_output = self.actor_rollout_wg.generate_sequences(gen_batch_output)
118.                         timing_raw.update(gen_batch_output.meta_info["timing"])
119.                         gen_batch_output.meta_info.pop("timing", None)
121.                     # 这个advatange 可以先忽略。RMAX需要先计算 贪心采样的sample的logits作为后序adv计算的baseline
122.                     if self.config.algorithm.adv_estimator == AdvantageEstimator.REMAX:
123.                         with marked_timer("gen_max", timing_raw, "red"):
124.                             gen_baseline_batch = deepcopy(gen_batch)
125.                             # 这里是贪心采样的baseline，do_sample = False
126.                             gen_baseline_batch.meta_info["do_sample"] = False
127.                             gen_baseline_output = self.actor_rollout_wg.generate_sequences(gen_baseline_batch)
129.                             new_batch = new_batch.union(gen_baseline_output)
130.                             # compute reward model score on new_batch
131.                             rm_scores = None
132.                             if self.use_rm and "rm_scores" not in new_batch.batch.keys():
133.                                 rm_scores = self.rm_wg.compute_rm_score(new_batch)
134.                                 new_batch = new_batch.union(rm_scores)
135.                             reward_baseline_tensor, _ = compute_reward(new_batch, self.reward_fn)
136.                             reward_baseline_tensor = reward_baseline_tensor.sum(dim=-1)
138.                             keys_to_pop = set(gen_baseline_output.batch.keys())
139.                             if rm_scores is not None:
140.                                 keys_to_pop.update(rm_scores.batch.keys())
141.                             new_batch.pop(batch_keys=list(keys_to_pop))
143.                             new_batch.batch["reward_baselines"] = reward_baseline_tensor
145.                             del rm_scores, gen_baseline_batch, gen_baseline_output
147.                     #################################################################
148.                     # new_batch的大小是gen_prompt_bsz
149.                     # 对每一个prompt设置一个专属的标识 uid
150.                                                                                 # 之所以设置uid，是因为之后对sample计算reward时，需要对同一个q的n个sample的reward标准化
151.                     #################################################################
152.                     new_batch.non_tensor_batch["uid"] = np.array(
153.                         [str(uuid.uuid4()) for _ in range(len(new_batch.batch))], dtype=object
154.                     )
155.                     # 对batch中的每个key进行repeat（这里应该主要是对uid进行repeat）
156.                     # repeat to align with repeated responses in rollout
157.                     new_batch = new_batch.repeat(repeat_times=self.config.actor_rollout_ref.rollout.n, interleave=True)
158.                     # 把采样完的放到new_batch中
159.                     new_batch = new_batch.union(gen_batch_output)
161.                     with marked_timer("reward", timing_raw, "yellow"):
162.                         # compute scores. Support both model and function-based.
163.                         # We first compute the scores using reward model. Then, we call reward_fn to combine
164.                         # the results from reward model and rule-based results.
165.                         if self.use_rm and "rm_scores" not in new_batch.batch.keys():
166.                             # we first compute reward model score
167.                             reward_tensor = self.rm_wg.compute_rm_score(new_batch)
168.                             new_batch = new_batch.union(reward_tensor)
170.                         # 计算new_batch各个采样的reward，根据设置好的self.reward_fn
171.                         # we combine with rule-based rm
172.                         reward_tensor, reward_extra_infos_dict = compute_reward(new_batch, self.reward_fn)
174.                         new_batch.batch["token_level_scores"] = reward_tensor
176.                         if reward_extra_infos_dict:
177.                             new_batch.non_tensor_batch.update(
178.                                 {k: np.array(v) for k, v in reward_extra_infos_dict.items()}
179.                             )
181.                         # compute rewards. apply_kl_penalty if available
182.                         if self.config.algorithm.use_kl_in_reward:
183.                             new_batch, kl_metrics = apply_kl_penalty(
184.                                 new_batch, kl_ctrl=self.kl_ctrl_in_reward, kl_penalty=self.config.algorithm.kl_penalty
185.                             )
186.                             metrics.update(
187.                                 kl_metrics
188.                             )  # TODO: This will be cleared if we use multiple genenration batches
189.                         else:
190.                             new_batch.batch["token_level_rewards"] = new_batch.batch["token_level_scores"]
191.                                                                                
192.                     #################################################################
193.                     # dapo的filter（dynamic sample）部分
194.                     #################################################################
195.                     if not self.config.algorithm.filter_groups.enable:
196.                         batch = new_batch
197.                     else:  # NOTE: When prompts after filtering is less than train batch size,
198.                         # we skip to the next generation batch
199.                         metric_name = self.config.algorithm.filter_groups.metric
200.                         if metric_name == "seq_final_reward":
201.                             # Turn to numpy for easier filtering
202.                             new_batch.non_tensor_batch["seq_final_reward"] = (
203.                                 new_batch.batch["token_level_rewards"].sum(dim=-1).numpy()
204.                             )
205.                         elif metric_name == "seq_reward":
206.                             new_batch.non_tensor_batch["seq_reward"] = (
207.                                 new_batch.batch["token_level_scores"].sum(dim=-1).numpy()
208.                             )
210.                         # {uid: [r1,r2,r3,...,rn], uid: [...], ...}，记录每个轨迹所有采样的reward
211.                         # Collect the sequence reward for each trajectory
212.                         prompt_uid2metric_vals = defaultdict(list)
213.                         for uid, metric_val in zip(
214.                             new_batch.non_tensor_batch["uid"], new_batch.non_tensor_batch[metric_name], strict=True
215.                         ):
216.                             prompt_uid2metric_vals[uid].append(metric_val)
218.                         # 每个q的reward的std
219.                         prompt_uid2metric_std = {}
220.                         for prompt_uid, metric_vals in prompt_uid2metric_vals.items():
221.                             prompt_uid2metric_std[prompt_uid] = np.std(metric_vals)
223.                         # 保留reward std不是0的q的uid
224.                         kept_prompt_uids = [
225.                             uid
226.                             for uid, std in prompt_uid2metric_std.items()
227.                             if std > 0 or len(prompt_uid2metric_vals[uid]) == 1
228.                         ]
229.                         # 累积std不是0的q
230.                         num_prompt_in_batch += len(kept_prompt_uids)
232.                         # 记录留下来的q的sample的idx
233.                         kept_traj_idxs = []
234.                         for idx, traj_from_prompt_uid in enumerate(new_batch.non_tensor_batch["uid"]):
235.                             if traj_from_prompt_uid in kept_prompt_uids:
236.                                 kept_traj_idxs.append(idx)
238.                         # 基于traj的id，检索对应的new_batch
239.                         new_batch = new_batch[kept_traj_idxs]
240.                         # batch是留下的traj数据的累积
241.                         batch = new_batch if batch is None else DataProto.concat([batch, new_batch])
243.                         # .sh文件配置的 可以训练的batch的最小大小（q的数量）
244.                         prompt_bsz = self.config.data.train_batch_size
245.                         # 如果现有的累积filter出来的q的数量小于 配置的最小数量，则continue继续使用下一个new_batch进行累积
246.                         if num_prompt_in_batch < prompt_bsz:
247.                             print(f"{num_prompt_in_batch=} < {prompt_bsz=}")
248.                             max_num_gen_batches = self.config.algorithm.filter_groups.max_num_gen_batches
249.                             # max_num_gen_batches是最多可以使用的gen_batch的个数
250.                             # 如果其小于0的话，即没有限制；若num_gen_batches < max_num_gen_batches则继续continue
251.                             if max_num_gen_batches <= 0 or num_gen_batches < max_num_gen_batches:
252.                                 print(f"{num_gen_batches=}. Keep generating...")
253.                                 self.gen_steps += 1
254.                                 is_last_step = self.global_steps >= self.total_training_steps
255.                                 continue
256.                             else:
257.                                 raise ValueError(
258.                                     f"{num_gen_batches=} >= {max_num_gen_batches=}."
259.                                     + " Generated too many. Please check if your data are too difficult."
260.                                     + " You could also try set max_num_gen_batches=0 to enable endless trials."
261.                                 )
262.                       # 累积的符合的q个个数>=最小的可以训练的batch的大小  
263.                       else:
264.                             # Align the batch
265.                             traj_bsz = self.config.data.train_batch_size * self.config.actor_rollout_ref.rollout.n
266.                             #################################################################
267.                             # 对齐一下，多余的轨迹会被抛弃，不知道会不会导致采样的利用效率不高，
268.                             # 会不会导致一些轨迹根本不会被训练到
269.                             #################################################################
270.                             batch = batch[:traj_bsz]
272.                     #################################################################
273.                     # actor模型更新
274.                     #################################################################
275.                     # === Updating ===
277.                     batch.batch["response_mask"] = compute_response_mask(batch)
279.                     # Balance the number of valid tokens across DP ranks.
280.                     # NOTE: This usually changes the order of data in the `batch`,
281.                     # which won't affect the advantage calculation (since it's based on uid),
282.                     # but might affect the loss calculation (due to the change of mini-batching).
283.                     # TODO: Decouple the DP balancing and mini-batching.
284.                     if self.config.trainer.balance_batch:
285.                         self._balance_batch(batch, metrics=metrics)
287.                     # compute global_valid tokens
288.                     batch.meta_info["global_token_num"] = torch.sum(batch.batch["attention_mask"], dim=-1).tolist()
290.                     #################################################################
291.                     # 记录filter后的batch的每个traj的采样时的logtis（token-level）
292.                     # 用于计算重要性采样的比值
293.                     #################################################################
294.                     # recompute old_log_probs
295.                     with marked_timer("old_log_prob", timing_raw, "blue"):
296.                         old_log_prob = self.actor_rollout_wg.compute_log_prob(batch)
297.                         entropys = old_log_prob.batch["entropys"]
298.                         response_masks = batch.batch["response_mask"]
299.                         loss_agg_mode = self.config.actor_rollout_ref.actor.loss_agg_mode
300.                         # 这里dapo的loss_agg_mode是“token_mean”
301.                         entropy_agg = agg_loss(loss_mat=entropys, loss_mask=response_masks, loss_agg_mode=loss_agg_mode)
302.                         old_log_prob_metrics = {"actor/entropy": entropy_agg.detach().item()}
303.                         metrics.update(old_log_prob_metrics)
304.                         old_log_prob.batch.pop("entropys")
305.                         batch = batch.union(old_log_prob)
307.                     if self.use_reference_policy:
308.                         # compute reference log_prob
309.                         with marked_timer("ref", timing_raw, "olive"):
310.                             ref_log_prob = self.ref_policy_wg.compute_ref_log_prob(batch)
311.                             batch = batch.union(ref_log_prob)
313.                     # compute values
314.                     if self.use_critic:
315.                         with marked_timer("values", timing_raw, "cyan"):
316.                             values = self.critic_wg.compute_values(batch)
317.                             batch = batch.union(values)
319.                     # 计算token_level的重要性采样
320.                     # Compute rollout IS weights and mismatch metrics (inherited from RayPPOTrainer)
321.                     batch, is_metrics = self.compute_rollout_importance_weights_and_add_to_batch(batch)
322.                     # IS and mismatch metrics already have mismatch/ prefix
323.                     metrics.update(is_metrics)
325.                     #################################################################
326.                     # 计算advantage
327.                     #################################################################
328.                     with marked_timer("adv", timing_raw, "brown"):
329.                         # compute advantages, executed on the driver process
330.                         norm_adv_by_std_in_grpo = self.config.algorithm.get("norm_adv_by_std_in_grpo", True)
331.                         batch = compute_advantage(
332.                             batch,
333.                             adv_estimator=self.config.algorithm.adv_estimator,
334.                             gamma=self.config.algorithm.gamma,
335.                             lam=self.config.algorithm.lam,
336.                             num_repeat=self.config.actor_rollout_ref.rollout.n,
337.                             norm_adv_by_std_in_grpo=norm_adv_by_std_in_grpo,
338.                         )
340.                     # update critic
341.                     if self.use_critic:
342.                         with marked_timer("update_critic", timing_raw, "pink"):
343.                             critic_output = self.critic_wg.update_critic(batch)
344.                         critic_output_metrics = reduce_metrics(critic_output.meta_info["metrics"])
345.                         metrics.update(critic_output_metrics)
347.                     # implement critic warmup
348.                     if self.config.trainer.critic_warmup <= self.global_steps:
349.                         #################################################################
350.                         # 更新actor model（batch的大小是train_prompt_size）
351.                         # 每个mini_bsz 更新一次模型（参数-累积梯度）
352.                         # 每个micro_bsz 累积一次梯度
353.                         #################################################################
354.                         # update actor
355.                         with marked_timer("update_actor", timing_raw, "red"):
356.                             actor_output = self.actor_rollout_wg.update_actor(batch)
357.                         actor_output_metrics = reduce_metrics(actor_output.meta_info["metrics"])
358.                         metrics.update(actor_output_metrics)
360.                     # Log rollout generations if enabled
361.                     rollout_data_dir = self.config.trainer.get("rollout_data_dir", None)
362.                     if rollout_data_dir:
363.                         self._log_rollout_data(batch, reward_extra_infos_dict, timing_raw, rollout_data_dir)
365.                 # validate
366.                 if (
367.                     self.val_reward_fn is not None
368.                     and self.config.trainer.test_freq > 0
369.                     and (is_last_step or self.global_steps % self.config.trainer.test_freq == 0)
370.                 ):
371.                     with marked_timer("testing", timing_raw, "green"):
372.                         val_metrics: dict = self._validate()
373.                         if is_last_step:
374.                             last_val_metrics = val_metrics
375.                     metrics.update(val_metrics)
377.                 if self.config.trainer.save_freq > 0 and (
378.                     is_last_step or self.global_steps % self.config.trainer.save_freq == 0
379.                 ):
380.                     with marked_timer("save_checkpoint", timing_raw, "green"):
381.                         self._save_checkpoint()
383.                 with marked_timer("stop_profile", timing_raw):
384.                     next_step_profile = (
385.                         self.global_steps + 1 in self.config.global_profiler.steps
386.                         if self.config.global_profiler.steps is not None
387.                         else False
388.                     )
389.                     self._stop_profiling(
390.                         curr_step_profile and not next_step_profile
391.                         if self.config.global_profiler.profile_continuous_steps
392.                         else curr_step_profile
393.                     )
394.                     prev_step_profile = curr_step_profile
395.                     curr_step_profile = next_step_profile
397.                 # collect metrics
398.                 metrics.update(compute_data_metrics(batch=batch, use_critic=self.use_critic))
399.                 metrics.update(compute_timing_metrics(batch=batch, timing_raw=timing_raw))
400.                 # TODO: implement actual tflpo and theoretical tflpo
401.                 n_gpus = self.resource_pool_manager.get_n_gpus()
402.                 metrics.update(compute_throughout_metrics(batch=batch, timing_raw=timing_raw, n_gpus=n_gpus))
403.                 timing_raw = defaultdict(float)  # clear timing
405.                 metrics["train/num_gen_batches"] = num_gen_batches
406.                 batch = None
407.                 num_prompt_in_batch = 0
408.                 num_gen_batches = 0
410.                 # TODO: make a canonical logger that supports various backend
411.                 logger.log(data=metrics, step=self.global_steps)
413.                 if is_last_step:
414.                     pprint(f"Final validation metrics: {last_val_metrics}")
415.                     progress_bar.close()
416.                     return
418.                 progress_bar.update(1)
419.                 self.global_steps += 1
420.                 self.gen_steps += 1
421.         # check if last step checkpint exists
422.         checkpoint_dir = os.path.join(self.config.trainer.default_local_dir, f"global_step_{self.global_steps}")
423.         if not os.path.exists(checkpoint_dir):
424.             # save last step checkpoint
425.             timing_raw = defaultdict(float)
426.             with marked_timer("save_checkpoint", timing_raw, "green"):
427.                 self._save_checkpoint()
428.             metrics = {f"timing/{k}": v for k, v in timing_raw.items()}
429.             logger.log(data=metrics, step=self.global_steps)

复制代码

这时候咱们再看一下dapo的reward manager实现：主要和ppo的区别在于使用了overlong_buffer，计算长度的reward
verl/verl/workers/reward_manager/dapo.py

1. #################################################################
2. # 这里使用dapo注册了DAPORewardManager，因此可以用
3. # reward_manager_cls = get_reward_manager_cls(reward_manager_name)得到
4. #################################################################
5. @register("dapo")
6. class DAPORewardManager(AbstractRewardManager):
7.     """The reward manager."""
9.     def __init__(
10.         self,
11.         tokenizer,
12.         num_examine,
13.         compute_score=None,
14.         reward_fn_key="data_source",
15.         max_resp_len=None,
16.         overlong_buffer_cfg=None,
17.     ) -> None:
18.         self.tokenizer = tokenizer
19.         self.num_examine = num_examine  # the number of batches of decoded responses to print to the console
20.         self.compute_score = compute_score or default_compute_score
21.         self.reward_fn_key = reward_fn_key
22.         self.overlong_buffer_cfg = overlong_buffer_cfg
23.         self.max_resp_len = max_resp_len
25.         if self.overlong_buffer_cfg is not None:
26.             assert self.max_resp_len is not None, (
27.                 f"max_resp_len must be provided if {overlong_buffer_cfg=}, but got None"
28.             )
29.             assert self.max_resp_len >= self.overlong_buffer_cfg.len, (
30.                 "max_resp_len must be larger than overlong_buffer.len"
31.             )
33.     #################################################################
34.     # DAPO reward manager的主要函数
35.     #################################################################
36.     def __call__(self, data: DataProto, return_dict: bool = False):
37.         """We will expand this function gradually based on the available datasets"""
39.         # If there is rm score, we directly return rm score. Otherwise, we compute via rm_score_fn
40.         if "rm_scores" in data.batch.keys():
41.             if return_dict:
42.                 reward_extra_keys = data.meta_info.get("reward_extra_keys", [])
43.                 reward_extra_info = {key: data.non_tensor_batch[key] for key in reward_extra_keys}
44.                 return {"reward_tensor": data.batch["rm_scores"], "reward_extra_info": reward_extra_info}
45.             else:
46.                 return data.batch["rm_scores"]
49.         reward_tensor = torch.zeros_like(data.batch["responses"], dtype=torch.float32)
50.         reward_extra_info = defaultdict(list)
52.         already_print_data_sources = {}
54.         for i in range(len(data)):
55.             data_item = data[i]  # DataProtoItem
57.             prompt_ids = data_item.batch["prompts"]
59.             prompt_length = prompt_ids.shape[-1]
61.             ########################################################
62.             # 值得注意的是。prompt_ids是左填充的
63.             # response_ids是右填充的
64.             ########################################################
65.             valid_prompt_length = data_item.batch["attention_mask"][:prompt_length].sum()
66.             valid_prompt_ids = prompt_ids[-valid_prompt_length:]
68.             response_ids = data_item.batch["responses"]
69.             valid_response_length = data_item.batch["attention_mask"][prompt_length:].sum()
70.             valid_response_ids = response_ids[:valid_response_length]
72.             # decode
73.             prompt_str = self.tokenizer.decode(valid_prompt_ids, skip_special_tokens=True)
74.             response_str = self.tokenizer.decode(valid_response_ids, skip_special_tokens=True)
75.             eos_token = self.tokenizer.eos_token
76.             if response_str.endswith(eos_token):
77.                 response_str = response_str[: -len(eos_token)]
79.             ground_truth = data_item.non_tensor_batch["reward_model"]["ground_truth"]
81.             data_source = data_item.non_tensor_batch[self.reward_fn_key]
83.             extra_info = data_item.non_tensor_batch.get("extra_info", {})
85.             rollout_reward_scores = data_item.non_tensor_batch.get("reward_scores", {})
87.             extra_info["rollout_reward_scores"] = rollout_reward_scores
89.             result = self.compute_score(
90.                 data_source=data_source,
91.                 solution_str=response_str,
92.                 ground_truth=ground_truth,
93.                 extra_info=extra_info,
94.             )
96.             score: float
97.             if isinstance(result, dict):
98.                 score = result["score"]
99.                 # Store the information including original reward
100.                 for key, value in result.items():
101.                     reward_extra_info[key].append(value)
102.             else:
103.                 score = result
104.                 reward_extra_info["acc"].append(score)
106.             reward = score
108.             ########################################################
109.             # 这里是overlong reward的计算
110.             ########################################################
111.             if self.overlong_buffer_cfg.enable:
112.                 overlong_buffer_len = self.overlong_buffer_cfg.len
113.                 expected_len = self.max_resp_len - overlong_buffer_len
114.                 exceed_len = valid_response_length - expected_len
115.                 overlong_penalty_factor = self.overlong_buffer_cfg.penalty_factor
116.                 overlong_reward = min(-exceed_len / overlong_buffer_len * overlong_penalty_factor, 0)
117.                 reward += overlong_reward
118.                 if self.overlong_buffer_cfg.log:
119.                     reward_extra_info["overlong_reward"].append(overlong_reward)
120.                     reward_extra_info["overlong"].append(overlong_reward < 0)
122.             reward_tensor[i, valid_response_length - 1] = reward
124.             if data_source not in already_print_data_sources:
125.                 already_print_data_sources[data_source] = 0
127.             if already_print_data_sources[data_source] < self.num_examine:
128.                 already_print_data_sources[data_source] += 1
129.                 print("[prompt]", prompt_str)
130.                 print("[response]", response_str)
131.                 print("[ground_truth]", ground_truth)
132.                 if isinstance(result, dict):
133.                     for key, value in result.items():
134.                         print(f"[{key}]", value)
135.                 else:
136.                     print("[score]", score)
138.         if return_dict:
139.             return {
140.                 "reward_tensor": reward_tensor,
141.                 "reward_extra_info": reward_extra_info,
142.             }
143.         else:
144.             return reward_tensor

复制代码

dapo和ppo的具体区别可进一步参考：dapo readme

来源：程序园用户自行投稿发布，如果侵权，请联系站长删除
免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！

这个好，看起来很实用

不错，里面软件多更新就更好了