代码阅读

主入口

fully_async_main

提示

代码

• fully_async_policy/fully_async_main.py

@hydra.main(config_path="config", config_name="fully_async_ppo_trainer", version_base=None)
def main(config):
    from verl.trainer.main_ppo import run_ppo

    # Ensure async training config exists
    if not hasattr(config, "async_training"):
        raise RuntimeError("must set async_training config")

    from time import time

    start_time = time()
    auto_set_device(config)  
    # TODO: unify rollout config with actor_rollout_ref
    config.actor_rollout_ref.rollout.nnodes = config.rollout.nnodes
    config.actor_rollout_ref.rollout.n_gpus_per_node = config.rollout.n_gpus_per_node
    config = migrate_legacy_reward_impl(config)
    run_ppo(config, task_runner_class=FullyAsyncTaskRunner)  
    print(f"total time: {time() - start_time:.2f} seconds")


if __name__ == "__main__":
    main()

trainer.main_ppo

提示

代码

• main_ppo.py
• 传入 task_runner_class，可自定义

def run_ppo(config, task_runner_class=None, wait_for_task=None) -> None: 
    """Initialize Ray cluster and run distributed PPO training process.

    Args:
        config: Training configuration object containing all necessary parameters
                for distributed PPO training including Ray initialization settings,
                model paths, and training hyperparameters.
        task_runner_class: For recipe to change TaskRunner.
        wait_for_task: Optional recipe callback waiting for the run ObjectRef.
    """
    # Check if Ray is not initialized
    if not ray.is_initialized():
        # Initialize Ray with a local cluster configuration
        # Set environment variables in the runtime environment to control tokenizer parallelism,
        # NCCL debug level, VLLM logging level, and allow runtime LoRA updating
        # `num_cpus` specifies the number of CPU cores Ray can use, obtained from the configuration
        default_runtime_env = get_ppo_ray_runtime_env()
        ray_init_kwargs = config.ray_kwargs.get("ray_init", {})
        runtime_env_kwargs = ray_init_kwargs.get("runtime_env", {})

        if config.transfer_queue.enable:
            # Add runtime environment variables for transfer queue
            runtime_env_vars = runtime_env_kwargs.get("env_vars", {})
            runtime_env_vars["TRANSFER_QUEUE_ENABLE"] = "1"
            runtime_env_kwargs["env_vars"] = runtime_env_vars

        runtime_env = OmegaConf.merge(default_runtime_env, runtime_env_kwargs)
        ray_init_kwargs = OmegaConf.create({**ray_init_kwargs, "runtime_env": runtime_env})
        print(f"ray init kwargs: {ray_init_kwargs}")
        ray.init(**OmegaConf.to_container(ray_init_kwargs))

    if task_runner_class is None:
        task_runner_class = ray.remote(num_cpus=1)(TaskRunner)  # please make sure main_task is not scheduled on head

    # Create a remote instance of the TaskRunner class, and
    # Execute the `run` method of the TaskRunner instance remotely and wait for it to complete
    if (
        is_cuda_available
        and config.global_profiler.tool == "nsys"
        and config.global_profiler.get("steps") is not None
        and len(config.global_profiler.get("steps", [])) > 0
    ):
        from verl.utils.import_utils import is_nvtx_available

        assert is_nvtx_available(), "nvtx is not available in CUDA platform. Please 'pip3 install nvtx'"
        nsight_options = OmegaConf.to_container(
            config.global_profiler.global_tool_config.nsys.controller_nsight_options
        )
        runner = task_runner_class.options(runtime_env={"nsight": nsight_options}).remote()  
    else:
        runner = task_runner_class.remote()
    task = runner.run.remote(config)  
    try:
        (wait_for_task or ray.get)(task)
    except BaseException:
        # The driver must exit on failure, even when the runner is blocked in a
        # distributed call. Ray then tears down this job's non-detached workers.
        try:
            ray.kill(runner, no_restart=True)
        except Exception:
            pass  # Preserve the original failure if GCS is also unavailable.
        raise

    # [Optional] get the path of the timeline trace file from the configuration, default to None
    # This file is used for performance analysis
    timeline_json_file = config.ray_kwargs.get("timeline_json_file", None)
    if timeline_json_file:
        ray.timeline(filename=timeline_json_file)

SeparateRayPPOTrainer (继承自RayPPOTrainer)

提示

代码

• experimental/separation/ray_trainer.py

初始化

class SeparateRayPPOTrainer(RayPPOTrainer):
  """
  Support for the initialization and fit process of Ray Trainer in the resource-separated scenario:
      - Fully async policy
      - One-step off-policy
  """

  def __init__(
      self,
      config,
      tokenizer,
      role_worker_mapping: dict[Role, WorkerType],
      resource_pool_manager: ResourcePoolManager,
      ray_worker_group_cls: type[RayWorkerGroup] = RayWorkerGroup,
      processor=None,
      reward_fn=None,
      val_reward_fn=None,
      train_dataset: Optional[Dataset] = None,
      val_dataset: Optional[Dataset] = None,
      collate_fn=None,
      train_sampler: Optional[Sampler] = None,
      device_name=None,
  ):
      super().__init__(
          config,
          tokenizer,
          role_worker_mapping,
          resource_pool_manager,
          ray_worker_group_cls,
          processor,
          train_dataset,
          val_dataset,
          collate_fn,
          train_sampler,
          device_name,
      )
      self.global_steps = 0
      self.epoch = 0
      self.max_steps_duration = 0
      self.progress_bar = None
      self.logger = None
      self.is_last_step = False
      self.prev_step_profile = False
      self.curr_step_profile = False
      self.next_step_profile = False
      self.last_val_metrics = {}
      self.metrics = {}
      self.timing_raw = {}
      # reward message
      self.reward_tensor = None
      self.reward_extra_infos_dict = {}
      self.checkpoint_manager = None

  def init_workers(self):
      """Initialize distributed training workers using Ray backend.

      Creates:
      1. Ray resource pools from configuration
      2. Worker groups for each role (actor, critic, etc.)
      """
      self._init_resource_pools()
      self._create_worker_classes() 
      self._init_worker_groups()
      self._init_models() 
      self._init_reward_loop()
      self._init_async_rollout_manager() 

      # Support custom CheckpointEngineManager via config
      checkpoint_manager_class_fqn = self.config.actor_rollout_ref.rollout.get("checkpoint_manager_class")
      if checkpoint_manager_class_fqn:
          CheckpointEngineManager = load_class_from_fqn(checkpoint_manager_class_fqn, "CheckpointEngineManager")
      else:
          from verl.checkpoint_engine import CheckpointEngineManager

      self.checkpoint_manager = CheckpointEngineManager(
          config=omega_conf_to_dataclass(self.config.actor_rollout_ref.rollout.checkpoint_engine),
          actor_wg=self.actor_rollout_wg,
          replicas=self.llm_server_manager.get_replicas(),
      )

  def _init_resource_pools(self):
      self.resource_pool_manager.create_resource_pool()
      self.resource_pool_to_cls = {pool: {} for pool in self.resource_pool_manager.resource_pool_dict.values()}

  def _create_worker_classes(self):
      self._create_actor_rollout_classes()
      self._create_critic_class()
      self._create_reference_policy_class()
      self._create_reward_model_class()

  def _create_actor_rollout_classes(self):  
      raise NotImplementedError

  def _create_critic_class(self):
      # create critic
      if self.use_critic:
          resource_pool = self.resource_pool_manager.get_resource_pool(Role.Critic)
          critic_cfg = omega_conf_to_dataclass(self.config.critic)

          # convert critic_cfg into TrainingWorkerConfig for the unified model engine worker
          from verl.workers.config import FSDPEngineConfig
          from verl.workers.engine_workers import TrainingWorkerConfig

          self.orig_critic_cfg = critic_cfg
          if self.orig_critic_cfg.strategy == "fsdp":
              engine_config: FSDPEngineConfig = self.orig_critic_cfg.model.fsdp_config 
              engine_config.infer_max_token_len_per_gpu = critic_cfg.ppo_infer_max_token_len_per_gpu
              engine_config.max_token_len_per_gpu = critic_cfg.ppo_max_token_len_per_gpu
          else:
              raise NotImplementedError(f"Unknown strategy {self.orig_critic_cfg.strategy=}")

          # Wire the critic profiler config via the hydra path (real dataclass tool_config), so the
          # standalone critic TrainingWorker gets a working DistProfiler instead of a silent no-op.
          critic_omega_profiler_config = self.config.critic.get("profiler", {})
          critic_profiler_config = (
              omega_conf_to_dataclass(critic_omega_profiler_config) if critic_omega_profiler_config else None
          )

          critic_cfg = TrainingWorkerConfig( 
              model_type="value_model",
              model_config=self.orig_critic_cfg.model,
              engine_config=engine_config,
              optimizer_config=self.orig_critic_cfg.optim,
              checkpoint_config=self.orig_critic_cfg.checkpoint,
              profiler_config=critic_profiler_config,
          )

          critic_cls = RayClassWithInitArgs(cls=self.role_worker_mapping[Role.Critic], config=critic_cfg)
          self.resource_pool_to_cls[resource_pool][str(Role.Critic)] = critic_cls

  def _create_reference_policy_class(self):
      # create reference policy if needed
      if self.use_reference_policy:
          resource_pool = self.resource_pool_manager.get_resource_pool(Role.RefPolicy)
          ref_policy_cls = RayClassWithInitArgs(
              self.role_worker_mapping[Role.RefPolicy],
              config=self.config.actor_rollout_ref,
              role=str(Role.RefPolicy),
              # profile_option=self.config.trainer.npu_profile.options,
          )
          self.resource_pool_to_cls[resource_pool][str(Role.RefPolicy)] = ref_policy_cls

  def _create_reward_model_class(self):
      # create a reward model if reward_fn is None
      if self.use_rm:
          # we create a RM here
          resource_pool = self.resource_pool_manager.get_resource_pool(Role.RewardModel)
          rm_cls = RayClassWithInitArgs(
              self.role_worker_mapping[Role.RewardModel], config=self.config.reward.reward_model
          )
          self.resource_pool_to_cls[resource_pool][str(Role.RewardModel)] = rm_cls

  def _init_worker_groups(self):
      # initialize WorkerGroup
      # NOTE: if you want to use a different resource pool for each role, which can support different parallel size,
      # you should not use `create_colocated_worker_cls`.
      # Instead, directly pass different resource pool to different worker groups.
      # See https://github.com/verl-project/verl/blob/master/examples/tutorial/ray/tutorial.ipynb
      # for more information.
      all_wg = {}
      wg_kwargs = {}  # Setting up kwargs for RayWorkerGroup
      if OmegaConf.select(self.config.trainer, "ray_wait_register_center_timeout") is not None:
          wg_kwargs["ray_wait_register_center_timeout"] = self.config.trainer.ray_wait_register_center_timeout
      if OmegaConf.select(self.config.global_profiler, "steps") is not None:
          wg_kwargs["profile_steps"] = OmegaConf.select(self.config.global_profiler, "steps")
          # Only require nsight worker options when tool is nsys
          if OmegaConf.select(self.config.global_profiler, "tool") == "nsys":
              assert (
                  OmegaConf.select(self.config.global_profiler.global_tool_config.nsys, "worker_nsight_options")
                  is not None
              ), "worker_nsight_options must be set when using nsys with profile_steps"
              wg_kwargs["worker_nsight_options"] = OmegaConf.to_container(
                  OmegaConf.select(self.config.global_profiler.global_tool_config.nsys, "worker_nsight_options")
              )
      wg_kwargs["device_name"] = self.device_name

      for resource_pool, class_dict in self.resource_pool_to_cls.items():
          worker_dict_cls = create_colocated_worker_cls(class_dict=class_dict)
          wg_dict = self.ray_worker_group_cls(
              resource_pool=resource_pool,
              ray_cls_with_init=worker_dict_cls,
              **wg_kwargs,
          )
          spawn_wg = wg_dict.spawn(prefix_set=class_dict.keys())
          all_wg.update(spawn_wg)
      self.all_wg = all_wg

  def _init_models(self):
      if self.use_critic:
          self.critic_wg = self.all_wg[str(Role.Critic)]
          self.critic_wg.reset()
          # assign critic loss
          from functools import partial

          from verl.workers.utils.losses import value_loss

          value_loss_ = partial(value_loss, config=self.orig_critic_cfg)
          self.critic_wg.set_loss_fn(value_loss_)

      if self.use_reference_policy and not self.ref_in_actor:
          self.ref_policy_wg = self.all_wg[str(Role.RefPolicy)]
          self.ref_policy_wg.init_model()

      if self.use_rm:
          self.rm_wg = self.all_wg[str(Role.RewardModel)]
          self.rm_wg.init_model()

      # we should create rollout at the end so that vllm can have a better estimation of kv cache memory
      self.actor_rollout_wg = self.all_wg[str(Role.ActorRollout)]
      self.actor_rollout_wg.init_model()

  def _init_reward_loop(self):
      from verl.experimental.reward_loop import RewardLoopManager

      # initalize reward loop manager
      # reward model (colocate or standalone): get resource_pool
      # no reward model: resource_pool = None
      resource_pool = self.resource_pool_manager.get_resource_pool(Role.RewardModel) if self.use_rm else None
      self.reward_loop_manager = RewardLoopManager(
          config=self.config,
          rm_resource_pool=resource_pool,
      )

  def _init_async_rollout_manager(self):
      pass

fit 主流程

class SeparateRayPPOTrainer(RayPPOTrainer):
  
  # ... 其他代码	
  
  def fit(self):
    """
    The training loop of PPO.
    The driver process only need to call the compute functions of the worker group through RPC
    to construct the PPO dataflow.
    The light-weight advantage computation is done on the driver process.

    !!!
    The logic of fit is consistent with that of fit_refactor;
    if any modifications are made, apply them to both methods simultaneously.
    """
    from omegaconf import OmegaConf

    from verl.utils.tracking import Tracking

    self.logger = Tracking(
        project_name=self.config.trainer.project_name,
        experiment_name=self.config.trainer.experiment_name,
        default_backend=self.config.trainer.logger,
        config=OmegaConf.to_container(self.config, resolve=True),
    )

    self.global_steps = 0

    # load checkpoint and update weights before doing anything
    self._load_checkpoint() 
    self.checkpoint_manager.update_weights(self.global_steps)

    current_epoch = self.global_steps // len(self.train_dataloader)

    # perform validation before training
    # currently, we only support validation using the reward_function.
    if self.config.trainer.get("val_before_train", True):
        val_metrics = self._validate() 
        assert val_metrics, f"{val_metrics=}"
        pprint(f"Initial validation metrics: {val_metrics}")
        self.logger.log(data=val_metrics, step=self.global_steps)
        if self.config.trainer.get("val_only", False):
            return

    # add tqdm
    self.progress_bar = tqdm(total=self.total_training_steps, initial=self.global_steps, desc="Training Progress")

    # we start from step 1
    self.global_steps += 1
    self.last_val_metrics = None
    self.max_steps_duration = 0

    self.prev_step_profile = False
    self.curr_step_profile = (
        self.global_steps in self.config.global_profiler.steps
        if self.config.global_profiler.steps is not None
        else False
    )
    self.next_step_profile = False

    for epoch in range(current_epoch, self.config.trainer.total_epochs): 
        for batch_dict in self.train_dataloader: 
            self.epoch = epoch
            self.fit_step(batch_dict) 
            if self.is_last_step:
                return

fit_step

核心流程

1.Pre-step processing

2.Get batch

3. Generate sequences

4. Compute reward

5. Compute log_prob

6. Compute ref_log_prob

7. Compute advantage

8. Update critic

9. Update actor

10. Post-step

class SeparateRayPPOTrainer(RayPPOTrainer):
  
  # 其他代码....
  
  def fit_step(self, batch_dict: Any = None):
    """
    Single-step training template method. Handles all logic for one training step.

    Flow:
    1. Pre-step processing -> 2. Get batch -> 3. Generate sequences ->
    4. Compute reward -> 5. Compute log_prob -> 6. Compute reward ->
    7. Compute advantage -> 8. Update critic -> 9. Update actor -> 10. Post-step processing

    Args:
        batch_dict: Raw data dictionary
    """
    self.metrics = {"training/global_step": self.global_steps, "training/epoch": self.epoch}
    self.timing_raw = {}
    # reward message
    self.reward_tensor = None
    self.reward_extra_infos_dict = {}

    self._fit_prepare_step()
    self._fit_start_profile()

    with marked_timer("step", self.timing_raw):
        batch = self._fit_get_batch(batch_dict)
        batch = self._fit_generate(batch)
        batch = self._fit_compute_reward(batch)
        batch = self._fit_compute_log_prob(batch)
        batch = self._fit_compute_ref_log_prob(batch)
        batch = self._fit_compute_critic(batch)
        batch = self._fit_compute_advantage(batch)
        batch = self._fit_update_critic(batch)
        batch = self._fit_update_actor(batch)
        self._fit_update_weights()
        self._fit_dump_data(batch)

    self._fit_validate()
    self._fit_save_checkpoint()
    self._fit_stop_profile()
    self._fit_collect_metrics(batch)
    self._fit_experimental(batch)
    self._fit_postprocess_step()

Get batch

def _fit_get_batch(self, batch_dict: dict) -> DataProto:
  	# 获取 batch
    batch = DataProto.from_single_dict(batch_dict) 
    batch.meta_info["temperature"] = self.config.actor_rollout_ref.rollout.temperature
    # add uid
    batch.non_tensor_batch["uid"] = np.array([str(uuid.uuid4()) for _ in range(len(batch.batch))], dtype=object)
    return batch

Generate Sequences

def _fit_generate(self, batch: DataProto = None) -> DataProto:
    metrics = self.metrics
    timing_raw = self.timing_raw
    gen_batch = self._get_gen_batch(batch)
    # pass global_steps to trace
    gen_batch.meta_info["global_steps"] = self.global_steps
    # 重复n次
    rollout_n = self.config.actor_rollout_ref.rollout.n 
    gen_batch_output = gen_batch.repeat(repeat_times=rollout_n, interleave=True) 

    if self.config.algorithm.adv_estimator == AdvantageEstimator.REMAX:
        # pass
    else:
        combined_gen_batch = gen_batch_output
        num_sampled_prompts = len(gen_batch_output)

    with marked_timer("gen", timing_raw, color="red"):
        if self.curr_step_profile:
            self.async_rollout_manager.start_profile(global_step=self.global_steps)
        # 调用 agentloop 生成一批轨迹
        combined_gen_output = self.async_rollout_manager.generate_sequences(combined_gen_batch)  
        self.checkpoint_manager.sleep_replicas()
        if self.curr_step_profile:
            self.llm_server_manager.stop_profile()

        timing_raw.update(combined_gen_output.meta_info["timing"])
        combined_gen_output.meta_info.pop("timing", None)

    gen_batch_output = combined_gen_output.slice(0, num_sampled_prompts)
    if "__do_sample__" in gen_batch_output.non_tensor_batch:
        gen_batch_output.pop(non_tensor_batch_keys=["__do_sample__"])

    if self.config.algorithm.adv_estimator == AdvantageEstimator.REMAX:
        # pass
    del combined_gen_batch, combined_gen_output
    # 拼接回去，repeat to align with repeated responses in rollout
    batch = batch.repeat(repeat_times=self.config.actor_rollout_ref.rollout.n, interleave=True) 
    # 合并gen_batch_output里的新字段
    batch = batch.union(gen_batch_output) 

    if "response_mask" not in batch.batch.keys():
        batch.batch["response_mask"] = compute_response_mask(batch)
    # Balance the number of valid tokens across DP ranks.
    # NOTE: This usually changes the order of data in the `batch`,
    # which won't affect the advantage calculation (since it's based on uid),
    # but might affect the loss calculation (due to the change of mini-batching).
    if self.config.trainer.balance_batch:
        self._balance_batch(batch, metrics=metrics)

    # compute global_valid tokens
    batch.meta_info["global_token_num"] = torch.sum(batch.batch["attention_mask"], dim=-1).tolist()
    # get images_seqlens
    images_seqlens_all = []
    for multi_modal_input in batch.non_tensor_batch["multi_modal_inputs"]:
        if "image_grid_thw" not in multi_modal_input.keys():
            continue
        images_seqlens_all.extend(multi_modal_input["images_seqlens"].tolist())
    batch.meta_info["images_seqlens"] = images_seqlens_all
    return batch

Compute reward

检查是否有rm_scores，如果没有，则调用reward_loop_manager去计算reward。

def _fit_compute_reward(self, batch: DataProto) -> DataProto:
    timing_raw = self.timing_raw
    with marked_timer("reward", timing_raw, color="yellow"):
        # compute reward model score
        if self.use_rm and "rm_scores" not in batch.batch.keys(): 
            batch_reward = self._compute_reward_colocate(batch) 
            batch = batch.union(batch_reward)

        # Compute or extract reward_tensor and reward_extra_infos_dict for training
        reward_tensor, reward_extra_infos_dict = extract_reward(batch) 
        self.reward_tensor = reward_tensor
        self.reward_extra_infos_dict = reward_extra_infos_dict
    return batch
  

def _compute_reward_colocate(self, batch: DataProto) -> tuple[torch.Tensor, dict[str, Any]] | torch.Tensor:
    """
    compute reward use colocate reward model
    """
    assert self.reward_loop_manager is not None, "RewardLoopManager is None"
    batch_reward = self.reward_loop_manager.compute_rm_score(batch) 
    return batch_reward

Compute log_prob

compute log_prob

Bypass Mode

• 不做重算，直接利用rollout_log_probs 作为old_log_probs
• 默认为False

Decoupled Mode

• 把已生成token序列，输入到当前Actor，重计算old_log_probs
  • old_log_probs：对于已发生动作，训练开始前的actor会给出多大概率？
  • old：相比于接下来的训练更新，它是旧的。

Routing Replay

• R2：训练重放 old policy 的路由
  • 计算 old_log_probs → 记录 [2, 9]
  • 更新 actor → 回放 [2, 9]
• Rollout Routing Replay：训练重放 rollout时的路由
  • rollout → 记录 [2, 7]
  • 计算 old_log_probs → 回放 [2, 7]
  • 更新 actor → 回放 [2, 7]

def _fit_compute_log_prob(self, batch: DataProto) -> DataProto:
    metrics = self.metrics
    timing_raw = self.timing_raw
    # Operating Mode Selection:
    # - Bypass mode: Sets old_log_probs = rollout_log_probs (2 policies: π_rollout, π_θ)
    # - Decoupled mode: Recomputes old_log_probs as proximal anchor (3 policies: π_rollout, π_old, π_θ)
    #   Note: π_old computed once per data batch, serves as stable reference during mini-batch updates
    rollout_corr_config = self.config.algorithm.get("rollout_correction", None)
    bypass_recomputing_logprobs = rollout_corr_config and rollout_corr_config.get("bypass_mode", False)
    if bypass_recomputing_logprobs:  # Use `rollout_log_probs`
        from verl.trainer.ppo.rollout_corr_helper import apply_bypass_mode

        apply_bypass_mode(
            batch=batch,
            rollout_corr_config=rollout_corr_config,
            policy_loss_config=self.config.actor_rollout_ref.actor.policy_loss,
        )
    else:  # Recompute old_log_probs
        with marked_timer("old_log_prob", timing_raw, color="blue"):
          	# 把rollout的token序列，输入到当前actor，重算 old_log_probs
            # 对于已发生的动作，训练开始前的actor会给出多大的概率？
            # old？指相当于接下来要做的更新，它是旧的
            old_log_prob, old_log_prob_mfu = self._compute_old_log_prob(batch) 
            entropys = old_log_prob.batch["entropys"]
            response_masks = batch.batch["response_mask"]
            actor_config = self.config.actor_rollout_ref.actor
            entropy_agg = agg_loss(
                loss_mat=entropys,
                loss_mask=response_masks,
                loss_agg_mode=actor_config.loss_agg_mode,
                loss_scale_factor=actor_config.loss_scale_factor,
            )
            old_log_prob_metrics = {
                "actor/entropy": entropy_agg.detach().item(),
                "perf/mfu/actor_infer": old_log_prob_mfu,
            }
            metrics.update(old_log_prob_metrics)
            old_log_prob.batch.pop("entropys")
            if "routed_experts" in batch.batch and "routed_experts" in old_log_prob.batch: 
                actor_cfg = self.config.actor_rollout_ref.actor
                if getattr(actor_cfg, "strategy", None) == "megatron":
                    engine_cfg = getattr(actor_cfg, "megatron", None)
                elif getattr(actor_cfg, "strategy", None) == "veomni":
                    engine_cfg = getattr(actor_cfg, "veomni", None)
                else:
                    engine_cfg = None
                router_mode = getattr(getattr(engine_cfg, "router_replay", None), "mode", "disabled")
                if router_mode == "R2":
                  	# R2，使用old里的路由回放，batch 原有的路由
                    batch.batch.pop("routed_experts") 
                else:
                  	# 其他，保留batch 原有路由，R3，不用old策略里的，使用rollout里的路由回放
                    old_log_prob.batch.pop("routed_experts") 
            batch = batch.union(old_log_prob)
            if "rollout_log_probs" in batch.batch.keys():
                # TODO: we may want to add diff of probs too.
                from verl.utils.debug.metrics import calculate_debug_metrics

                metrics.update(calculate_debug_metrics(batch))

    assert "old_log_probs" in batch.batch, f'"old_log_prob" not in {batch.batch.keys()=}'
    return batch

ppo/ray_trainer，Compute LogProb (最核心，所有内容都在这)

def _compute_old_log_prob(self, batch: DataProto):
    # TODO: remove step 1, 2, 4 after we make the whole training tensordict and padding free
    # step 1: convert dataproto to tensordict.
    batch_td = batch.to_tensordict()
    # step 2: convert from padding to nopadding
    batch_td = left_right_2_no_padding(batch_td)
    # step 3: add meta info
    calculate_sum_pi_squared = self.config.actor_rollout_ref.actor.get("calculate_sum_pi_squared", False)
    tu.assign_non_tensor(
        batch_td,
        calculate_entropy=True,
        calculate_sum_pi_squared=calculate_sum_pi_squared,
        compute_loss=False,
    )
    # 计算 log_prob
    output = self.actor_rollout_wg.compute_log_prob(batch_td) 
    # gather output
    entropy = tu.get(output, "entropy")
    log_probs = tu.get(output, "log_probs")
    routed_experts = tu.get(output, "routed_experts")
    sum_pi_squared = tu.get(output, "sum_pi_squared") if calculate_sum_pi_squared else None

    old_log_prob_mfu = tu.get(output, "metrics")["mfu"]
    # step 4. No padding to padding
    entropy = no_padding_2_padding(entropy, batch_td)
    log_probs = no_padding_2_padding(log_probs, batch_td)
    if sum_pi_squared is not None:
        sum_pi_squared = no_padding_2_padding(sum_pi_squared, batch_td)
    # step 5: rebuild a tensordict and convert to dataproto
    result = {"old_log_probs": log_probs.float(), "entropys": entropy.float()} 
    if routed_experts is not None:
        result["routed_experts"] = routed_experts
    if sum_pi_squared is not None:
        result["sum_pi_squared"] = sum_pi_squared.float()
    old_log_prob = tu.get_tensordict(result)
    old_log_prob = DataProto.from_tensordict(old_log_prob)
    return old_log_prob, old_log_prob_mfu

Compute ref_log_prob

def _fit_compute_ref_log_prob(self, batch: DataProto) -> DataProto:
  timing_raw = self.timing_raw
  if self.use_reference_policy:
      with marked_timer(str(Role.RefPolicy), timing_raw, color="olive"):
          ref_log_prob = self._compute_ref_log_prob(batch)  
          batch = batch.union(ref_log_prob)
  return batch

RayPPOTrainer 基类里的函数，具体里面就得去看 这里的笔记： Compute LogProb (最核心,所有内容都在这)

def _compute_ref_log_prob(self, batch: DataProto) -> DataProto:
    # step 1: convert dataproto to tensordict.
    batch_td = batch.to_tensordict()
    # step 2: convert from padding to nopadding
    batch_td = left_right_2_no_padding(batch_td)
    # step 3: add meta info
    metadata = {"calculate_entropy": False, "compute_loss": False}
    if self.ref_in_actor:
        metadata["no_lora_adapter"] = True
    tu.assign_non_tensor(batch_td, **metadata)
    if self.ref_in_actor:
      	# LORA
        output = self.actor_rollout_wg.compute_log_prob(batch_td)
    else:
      	# 保留冻结的reference model，保持冻结
        output = self.ref_policy_wg.compute_ref_log_prob(batch_td)  
    # gather output
    log_probs = tu.get(output, "log_probs")
    # step 4. No padding to padding
    log_probs = no_padding_2_padding(log_probs, batch_td)
    # step 5: rebuild a tensordict and convert to dataproto
    ref_log_prob = tu.get_tensordict({"ref_log_prob": log_probs.float()})
    ref_log_prob = DataProto.from_tensordict(ref_log_prob)

    return ref_log_prob

Compute Advantage

具体就得进这里去看了：core_algos优势计算主入口

def _fit_compute_advantage(self, batch) -> DataProto:
    metrics = self.metrics
    timing_raw = self.timing_raw
    reward_tensor = self.reward_tensor
    reward_extra_infos_dict = self.reward_extra_infos_dict

    with marked_timer("adv", timing_raw, color="brown"):
        # we combine with rule-based rm
        reward_extra_infos_dict: dict[str, list]
        batch.batch["token_level_scores"] = reward_tensor

        if reward_extra_infos_dict:
            batch.non_tensor_batch.update({k: np.array(v) for k, v in reward_extra_infos_dict.items()})

        # compute rewards. apply_kl_penalty if available
        if self.config.algorithm.use_kl_in_reward:
            batch, kl_metrics = apply_kl_penalty(
                batch, kl_ctrl=self.kl_ctrl_in_reward, kl_penalty=self.config.algorithm.kl_penalty
            )
            metrics.update(kl_metrics)
        else:
            batch.batch["token_level_rewards"] = batch.batch["token_level_scores"]

        # Compute rollout correction: IS weights, rejection sampling, and metrics
        # Only runs in decoupled mode (computes once per batch using stable π_old)
        # In bypass mode, this is skipped - actor computes metrics from evolving π_θ vs π_rollout
        rollout_corr_config = self.config.algorithm.get("rollout_correction", None)
        bypass_recomputing_logprobs = rollout_corr_config and rollout_corr_config.get("bypass_mode", False)
        if (
            rollout_corr_config is not None
            and "rollout_log_probs" in batch.batch
            and not bypass_recomputing_logprobs  # Only in decoupled mode
        ):
            from verl.trainer.ppo.rollout_corr_helper import compute_rollout_correction_and_add_to_batch

            # Compute IS weights, apply rejection sampling, compute metrics
            # 计算了 rollout IS 权重
            batch, is_metrics = compute_rollout_correction_and_add_to_batch(batch, rollout_corr_config)  
            # IS and off-policy metrics already have rollout_corr/ prefix
            metrics.update(is_metrics)

        # compute advantages, executed on the driver process
        norm_adv_by_std_in_grpo = self.config.algorithm.get(
            "norm_adv_by_std_in_grpo", True
        )  # GRPO adv normalization factor

        batch = compute_advantage(  
            batch,
            adv_estimator=self.config.algorithm.adv_estimator,
            gamma=self.config.algorithm.gamma,
            lam=self.config.algorithm.lam,
            num_repeat=self.config.actor_rollout_ref.rollout.n,
            norm_adv_by_std_in_grpo=norm_adv_by_std_in_grpo,
            config=self.config.algorithm,
        )
    return batch

Update critic

def _fit_update_critic(self, batch: DataProto) -> DataProto:
    metrics = self.metrics
    timing_raw = self.timing_raw
    if self.use_critic:
        with marked_timer("update_critic", timing_raw, color="pink"):
            critic_output = self._update_critic(batch)
        critic_output_metrics = reduce_metrics(critic_output.meta_info["metrics"])
        metrics.update(critic_output_metrics)
    return batch

def _update_critic(self, batch: DataProto) -> DataProto:
    batch_td = batch.to_tensordict()
    # step 2: convert from padding to no-padding
    batch_td = left_right_2_no_padding(batch_td)
    ppo_mini_batch_size = self.config.critic.ppo_mini_batch_size
    ppo_mini_batch_size = ppo_mini_batch_size * self.config.actor_rollout_ref.rollout.n  
    ppo_epochs = self.config.critic.ppo_epochs
    seed = self.config.critic.data_loader_seed
    shuffle = self.config.critic.shuffle
    tu.assign_non_tensor(
        batch_td,
        global_batch_size=ppo_mini_batch_size,
        mini_batch_size=ppo_mini_batch_size,
        epochs=ppo_epochs,
        seed=seed,
        dataloader_kwargs={"shuffle": shuffle},
    )

    output = self.critic_wg.train_mini_batch(batch_td)  
    output = output.get()
    output = tu.get(output, "metrics")
    output = rename_dict(output, "critic/")
    # modify key name
    output["perf/mfu/critic"] = output.pop("critic/mfu")
    critic_output = DataProto.from_single_dict(data={}, meta_info={"metrics": output})
    return critic_output

Update Actor

def _fit_update_actor(self, batch: DataProto) -> DataProto:
    metrics = self.metrics
    timing_raw = self.timing_raw
    # implement critic warmup
    if self.config.trainer.critic_warmup <= self.global_steps:
        # update actor
        with marked_timer("update_actor", timing_raw, color="red"):
            actor_output = self._update_actor(batch)

        actor_output_metrics = reduce_metrics(actor_output.meta_info["metrics"])
        metrics.update(actor_output_metrics)
    return batch

def _update_actor(self, batch: DataProto) -> DataProto:
    rollout_config = self.config.actor_rollout_ref.rollout
    batch.meta_info["multi_turn"] = rollout_config.multi_turn.enable
    # TODO: Make "temperature" single source of truth from generation.
    batch.meta_info["temperature"] = rollout_config.temperature
    # update actor
    batch_td = batch.to_tensordict()
    # step 2: convert from padding to no-padding
    batch_td = left_right_2_no_padding(batch_td)
    calculate_entropy = self.config.actor_rollout_ref.actor.calculate_entropy or (
        self.config.actor_rollout_ref.actor.entropy_coeff != 0.0
    )
    distillation_use_topk = (
        self.distillation_config.distillation_loss.loss_settings.use_topk
        if is_distillation_enabled(self.config.get("distillation"))
        else False
    )
    ppo_mini_batch_size = self.config.actor_rollout_ref.actor.ppo_mini_batch_size
    ppo_mini_batch_size = ppo_mini_batch_size * self.config.actor_rollout_ref.rollout.n  
    ppo_epochs = self.config.actor_rollout_ref.actor.ppo_epochs
    seed = self.config.actor_rollout_ref.actor.data_loader_seed
    shuffle = self.config.actor_rollout_ref.actor.shuffle
    tu.assign_non_tensor(
        batch_td,
        calculate_entropy=calculate_entropy,
        distillation_use_topk=distillation_use_topk,
        global_batch_size=ppo_mini_batch_size,
        mini_batch_size=ppo_mini_batch_size,
        epochs=ppo_epochs,
        seed=seed,
        dataloader_kwargs={"shuffle": shuffle},
        compute_loss=True,
    )
    actor_output = self.actor_rollout_wg.update_actor(batch_td)  
    actor_output = tu.get(actor_output, "metrics")
    actor_output = rename_dict(actor_output, "actor/")
    # modify key name
    actor_output["perf/mfu/actor"] = actor_output.pop("actor/mfu")
    actor_output = DataProto.from_single_dict(data={}, meta_info={"metrics": actor_output})

    return actor_output

fit_validate

def _fit_validate(self):
  metrics = self.metrics
  timing_raw = self.timing_raw
  if self.config.trainer.test_freq > 0 and (
      self.is_last_step or self.global_steps % self.config.trainer.test_freq == 0
  ):
      with marked_timer("testing", timing_raw, color="green"):
          val_metrics: dict = self._validate()  
          if self.is_last_step:
              self.last_val_metrics = val_metrics
      metrics.update(val_metrics)

def _validate(self, merged: bool = False):
    data_source_lst = []
    reward_extra_infos_dict: dict[str, list] = defaultdict(list)

    # Lists to collect samples for the table
    sample_inputs = []
    sample_outputs = []
    sample_gts = []
    sample_scores = []
    sample_turns = []
    sample_uids = []

    for test_data in self.val_dataloader:  
        test_batch = DataProto.from_single_dict(test_data)

        if "uid" not in test_batch.non_tensor_batch:
            test_batch.non_tensor_batch["uid"] = np.array(
                [str(uuid.uuid4()) for _ in range(len(test_batch.batch))], dtype=object
            )

        # repeat test batch
        test_batch = test_batch.repeat(
            repeat_times=self.config.actor_rollout_ref.rollout.val_kwargs.n, interleave=True
        )

        ground_truths = [
            item.non_tensor_batch.get("reward_model", {}).get("ground_truth", None) for item in test_batch
        ]
        sample_gts.extend(ground_truths)

        test_gen_batch = self._get_gen_batch(test_batch)
        test_gen_batch.meta_info = {
            "eos_token_id": self.tokenizer.eos_token_id,
            "pad_token_id": self.tokenizer.pad_token_id,
            "recompute_log_prob": False,
            "do_sample": self.config.actor_rollout_ref.rollout.val_kwargs.do_sample,
            "validate": True,
            "global_steps": self.global_steps,
        }
        print(f"test_gen_batch meta info: {test_gen_batch.meta_info}")

        # pad to be divisible by dp_size
        size_divisor = self.config.actor_rollout_ref.rollout.agent.num_workers
        test_gen_batch_padded, pad_size = pad_dataproto_to_divisor(test_gen_batch, size_divisor)
        test_output_gen_batch_padded = self.async_rollout_manager.generate_sequences(test_gen_batch_padded)

        if self.use_rm and "rm_scores" not in test_output_gen_batch_padded.batch.keys():
            # for colocate reward models, we need to sleep rollout model
            # to spare GPU memory for reward model
            self.checkpoint_manager.sleep_replicas()
            batch_reward = self._compute_reward_colocate(test_output_gen_batch_padded)
            test_output_gen_batch_padded = test_output_gen_batch_padded.union(batch_reward)
            # wake up rollout model
            # replace with wake_up method once supported
            self.checkpoint_manager.update_weights(self.global_steps)

        # unpad
        test_output_gen_batch = unpad_dataproto(test_output_gen_batch_padded, pad_size=pad_size)

        print("validation generation end")

        # Store generated outputs
        output_ids = test_output_gen_batch.batch["responses"]
        output_texts = [self.tokenizer.decode(ids, skip_special_tokens=True) for ids in output_ids]
        sample_outputs.extend(output_texts)

        test_batch = test_batch.union(test_output_gen_batch)
        test_batch.meta_info["validate"] = True

        # Store original inputs
        input_ids = test_batch.batch["prompts"]
        # TODO: Can we keep special tokens except for padding tokens?
        input_texts = [self.tokenizer.decode(ids, skip_special_tokens=True) for ids in input_ids]
        sample_inputs.extend(input_texts)
        sample_uids.extend(test_batch.non_tensor_batch["uid"])

        # evaluate using reward_function
        reward_tensor, reward_extra_info = extract_reward(test_batch)

        scores = reward_tensor.sum(-1).cpu().tolist()
        sample_scores.extend(scores)

        reward_extra_infos_dict["reward"].extend(scores)
        for key, values in reward_extra_info.items():
            if key not in reward_extra_infos_dict:
                reward_extra_infos_dict[key] = []
            if isinstance(values, np.ndarray):
                reward_extra_infos_dict[key].extend(values.tolist())
            else:
                reward_extra_infos_dict[key].extend(values if isinstance(values, list) else [values])

        # collect num_turns of each prompt
        if "__num_turns__" in test_batch.non_tensor_batch:
            sample_turns.append(test_batch.non_tensor_batch["__num_turns__"])

        data_source_lst.append(test_batch.non_tensor_batch.get("data_source", ["unknown"] * reward_tensor.shape[0]))

    self._maybe_log_val_generations(inputs=sample_inputs, outputs=sample_outputs, scores=sample_scores)

    # dump generations
    val_data_dir = self.config.trainer.get("validation_data_dir", None)
    if val_data_dir:
        self._dump_generations(
            inputs=sample_inputs,
            outputs=sample_outputs,
            gts=sample_gts,
            scores=sample_scores,
            reward_extra_infos_dict=reward_extra_infos_dict,
            dump_path=val_data_dir,
        )

    for key_info, lst in reward_extra_infos_dict.items():
        assert len(lst) == 0 or len(lst) == len(sample_scores), f"{key_info}: {len(lst)=}, {len(sample_scores)=}"

    if merged:
        print("_merge_validation_results validate result will be merged")
        return {
            "data_sources": data_source_lst,
            "sample_uids": sample_uids,
            "sample_turns": sample_turns,
            "reward_extra_infos_dict": reward_extra_infos_dict,
        }
    data_sources = np.concatenate(data_source_lst, axis=0)
    return self._val_metrics_update(data_sources, sample_uids, reward_extra_infos_dict, sample_turns)

FullyAsyncTaskRunner

提示

代码

• fully_async_main.py

初始化加载组件(trainer/rollouter)

FullyAsyncTaskRunner 初始化

• components：tokenizer， processor，
• 初始化 trainer
• 初始化 rollouter
• 加载 checkpoint

@ray.remote(num_cpus=1)
class FullyAsyncTaskRunner:
    """
    Ray remote class for executing distributed PPO training tasks.
    """

    def __init__(self):
        self.running = False
        self.components = {}
        self.shutdown_event = threading.Event()

    def run(self, config):
        print("[ASYNC MAIN] Starting fully async PPO training...")
        self._initialize_components(config)  
        self._run_training_loop()  

    def _initialize_components(self, config) -> None:
        print(f"[ASYNC MAIN] TaskRunner hostname: {socket.gethostname()}, PID: {os.getpid()}")
        pprint(OmegaConf.to_container(config, resolve=True))
        OmegaConf.resolve(config)

        print("[ASYNC MAIN] Initializing model and tokenizer...")
        local_path = copy_to_local(
            config.actor_rollout_ref.model.path, use_shm=config.actor_rollout_ref.model.get("use_shm", False)
        )
        from verl.utils import hf_processor, hf_tokenizer

        trust_remote_code = config.data.get("trust_remote_code", False)
        tokenizer = hf_tokenizer(local_path, trust_remote_code=trust_remote_code)  

        # Used for multimodal LLM, could be None
        processor = hf_processor(local_path, trust_remote_code=trust_remote_code, use_fast=True)

        self.components["tokenizer"] = tokenizer  
        self.components["processor"] = processor
        self.components["config"] = config  

        print("[ASYNC MAIN] Creating worker mapping and resource pools...")
        role_worker_mapping, ray_worker_group_cls = create_role_worker_mapping(config)
        self.components["role_worker_mapping"] = role_worker_mapping
        self.components["ray_worker_group_cls"] = ray_worker_group_cls

        print("[ASYNC MAIN] Creating FullyAsyncTrainer first (needed for hybrid worker group injection)...")
        self._create_trainer(config)  

        print("[ASYNC MAIN] Injecting trainer's worker group into rollouter for hybrid replicas...")
        self._setup_hybrid_worker_group(config)

        print("[ASYNC MAIN] Creating FullyAsyncRollouter...")
        self._create_rollouter(config)  

        print("[ASYNC MAIN] Setting up rollouter reference on trainer")
        ray.get(self.components["trainer"].set_rollouter.remote(self.components["rollouter"]))

        # sync total_train_steps between rollouter and trainer
        total_train_steps = ray.get(self.components["rollouter"].get_total_train_steps.remote()) 
        print(f"total_train_steps {total_train_steps}")
        ray.get(self.components["trainer"].set_total_train_steps.remote(total_train_steps))

        # max_queue_size
        max_queue_size = ray.get(self.components["rollouter"].get_max_queue_size.remote()) 
        print(f"[ASYNC MAIN] Creating MessageQueue... max_queue_size {max_queue_size}")
        message_queue = MessageQueue.remote(config, max_queue_size)
        message_queue_client = MessageQueueClient(message_queue)
        self.components["message_queue"] = message_queue
        self.components["message_queue_client"] = message_queue_client

        ray.get(self.components["rollouter"].set_message_queue_client.remote(self.components["message_queue_client"]))
        ray.get(self.components["trainer"].set_message_queue_client.remote(self.components["message_queue_client"]))

        # 加载参数 param_version resume from ckpt or default 0
        ray.get(self.components["trainer"].load_checkpoint.remote())  
        ray.get(self.components["rollouter"].load_checkpoint.remote())  

        print("[ASYNC MAIN] Param sync before fit..")
        ray.get(self.components["trainer"]._fit_update_weights.remote())  

        if config.trainer.get("val_before_train", True):
            ray.get(self.components["trainer"]._fit_validate.remote(True))  

        print("[ASYNC MAIN] All components initialized successfully")

    def _create_rollouter(self, config) -> None:
        print("[ASYNC MAIN] Starting create rollouter...")
        rollouter = FullyAsyncRollouter.remote( 
            config=config, 
            tokenizer=self.components["tokenizer"],
            processor=self.components["processor"],
            device_name=config.trainer.device,
        )

        # set_hybrid_worker_group must be called BEFORE init_workers() so that
        # _init_async_rollout_manager can pass the hybrid WG to ALM.create().
        if "hybrid_worker_group" in self.components:
            ray.get(rollouter.set_hybrid_worker_group.remote(self.components["hybrid_worker_group"]))
            print("[ASYNC MAIN] Hybrid worker group injected into rollouter")

        ray.get(rollouter.init_workers.remote())
        ray.get(rollouter.set_max_required_samples.remote())

        self.components["rollouter"] = rollouter
        print("[ASYNC MAIN] Rollouter created and initialized successfully")

    def _create_trainer(self, config) -> None:
        print("[ASYNC MAIN] Starting create trainer...")
        trainer_role_mapping = {
            role: worker_cls
            for role, worker_cls in self.components["role_worker_mapping"].items()
            if role != Role.Rollout
        }

        trainer = FullyAsyncTrainer.remote( 
            config=config, 
            tokenizer=self.components["tokenizer"], 
            role_worker_mapping=trainer_role_mapping,
            resource_pool_manager=create_resource_pool_manager(config, roles=list(trainer_role_mapping.keys())),
            ray_worker_group_cls=self.components["ray_worker_group_cls"],
            device_name=config.trainer.device,
        )

        ray.get(trainer.init_workers.remote())
        self.components["trainer"] = trainer
        print("[ASYNC MAIN] FullyAsyncTrainer created and initialized successfully")

    def _setup_hybrid_worker_group(self, config) -> None:
        """Extract the trainer's actor_rollout_wg for hybrid rollout replicas.

        Hybrid replicas are needed when *either*:
        - ``use_trainer_do_validate`` is enabled (trainer-side validation), or
        - ``use_dynamic_resource_scheduling`` is enabled (dynamic resource scheduling).

        The extracted worker group is stored in ``self.components["hybrid_worker_group"]``
        and injected into the rollouter before ``init_workers()`` is called.
        """
        trainer = self.components["trainer"]
        needs_hybrid = config.async_training.use_trainer_do_validate or config.async_training.get(
            "use_dynamic_resource_scheduling", False
        )
        if needs_hybrid:
            trainer_wg = ray.get(trainer.get_actor_wg.remote())
            self.components["hybrid_worker_group"] = trainer_wg
            print(
                f"[ASYNC MAIN] Hybrid worker group extracted from trainer "
                f"(world_size={getattr(trainer_wg, 'world_size', '?')})"
            )
        else:
            print(
                "[ASYNC MAIN] Neither use_trainer_do_validate nor use_dynamic_resource_scheduling enabled, "
                "skipping hybrid worker group setup"
            )
    # 省略run agent loop 入口

run_agent_loop

@ray.remote(num_cpus=1)
class FullyAsyncTaskRunner:
  """
  Ray remote class for executing distributed PPO training tasks.
  """
  """ 省略上述代码
  """

  def _run_training_loop(self):
      self.running = True

      print("[ASYNC MAIN] Starting Rollouter and Trainer...")
      rollouter_future = self.components["rollouter"].fit.remote() 
      trainer_future = self.components["trainer"].fit.remote() 

      futures = [rollouter_future, trainer_future]

      try:
          while futures:
              # Use ray.wait to monitor all futures and return when any one is completed.
              done_futures, remaining_futures = ray.wait(futures, num_returns=1, timeout=None)

              for future in done_futures:
                  try:
                      ray.get(future)
                      print("[ASYNC MAIN] One component completed successfully")
                  except Exception as e:
                      print(f"[ASYNC MAIN] Component failed with error: {e}")
                      for remaining_future in remaining_futures:
                          ray.cancel(remaining_future)
                      raise e

              futures = remaining_futures

      except Exception as e:
          print(f"[ASYNC MAIN] Training failed: {e}")
          for future in futures:
              ray.cancel(future)
          raise
      finally:
          asyncio.run(self.components["message_queue_client"].clear_queue())
          print("[ASYNC MAIN] Training completed or interrupted")

FullyAsyncTrainer (继承SeparateRayPPOTrainer)

提示

代码

• fully_async_policy/fully_async_trainer.py

初始化

@ray.remote(num_cpus=10)
class FullyAsyncTrainer(SeparateRayPPOTrainer):
  """
  A fully asynchronous PPO trainer that obtains samples from a MessageQueue for training.
  Based on an improved implementation of OneStepOffRayTrainer
  """

  def __init__(
      self,
      config,
      tokenizer,
      role_worker_mapping: dict[Role, WorkerType],
      resource_pool_manager: ResourcePoolManager,
      ray_worker_group_cls: RayWorkerGroup = RayWorkerGroup,
      device_name=None,
  ):
      # ==================== RayPPOTrainer config ====================

      # Store the tokenizer for text processing
      self.tokenizer = tokenizer
      self.config = config

      self.hybrid_engine = config.actor_rollout_ref.hybrid_engine
      assert not self.hybrid_engine

      self.role_worker_mapping = role_worker_mapping
      self.resource_pool_manager = resource_pool_manager
      self.use_reference_policy = need_reference_policy(self.config)

      self.use_rm = need_reward_model(self.config)

      # distillation config needed by _update_actor in ray_trainer.py
      from verl.trainer.distillation.losses import is_distillation_enabled

      if is_distillation_enabled(self.config.get("distillation")):
          self.distillation_config = omega_conf_to_dataclass(self.config.distillation)
      else:
          self.distillation_config = None

      self.use_critic = need_critic(self.config)
      self.ray_worker_group_cls = ray_worker_group_cls
      self.device_name = device_name if device_name else self.config.trainer.device

      # if ref_in_actor is True, the reference policy will be actor without lora applied
      lora_rank = config.actor_rollout_ref.model.get("lora", {}).get("rank", 0)
      if lora_rank <= 0:
          lora_rank = config.actor_rollout_ref.model.get("lora_rank", 0)
      self.ref_in_actor = lora_rank > 0 or config.actor_rollout_ref.model.get("lora_adapter_path") is not None

      # define in-reward KL control
      # kl loss control currently not suppoorted
      if self.config.algorithm.use_kl_in_reward:
          self.kl_ctrl_in_reward = core_algos.get_kl_controller(self.config.algorithm.kl_ctrl)

      self.use_prefix_grouper = self.config.actor_rollout_ref.actor.get("use_prefix_grouper", False)

      # ==================== SeparateRayPPOTrainer config ====================
      self.global_steps = 0
      self.epoch = 0
      self._init_dump_executor()
      self.validation_generations_logger = None
      self.max_steps_duration = 0
      self.progress_bar = None
      self.is_last_step = False
      self.prev_step_profile = False
      self.curr_step_profile = False
      self.next_step_profile = False
      self.last_val_metrics = {}
      self.metrics = {}
      self.timing_raw = {}
      # reward message
      self.future_reward = None
      self.reward_tensor = None
      self.reward_extra_infos_dict = {}

      self.logger = Tracking(
          project_name=self.config.trainer.project_name,
          experiment_name=self.config.trainer.experiment_name,
          default_backend=self.config.trainer.logger,
          config=OmegaConf.to_container(self.config, resolve=True),
      )

      # ==================== fully async config ====================

      self.message_queue_client = None

      # Statistics
      self.local_trigger_step = 1
      self.processed_samples = 0
      self.stale_trajectory_processed = 0
      self.current_param_version = 0
      self.total_train_steps = None
      self.progress_bar = None
      self.trigger_parameter_sync_step = config.async_training.trigger_parameter_sync_step 
      self.last_ckpt_version = 0
      # When use_trainer_do_validate OR use_dynamic_resource_scheduling is enabled, trainer
      # workers must carry a rollout engine (Role.ActorRollout) so that the master rank
      # can push weight updates directly to the colocated hybrid rollout instance via the
      # naive path (hybrid_checkpoint_manager).
      needs_hybrid_rollout = config.async_training.use_trainer_do_validate or config.async_training.get(
          "use_dynamic_resource_scheduling", False
      )
      self.train_role = Role.ActorRollout if needs_hybrid_rollout else Role.Actor

      # required_samples use ppo_mini_batch_size*require_batches as the minimum number of samples.
      self.require_batches = config.async_training.require_batches  
      self.required_samples = config.actor_rollout_ref.actor.ppo_mini_batch_size * self.require_batches  
      self._step_wait_times: list[float] = []  # per-collection wait times within the current step (seconds)
      # Per-collection count of samples that actually had to be waited on (not
      # already sitting in the queue at collection start). Parallel to
      # _step_wait_times; see _get_samples_from_queue().
      self._step_wait_samples: list[int] = []
      # Hybrid GPUs (trainer-node GPUs that switch between rollout/train under dynamic
      # resource scheduling) and standalone GPUs (dedicated rollout-node GPUs, always
      # doing rollout). Used both for the existing throughput metric and to combine
      # dynamic_resource/{train,rollout}_resource_utilization into a single
      # dynamic_resource/resource_utilization metric (see MetricsAggregator).
      hybrid_gpus = config.trainer.nnodes * config.trainer.n_gpus_per_node
      standalone_gpus = config.rollout.nnodes * config.rollout.n_gpus_per_node
      total_gpus = hybrid_gpus + standalone_gpus
      self.metrics_aggregator = MetricsAggregator(
          total_gpus=total_gpus, hybrid_gpus=hybrid_gpus, standalone_gpus=standalone_gpus
      )

      # Reference to rollouter for parameter synchronization
      self.rollouter = None
      self.checkpoint_manager = None

      # Hybrid checkpoint manager for trainer-side validation (use_trainer_do_validate)
      # and/or dynamic resource scheduling (use_dynamic_resource_scheduling).
      # Uses naive backend to sync weights from trainer to hybrid rollout replicas.
      # Initialized in _setup_hybrid_checkpoint_manager() via set_rollouter().
      self.hybrid_checkpoint_manager = None

      # Dynamic resource controller — activated when use_dynamic_resource_scheduling=True.
      self.dynamic_resource_controller = None
      self.dynamic_schedule_enabled: bool = config.async_training.get("use_dynamic_resource_scheduling", False)
      # Name of the scheduling policy (resolved in _setup_dynamic_resource_controller).
      self._dynamic_schedule_policy_name: str = config.async_training.get("dynamic_schedule_policy", "default")
      # Initial deactivate_ratio forwarded to the policy constructor.
      self._dynamic_schedule_deactivate_ratio_init: float = config.async_training.get(
          "dynamic_schedule_deactivate_ratio", 0.3
      )
      # Whether to enable request rebalancing (abort + clear sticky cache +
      # resume) after hybrid replica activation. Default False.
      self._dynamic_schedule_enable_rebalance: bool = config.async_training.get(
          "dynamic_schedule_enable_rebalance", True
      )
      self.staleness_threshold: float = config.async_training.get("staleness_threshold", 1)

      # When standalone rollout resources are 0 (rollout.nnodes == 0), there are no
      # standalone replicas: all rollout happens on hybrid (trainer-side) GPUs.
      self.only_hybrid: bool = self.dynamic_schedule_enabled and config.rollout.nnodes == 0

      # Per-step dynamic scheduling context — built once at init, mutable fields updated each step.
      self.dynamic_schedule_ctx = DynamicScheduleContext(
          required_samples=self.required_samples,
          trigger_parameter_sync_step=self.trigger_parameter_sync_step,
          total_generated_samples=0,
          expected_samples=0,
          buffer_samples=0,
          only_hybrid=self.only_hybrid,
      )

Set rollouter

async def set_rollouter(self, rollouter):
  """Set rollouter reference and initialize all checkpoint managers."""
  self.rollouter = rollouter
  # Setup checkpoint manager after rollouter is set
  await self._setup_checkpoint_manager()
  await self._setup_hybrid_checkpoint_manager()
  # Setup dynamic resource controller if enabled
  if self.dynamic_schedule_enabled:
      await self._setup_dynamic_resource_controller()

fit 主流程

async def fit(self):
    """
    The training loop of PPO.
    The driver process only need to call the compute functions of the worker group through RPC
    to construct the PPO dataflow.
    The light-weight advantage computation is done on the driver process.
    """
    print("[FullyAsyncTrainer] Starting FullyAsyncTrainer...")
    if self.message_queue_client is None:
        raise ValueError("MessageQueue client not set. Call set_message_queue_client() first.")
    if self.rollouter is None:
        raise ValueError("rollouter not set. Call set_rollouter() first.")

    self.max_steps_duration = 0

    self.global_steps += 1

    self.prev_step_profile = False
    self.curr_step_profile = False
    self.next_step_profile = False

    # Use queue mode, no need for traditional dataloader iterator
    # Initialize to get the first batch of data
    while True:
        try:
            await self.fit_step()  
        except TrainingStopException:
            print("[FullyAsyncTrainer] Training stopped by queue termination signal")
            break

    self.progress_bar.close()
    if self.current_param_version % self.config.trainer.test_freq != 0 or self.local_trigger_step > 1:
        rollout_reset_timing_raw = await self._fit_update_weights()
        if rollout_reset_timing_raw is not None:
            self._fit_log_aggregated_training_metrics(rollout_reset_timing_raw)
        await self._fit_validate()
    self._fit_save_checkpoint(force=True)

fit_step

async def fit_step(self, batch_dict: dict = None):
    """
    Single-step training template method. Handles all logic for one training step.

    Flow:
    1. Pre-step processing -> 2. Get batch -> 3. Generate sequences ->
    4. Compute reward -> 5. Compute log_prob -> 6. Compute reward ->
    7. Compute advantage -> 8. Update critic -> 9. Update actor -> 10. Post-step processing

    Args:
        batch_dict: Raw data dictionary
    """
    self.metrics = {"training/global_step": self.global_steps, "training/epoch": self.epoch}
    self.timing_raw = {}
    # reward message
    self.future_reward = None
    self.reward_tensor = None
    self.reward_extra_infos_dict = {}

    steps = self.config.global_profiler.steps
    should_profile = steps is not None and (self.current_param_version + 1) in steps
    self._fit_start_profile(should_profiler=should_profile)

    with marked_timer("step", self.timing_raw):
        ctrl = self.dynamic_resource_controller
        if self.dynamic_schedule_enabled and ctrl.policy.should_deactivate(
            global_steps=self.current_param_version,
            is_hybrid_active=ctrl.is_hybrid_active,
            ctx=self.dynamic_schedule_ctx,
        ):
            threshold_samples = ctrl.policy.deactivate_wait_samples(self.dynamic_schedule_ctx)
            with marked_timer("wait_for_enough_samples", self.timing_raw):
                _ = ray.get(self.rollouter.wait_for_enough_samples.remote(threshold_samples))
            _deact_start = time.time()
            await ctrl.deactivate_hybrid_replicas(self.current_param_version)
            deactivate_duration = time.time() - _deact_start
            self.dynamic_schedule_ctx.last_deactivate_duration_s += deactivate_duration
            print(
                f"[FullyAsyncTrainer] step={self.current_param_version} "
                f"deactivation took {deactivate_duration:.2f}s "
                f"(accumulated this cycle: {self.dynamic_schedule_ctx.last_deactivate_duration_s:.2f}s)",
                flush=True,
            )
        elif self.dynamic_schedule_enabled:
            # Not deactivating this step: still emit the metric as 0 so the
            # timing_s/wait_for_enough_samples curve stays continuous.
            self.timing_raw["wait_for_enough_samples"] = 0.0

        _allocated_start = time.time()
        batch = await self._fit_generate(None)
        batch = self._fit_compute_reward(batch)
        batch = self._fit_compute_log_prob(batch)
        batch = self._fit_compute_ref_log_prob(batch)
        batch = self._fit_compute_critic(batch)
        batch = self._fit_compute_advantage(batch)
        batch = self._fit_update_critic(batch)
        batch = self._fit_update_actor(batch)
        self._fit_update_local_step()
        rollout_reset_timing_raw = await self._fit_update_weights()
        self._fit_dump_data(batch)
        self._record_train_resource_utilization(allocated_time=time.time() - _allocated_start)

    await self._fit_validate()
    self._fit_save_checkpoint()
    self._fit_stop_profile(should_profiler=should_profile)
    self._fit_collect_metrics(batch)
    if rollout_reset_timing_raw is not None:
        self._fit_log_aggregated_training_metrics(rollout_reset_timing_raw)
    self._fit_postprocess_step()

_fit_generate

提示

• 从队列里取batch

async def _fit_generate(self, batch: DataProto = None) -> DataProto | None:
    metrics = self.metrics
    timing_raw = self.timing_raw
    with marked_timer("gen", timing_raw, color="red"):
        epoch, batch = await self._get_samples_from_queue()  
        if batch is None:
            raise TrainingStopException("Training terminated: queue returned None")
        self._collect_metrics_from_samples(batch, metrics)
    batch.meta_info["temperature"] = self.config.actor_rollout_ref.rollout.temperature
    return batch

async def _get_samples_from_queue(self) -> tuple[None, None] | tuple[int, Any]:
    """
    Get samples from message queue and compose gen_batch_output
    Uses a loop to continuously collect samples until enough are gathered

    Returns:
        tuple: (epoch, batch_dict, gen_batch_output)
    """
    print(
        f"[FullyAsyncTrainer] Requesting {self.required_samples} samples from queue",
        flush=True,
    )

    # Collect samples using a simple loop calling get_sample
    consumer_start = time.time()
    # Snapshot the queue backlog at collection start: samples already sitting
    # in the queue are served instantly and don't reflect actual generation
    # rate, so they must be excluded from the wait-time-per-sample estimate.
    # Only queried when dynamic scheduling is enabled, since it's the sole consumer
    # of this signal and the extra RPC would otherwise be pure overhead.
    if self.dynamic_schedule_enabled:
        queue_size_at_start = await self.message_queue_client.get_queue_size()
        pending_wait_samples = max(0, self.required_samples - queue_size_at_start)
    else:
        pending_wait_samples = 0
    queue_samples = []
    queue_len = 0
    while len(queue_samples) < self.required_samples:
        # Get a single sample and wait until there is a sample or None is received
        sample, queue_len = await self.message_queue_client.get_sample()  

        if sample is None:
            print(
                f"[FullyAsyncTrainer] Detected termination signal (None), stopping sample collection. "
                f"Collected {len(queue_samples)}/{self.required_samples} samples"
            )
            break

        queue_samples.append(sample)
        print(
            f"[FullyAsyncTrainer] sample collected {len(queue_samples)}/{self.required_samples}. "
            f"mq_len: {queue_len}",
            flush=True,
        )

        if len(queue_samples) % 64 == 0:
            print(
                f"[FullyAsyncTrainer] Collected {len(queue_samples)}/{self.required_samples} samples. "
                f"mq_len: {queue_len}"
            )

    consumer_end = time.time()

    if not queue_samples or len(queue_samples) < self.required_samples:
        print("[FullyAsyncTrainer] not enough samples collected after loop")
        return None, None
    total_wait_time = consumer_end - consumer_start

    print(
        f"[FullyAsyncTrainer] Loop collection completed: {len(queue_samples)}/{self.required_samples} samples, "
        f"total wait time: {total_wait_time:.2f} seconds. "
        f"mq_len: {queue_len}"
    )

    queue_samples = [ray.cloudpickle.loads(x) for x in queue_samples]
    # Assemble batch - now working directly with RolloutSample objects
    if self.config.trainer.balance_batch:
        batch = assemble_batch_from_rollout_samples(queue_samples, self.tokenizer, self.config, self._balance_batch)
    else:
        batch = assemble_batch_from_rollout_samples(queue_samples, self.tokenizer, self.config, None)

    batch.meta_info["fully_async/total_wait_time"] = total_wait_time
    self._step_wait_times.append(total_wait_time)
    # pending_wait_samples may be 0 when this collection was served entirely
    # from queue backlog (no real waiting for generation happened); the policy
    # layer special-cases that when estimating the generation rate.
    self._step_wait_samples.append(pending_wait_samples)
    return 0, batch

_fit_compute_reward

def _fit_compute_reward(self, batch: DataProto) -> DataProto:
    timing_raw = self.timing_raw
    with marked_timer("reward", timing_raw, color="yellow"):
        # compute reward model score
        if self.use_rm and "rm_scores" not in batch.batch.keys():
            batch_reward = self._compute_reward_colocate(batch)  
            batch = batch.union(batch_reward)

        # Compute or extract reward_tensor and reward_extra_infos_dict for training
        reward_tensor, reward_extra_infos_dict = extract_reward(batch)
        self.reward_tensor = reward_tensor
        self.reward_extra_infos_dict = reward_extra_infos_dict
    return batch

def _compute_reward_colocate(self, batch: DataProto) -> tuple[torch.Tensor, dict[str, Any]] | torch.Tensor:
    """
    compute reward use colocate reward model
    """
    assert self.reward_loop_manager is not None, "RewardLoopManager is None"
    batch_reward = self.reward_loop_manager.compute_rm_score(batch) 
    return batch_reward

FullyAsyncRollouter

异步参数相关笔记

链接

• fully_async 异步参数笔记
• fully_async 参数源代码笔记
• fully_async_policy/fully_async_rollouter.py

初始化

初始化参数训练数据集

@ray.remote(num_cpus=10, max_concurrency=100)
class FullyAsyncRollouter(SeparateRayPPOTrainer):
    """
    Asynchronous sample generator, responsible for continuously generating training samples
    and putting them into MessageQueue
    Based on the mature implementation improvements of OneStepOffRayTrainer
    """

    def __init__(
        self,
        config,
        tokenizer,
        processor=None,
        device_name=None,
    ):
        # Store the tokenizer for text processing
        self.tokenizer = tokenizer
        self.processor = processor
        self.config = config
        self.hybrid_engine = config.actor_rollout_ref.hybrid_engine

        assert not self.hybrid_engine
        assert self.config.data.train_batch_size == 0, "train_batch_size must be zero"
        assert self.config.data.gen_batch_size == 1, "gen_batch_size must be one"
        assert self.config.async_training.staleness_threshold >= 0, "staleness_threshold must larger than 0"
        assert self.config.async_training.trigger_parameter_sync_step >= 1, (
            "trigger_parameter_sync_step must larger or equal than 1"
        )

        self.use_reference_policy = False

        self.use_rm = need_reward_model(self.config)
        if self.use_rm:
            assert self.config.reward.reward_model.enable_resource_pool, (
                "GenRM/DisRM in fully async mode requires standalone mode (enable_resource_pool=True). "
                "Colocate mode is not supported because async rollout never pauses."
            )

        self.use_critic = False
        self.device_name = device_name if device_name else self.config.trainer.device
        self.validation_generations_logger = ValidationGenerationsLogger(
            project_name=self.config.trainer.project_name,
            experiment_name=self.config.trainer.experiment_name,
        )

        self.ref_in_actor = False
        self.kl_ctrl_in_reward = False

        self.use_prefix_grouper = self.config.actor_rollout_ref.actor.get("use_prefix_grouper", False)

        self._init_dump_executor()

        # ==================== fully async config ====================

        print("[FullyAsyncRollouter] Creating datasets...")
        from verl.utils.dataset.rl_dataset import collate_fn

        train_dataset = create_rl_dataset( 
            config.data.train_files,
            config.data,
            tokenizer,
            processor,
            max_samples=config.data.get("train_max_samples", -1),
        )
        val_dataset = create_rl_dataset( 
            config.data.val_files,
            config.data,
            tokenizer,
            processor,
            max_samples=config.data.get("val_max_samples", -1),
        )
        train_sampler = create_rl_sampler(config.data, train_dataset)

        self._validate_config()
        self._create_dataloader(train_dataset, val_dataset, collate_fn, train_sampler)

        self.total_rollout_steps = len(self.train_dataloader) * self.config.trainer.total_epochs  
        if self.config.rollout.total_rollout_steps is not None:
            self.total_rollout_steps = min(self.config.rollout.total_rollout_steps, self.total_rollout_steps)
        print(f"[FullyAsyncRollouter] Total rollout steps: {self.total_rollout_steps}")
        self.total_train_steps = None

        # Rollouter parameter configuration
        self.message_queue_client = None

        self.async_rollout_manager = None

        # Elastic worker group (injected via set_hybrid_worker_group before init_workers)
        # When set, its GPUs back hybrid replicas for trainer-side validation.
        self._hybrid_worker_group = None

        # Config
        self.staleness_threshold: float = config.async_training.get("staleness_threshold", 1)
        # required_samples use ppo_mini_batch_size*require_batches as the minimum number of samples.
        self.require_batches = config.async_training.require_batches  
        self.required_samples = config.actor_rollout_ref.actor.ppo_mini_batch_size * self.require_batches  
        self.concurrent_samples_per_replica: int = config.async_training.get("concurrent_samples_per_replica", 16)  
        self.max_required_samples = None
        self.max_concurrent_samples = None
        # queue size
        self.max_queue_size = None

        # Statistics
        self.total_generated_samples = 0
        self.staleness_samples = 0
        self.dropped_stale_samples = 0
        self.processed_sample_count = 0
        # Per-step sample counter: counts fully-generated samples in the current param version.
        # Reset to 0 at each reset_staleness() call.
        self._step_generated_samples: int = 0
        # Rolling history of per-step sample counts: (param_version, sample_count).
        # Keeps the most recent _STEP_HISTORY_SIZE entries for throughput diff analysis.
        self._STEP_HISTORY_SIZE: int = 10
        self._step_samples_history: collections.deque = collections.deque(maxlen=self._STEP_HISTORY_SIZE)
        # Monotonically increasing counter of completed param-sync steps.
        # Must NOT be capped by _STEP_HISTORY_SIZE; used by get_completed_steps() to
        # compute expected_samples in the trainer, which grows without bound.
        self._completed_steps: int = 1
        # we start from step 1
        self.global_steps = 1
        self.idle_start_time = time.time()
        self.step_start_time = time.time()

        # Concurrency control
        # Modified by self.pause() or self._should_pause_generation()
        self.paused = False
        self.running = True

        # Add dataloader lock
        self.dataloader_lock = asyncio.Lock()

        # Initialize async queues
        self.pending_queue = asyncio.Queue(maxsize=128)
        self.active_tasks = set()

        # Event-driven history of (len(self.active_tasks), max_concurrent_samples) over time,
        # used to compute dynamic_resource/rollout_resource_utilization in reset_staleness().
        # Each entry is (timestamp, active_count_after_change, max_concurrent_at_that_time),
        # appended whenever active_tasks' size changes (a sample is submitted, completes, or is
        # drained during a pause) OR when max_concurrent_samples itself changes (replicas
        # added/removed under dynamic resource scheduling). Recording the capacity alongside the
        # active count is required because max_concurrent_samples is NOT constant over a step
        # window when rollout resources are dynamically (de)activated — using a single
        # end-of-window capacity for the whole window would misattribute utilization for
        # intervals that had a different capacity. Seeded with a single (t0, 0, 0) point so the
        # very first interval has a well-defined start; the capacity will be corrected to the
        # real value as soon as max_concurrent_samples is known (see set_max_required_samples).
        self._active_count_history: list[tuple[float, int, int]] = [(time.time(), 0, 0)]

    def _init_async_objects(self):
        # Initialize asyncio synchronization primitives.
        # `lock` protects shared state: paused / active_tasks / staleness_samples / timing fields.
        self.lock = asyncio.Lock()
        # `_resume_event` signals that the rollouter is currently running (paused == False).
        self._resume_event = asyncio.Event()
        self._resume_event.set()

    async def set_message_queue_client(self, message_queue_client: MessageQueueClient):
        """Set message queue client"""
        async with self.lock:
            self.message_queue_client = message_queue_client

设置rollouter参数(最大样本数/最大并发样本数/训练步数/队列大小)

异步参数相关笔记

链接

• fully_async 异步参数笔记
• fully_async 参数源代码笔记
• fully_async_policy/fully_async_rollouter.py

核心

• 这里samples，都是组的概念
• 在后续 prepare_single_generation_data 里，才会对 batch 扩充 rollout_n 倍。

async def set_max_required_samples(self):
  async with self.lock:
      self.max_required_samples = int(  
          self.required_samples
          * (self.staleness_threshold + 1)
          * self.config.async_training.trigger_parameter_sync_step
      )
      self.total_train_steps = int(  
          self.total_rollout_steps
          / (self.required_samples * self.config.async_training.trigger_parameter_sync_step)
      )

      self.max_concurrent_samples = (  
          self.llm_server_manager.get_active_server_count() * self.concurrent_samples_per_replica
      )
      self.max_concurrent_samples = min(self.max_concurrent_samples, self.max_required_samples)  
      self.max_queue_size = self.max_required_samples

      print(
          f"[FullyAsyncRollouter] required_samples : {self.required_samples} "
          f"max_required_samples: {self.max_required_samples} "
          f"max_queue_size: {self.max_queue_size} "
          f"total_train_steps: {self.total_train_steps} "
          f"total_rollout_steps: {self.total_rollout_steps} "
          f"max_concurrent_samples: {self.max_concurrent_samples} "
      )

      # The initial seed point in _active_count_history was recorded in __init__ before
      # max_concurrent_samples was known (placeholder capacity 0). Now that the real
      # capacity is available, correct that seed point so the very first interval isn't
      # spuriously treated as having zero capacity.
      if len(self._active_count_history) == 1:
          t0, active0, _capacity0 = self._active_count_history[0]
          self._active_count_history[0] = (t0, active0, self.max_concurrent_samples)

detach_utils.py 里，再对 batch 扩充 rollout_n倍

def prepare_single_generation_data(batch_dict, config) -> DataProto:
    """
    Similar to the logic of ray_trainer._prepare_generate_batch, but for a single sample.
    Separate the data used for generation from the original data.

    Returns:
        tuple: (original_batch_dict, gen_data_for_single_sample)
    """

    full_batch = DataProto.from_single_dict(batch_dict)

    batch_keys_to_pop = []
    non_tensor_batch_keys_to_pop = []

    existing_batch_keys = [k for k in batch_keys_to_pop if k in full_batch.batch.keys()]
    existing_non_tensor_keys = [k for k in non_tensor_batch_keys_to_pop if k in full_batch.non_tensor_batch.keys()]

    if existing_batch_keys or existing_non_tensor_keys:
        full_batch.pop(
            batch_keys=existing_batch_keys,
            non_tensor_batch_keys=existing_non_tensor_keys,
        )

    # Setting selected agent, that supports partial
    if not config.actor_rollout_ref.rollout.multi_turn.enable:
        full_batch.non_tensor_batch["agent_name"] = np.array(["single_turn_agent"] * len(full_batch), dtype=object)

    # Add global step count to generated data
    full_batch = full_batch.repeat(repeat_times=config.actor_rollout_ref.rollout.n, interleave=True)  
    return full_batch

初始化workers

async def init_workers(self):
  """Initialize distributed training workers using Ray backend.

  Creates:
  1. Ray resource pools from configuration
  2. Worker groups for each role (actor, critic, etc.)
  """
  self._init_async_objects()
  self._create_worker_classes()
  await self._create_reward_loop_manager()
  await self._create_teacher_model_manager()
  await self._init_async_rollout_manager()  
  SkipManager.init(self.config)

初始化aysnc_rollout_manger

async def _init_async_rollout_manager(self):
    """
    Create the server manager and agent loop manager for fully async training.

    Uses :class:`FullyAsyncLLMServerManager` which supports two-phase init:
    - Phase 1: hybrid replicas on trainer GPUs (sleeping)
    - Phase 2: standalone replicas on rollout GPUs

    The ``GlobalRequestLoadBalancer`` (which also holds the server-handle
    registry) serves as the single source of truth for handle mapping and
    routing.  Clients look up handles atomically — no per-worker notification
    needed on hybrid add/remove.
    """
    # infrastructure overview: https://verl.readthedocs.io/en/latest/advance/reward_loop.html#architecture-design
    # agent_reward_loop: streaming reward computation with actor rollout
    # two conditions satisfied: (1) no reward model, or (2) reward model with extra resource pool
    enable_agent_reward_loop = not self.use_rm or self.config.reward.reward_model.enable_resource_pool

    # if enable_agent_reward_loop, we directly pass reward_loop_workers to agent loop manager
    # to stream reward computation with actor rollout
    reward_loop_worker_handles = self.reward_loop_manager.reward_loop_workers if enable_agent_reward_loop else None

    # create async rollout manager and request scheduler
    assert self.config.actor_rollout_ref.rollout.mode == "async"

    self.async_rollout_mode = True
    # Use FullyAsyncLLMServerManager for two-phase (hybrid + standalone) init.
    # It creates GlobalRequestLoadBalancer (with merged handle registry) internally.
    self.llm_server_manager = await FullyAsyncLLMServerManager.create(  
        config=self.config,
        worker_group=self.get_hybrid_worker_group(),
    )
    self.async_rollout_manager = await FullyAsyncAgentLoopManager.create(  
        config=self.config,
        llm_client=self.llm_server_manager.get_client(client_cls=FullyAsyncLLMServerClient),
        reward_loop_worker_handles=reward_loop_worker_handles,
        teacher_client=self.teacher_model_manager.get_client() if self.teacher_model_manager else None,
    )

FullyAsyncLLMServerManager，llm server 初始化 - SGLangHttpServer (SGLangReplica) 笔记

class FullyAsyncLLMServerManager(LLMServerManager):
    """Extension of :class:`LLMServerManager` for fully async training with hybrid scheduling."""

    def __init__(
        self,
        config: DictConfig,
        worker_group: RayWorkerGroup = None,
        rollout_resource_pool: RayResourcePool = None,
    ):
        super().__init__(config, worker_group, rollout_resource_pool)
        # Pre-registered hybrid replicas: bound at init time but still sleeping.
        # Keyed by resource_id; populated during _initialize_llm_servers().
        self.hybrid_replicas: dict[str, RolloutReplica] = {}
        # Currently active (awake + in LB) subset of hybrid replicas.
        self.alive_replicas: dict[str, RolloutReplica] = {}
        # resource_id → server_address for alive hybrid replicas.
        self.alive_addresses: dict[str, str] = {}
        # Prometheus server addresses
        self.prometheus_server_addresses = []

        # Timing / counters
        self.last_hybrid_add_time: float = 0.0
        self.last_hybrid_remove_time: float = 0.0

    async def _initialize_llm_servers(self, start_rank: int = 0):
        # ── Step 1: hybrid replicas first (replica_rank 0 … N_e-1) ──────────
        # Use parent class to create + init_hybrid all hybrid replicas, then
        # migrate them from rollup_replicas → hybrid_replicas (sleeping, not
        # yet in the load balancer).  Starting from rank 0 gives hybrid actors
        # the lowest-numbered placement-group bundles which are co-located with
        # the training engine, maximising GPU affinity on multi-node deployments.
        num_hybrid = 0
        if self.worker_group is not None:
            await super()._initialize_llm_servers(start_rank=0)  
            num_hybrid = len(self.rollout_replicas)
            # Migrate hybrid replicas out of the parent's tracking lists.
            for i, replica in enumerate(self.rollout_replicas):
                resource_id = f"hybrid_{i}"
                self.hybrid_replicas[resource_id] = replica
                print(
                    f"[FullyAsyncAgentLoopManager] Hybrid replica '{resource_id}' "
                    f"(rank={i}) initialised at {replica._server_address} "
                )
            self.prometheus_server_addresses.extend(self.server_addresses)
            print(f"AgentLoopManager Hybrid: {self.server_addresses}")
            # Clear parent state so Step 2 starts clean.
            self.rollout_replicas = []
            self.server_handles = []
            self.server_addresses = []

        # ── Step 2: standalone replicas via parent class ─────────────────────
        # Temporarily clear worker_group so that super()._initialize_llm_servers()
        # takes the standalone branch (init_standalone).  Pass start_rank=num_hybrid
        # so that Ray actor names remain globally unique and never collide with the
        # hybrid actors created above.
        #
        # If standalone_gpu_memory_utilization is configured, temporarily override
        # gpu_memory_utilization so that standalone replicas can use a higher value
        # (e.g. 0.85) than hybrid replicas (which must share GPU with the training
        # engine).  The original value is restored in the finally block regardless
        # of whether the initialisation succeeds or fails.
        standalone_gmu = self.rollout_config.get("standalone_gpu_memory_utilization", None)
        original_gmu = self.rollout_config.get("gpu_memory_utilization", None)
        saved_worker_group = self.worker_group
        self.worker_group = None
        try:
            if standalone_gmu is not None and original_gmu is not None:
                with open_dict(self.rollout_config):
                    self.rollout_config.gpu_memory_utilization = standalone_gmu
            await super()._initialize_llm_servers(start_rank=num_hybrid)
        finally:
            # Always restore the original value to avoid affecting downstream logic.
            if standalone_gmu is not None and original_gmu is not None:
                with open_dict(self.rollout_config):
                    self.rollout_config.gpu_memory_utilization = original_gmu
            self.worker_group = saved_worker_group

        # Update Prometheus with the final (standalone) addresses.
        if self.rollout_config.prometheus.enable:
            if self.rollout_config.disable_log_stats:
                raise ValueError("PROMETHEUS needs disable_log_stats==False, but it is currently True.")
            all_addresses = self.prometheus_server_addresses + self.server_addresses
            update_prometheus_config(self.rollout_config.prometheus, all_addresses, self.rollout_config.name)

        print(
            f"[FullyAsyncLLMServerManager] Created: "
            f"{len(self.rollout_replicas)} standalone replicas (rank {num_hybrid}+), "
            f"{num_hybrid} hybrid replicas registered (sleeping, rank 0-{num_hybrid - 1})"
        )

FullyAsyncAgentLoopManager，继承自AgentLoopManager，笔记见：AgentLoopManager 初始化 (SGLangHttpServer,AgentLoopWorker)

class FullyAsyncAgentLoopManager(AgentLoopManager):
    @SkipManager.annotate(role="async_rollout")
    async def generate_sequences_single(self, prompts: DataProto) -> DataProto:
        """Split input batch and dispatch to agent loop workers.

        Args:
            prompts (DataProto): Input batch. Single sample data
        Returns:
            DataProto: Output batch.
        """
        worker = self._select_best_worker()
        output_future = worker.generate_sequences.remote(prompts)
        return await asyncio.wrap_future(output_future.future())

    def _select_best_worker(self):
        """Select the best worker, simple round-robin load balancing"""
        if not hasattr(self, "_worker_index"):
            self._worker_index = 0

        worker = self.agent_loop_workers[self._worker_index]
        self._worker_index = (self._worker_index + 1) % len(self.agent_loop_workers)
        return worker