从零开始学 verl 框架

"""
- trainer
  - config
    - ppo_trainer.yaml
  - main_ppo.py
"""
import hydra
from omegaconf import DictConfig

@hydra.main(config_path="config", config_name="ppo_trainer", version_base=None)
def main(config):
	# ...
    run_ppo(config)

@ray.remote
class Accumulator:
    def __init__(self):
        self.value = 0

    def add(self, x):
        self.value += x

    def get_value(self):
        return self.value

accumulator = Accumulator.remote()
value_ref = accumulator.get_value.remote()
# ObjectRef(16310a0f0a45af5c9e4edcae28cff4ef250feac40100000001000000)
value = ray.get(value_ref)  
# 0

@ray.remote
class GPUAccumulator(Worker):
    def __init__(self):
        super().__init__()
        self.value = torch.zeros(1, device="cuda") + self.rank

    @register(Dispatch.ONE_TO_ALL)
    def add(self, x):
        self.value += x
        return self.value.cpu()

resource_pool = RayResourcePool([1], use_gpu=True)
class_with_args = RayClassWithInitArgs(cls=GPUAccumulator)
worker_group = RayWorkerGroup(resource_pool, class_with_args)
print(worker_group.add(x=10))

class DataProto:
    batch: TensorDict          # tensor 类型的字段（input_ids, attention_mask 等）
    non_tensor_batch: dict     # 非 tensor 字段（data_source, ground_truth 等）
    meta_info: dict            # 元信息（seqlen 等）

def train(self) -> None:
    """
    标准 RLHF PPO pipeline
    1. 用 actor 根据 prompt 生成回答
    2. 用 actor/ref 计算每个 token 的 logprobs
    3. 用 reward model 给整条回复打分
    4. 用 critic 估计 value 并计算优势（GAE）
    5. 用 PPO 更新 actor 和 critic
    """
    for step, batch in enumerate(self.dataloader):
        # 生成样本
        samples = self.get_sample(batch)
        # 生成经验，包括 rollout 和 evaluation 阶段的 adv rewards 等等
        experience = self.get_experience(samples)
        torch.cuda.empty_cache()
        # 每个 prompt 都要进行多轮训练
        for epoch in range(self.config.ppo_epochs):
            self.train_step(experience)

定义: Actor rollout → ref/reward inference → critic inference → critic training → actor training
执行: verl 自动调度，分配 GPU，切换并行模式

def init_workers(self):
	all_wg = {}
	self.wg_dicts = []
	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)
	    spawn_wg = wg_dict.spawn(prefix_set=class_dict.keys())
	    all_wg.update(spawn_wg)
	    # keep the referece of WorkerDict to support ray >= 2.31. Ref: https://github.com/ray-project/ray/pull/45699
	    self.wg_dicts.append(wg_dict)

ray_trainer.py / fit()          ← 第4层：single-controller，只管调度逻辑
    ↓ 调用
RayWorkerGroup                  ← 第3层：multi-controller，类似 torchrun，把方法分布式广播给组内所有 Worker
    ↓ 调用
fsdp_workers.py                 ← 第2层：ActorRolloutRefWorker，调度层
  内部持有 ↓                          把 rollout、actor training、ref 三个角色合并管理
dp_actor.py                     ← 第1层：DataParallelPPOActor，执行层
                                      真正跑前向/反向传播的地方

fit()
├── 初始化（加载checkpoint、验证）
└── 训练主循环（epoch → batch）
    ├── 生成阶段（Rollout）
    ├── 奖励计算阶段
    ├── 对数概率计算阶段
    ├── 优势估计阶段（Advantage）
    ├── Critic/Actor更新阶段
    └── 日志/验证/保存

# 1. 准备数据，只取 prompt 部分去掉 response
batch = DataProto.from_single_dict(batch_dict)
gen_batch = self._get_gen_batch(batch)

# 2. rollout 生成回答
#    底层通过 Ray 进程间通信把 DataProto 传给 rollout worker，计算完再传回
gen_batch_output = self.async_rollout_manager.generate_sequences(gen_batch)

# 3. 把生成结果合并回 batch
batch = batch.union(gen_batch_output)

# 4. 奖励 → KL 惩罚 → 优势估计
reward_tensor = extract_reward(batch)
batch = apply_kl_penalty(batch)
batch = compute_advantage(batch)    # GAE / GRPO 等

# 5. 更新网络
critic_output = self._update_critic(batch)
actor_output  = self._update_actor(batch)
self.checkpoint_manager.update_weights(...)  # 同步最新权重到 rollout worker

ActorRolloutRefWorker
├── Actor    → 用 PPO 更新模型参数（训练）
├── Rollout  → 用当前模型生成回答（采样）
└── Ref      → 用原始模型算 KL 散度用的参考概率（不训练）

async def rollout_mode(self):
    # 把 FSDP 分片的参数重新聚合成完整参数，推送给推理引擎
    params = self.actor_module_fsdp.state_dict()
    await self.rollout.update_weights(params)
    await self.rollout.resume(tags=["kv_cache"])   # 恢复 KV cache，准备生成

def generate_sequences(self, prompts):
    loop.run_until_complete(self.rollout_mode())           # 切换到推理模式
    output = self.rollout.generate_sequences(prompts)      # vLLM/SGLang 生成
    loop.run_until_complete(self.trainer_mode())           # 切回训练模式
    return output

DataProto (全量 batch)
    ↓  split
mini_batch × ppo_epochs        ← PPO 多轮更新
    ↓  split  
micro_batch × gradient_accum   ← 梯度累积（省显存）
    ↓
_forward_micro_batch()          ← 实际 forward
    ↓
loss.backward()
    ↓
_optimizer_step()               ← clip grad + step

for epoch in ppo_epochs:
    for mini_batch in mini_batches:
        optimizer.zero_grad()
        for micro_batch in micro_batches:           # 梯度累积
            log_prob = forward(micro_batch)
            pg_loss = policy_loss_fn(old_log_prob, log_prob, advantages)
            loss = pg_loss - entropy * coeff + kl_loss * coeff
            (loss * scale).backward()
        clip_grad_norm + optimizer.step()

def _init_with_resource_pool(self, resource_pool, ray_cls_with_init):
  # ...
  rank = -1
  for pg_idx, pg in enumerate(sort_placement_group_by_node_ip(pgs)): # Node
    for local_rank in range(local_world_size): # GPU
      rank += 1
      env_vars = {
        'WORLD_SIZE': str(world_size), 'RANK': str(rank), # More env vars ...
      }
      ray_cls_with_init.update_options(
        {'runtime_env': {'env_vars': env_vars}})
      # ...
      worker = ray_cls_with_init(placement_group=pg,
                                 placement_group_bundle_idx=local_rank)
      self._workers.append(worker)
  # ...

def register(
	dispatch_mode=Dispatch.ALL_TO_ALL,
    execute_mode=Execute.ALL,
    blocking=True,
    materialize_futures=True
):

    def decorator(func):
        @wraps(func)
        def inner(*args, **kwargs):
            if materialize_futures:
                args, kwargs = _materialize_futures(*args, **kwargs)
            return func(*args, **kwargs)

        attrs = {
            'dispatch_mode':  dispatch_mode,
            'execute_mode':   execute_mode,
            'blocking':       blocking,
        }
        setattr(inner, MAGIC_ATTR, attrs)
        return inner

    return decorator

predefined_dispatch_mode_fn = {
    Dispatch.ONE_TO_ALL: {
        'dispatch_fn': dispatch_one_to_all,
        'collect_fn': collect_all_to_all,
    },
    Dispatch.ALL_TO_ALL: {
        'dispatch_fn': dispatch_all_to_all,
        'collect_fn': collect_all_to_all,
    },	
    Dispatch.DP_COMPUTE_PROTO: {
        'dispatch_fn': dispatch_dp_compute_data_proto,
        'collect_fn': collect_dp_compute_data_proto
    },
    # ...
}

# dispatch：把 DataProto 按 world_size 均匀切分，每个 Worker 拿到一片
def dispatch_dp_compute_data_proto(worker_group, *args, **kwargs):
    splitted_args, splitted_kwargs = _split_args_kwargs_data_proto(
        worker_group.world_size, *args, **kwargs
    )
    return splitted_args, splitted_kwargs

# collect：把所有 Worker 返回的 DataProto 片段拼回完整数据
def collect_dp_compute_data_proto(worker_group, output):
    return _concat_data_proto_or_future(output)

def _concat_data_proto_or_future(output):
    o = output[0]
    if isinstance(o, DataProto):
        return DataProto.concat(output)
    elif isinstance(o, ray.ObjectRef):
        return DataProtoFuture.concat(output)
    else:
        raise NotImplementedError

def execute_all_async(self, method_name: str, *args, **kwargs):
    # 这里我们假设，如果 args 和 kwargs 里面所有的参数都是 list，且所有的 list 长度都与 len(self._workers) 一致的话，我们会把
    # list 中的每一个分别发到对应的 worker 上去
    # print(f"execute_all_async: method {method_name}({args}, {kwargs})")
    length = len(self._workers)
    if all(isinstance(arg, list) for arg in args) and all(isinstance(kwarg, list) for kwarg in kwargs.values()):
        if all(len(arg) == length for arg in args) and all(len(kwarg) == length for kwarg in kwargs.values()):
            # print(f"splitting args and kwargs into {length} shards")
            result = []
            for i in range(length):
                sliced_args = tuple(arg[i] for arg in args)
                sliced_kwargs = {k: v[i] for k, v in kwargs.items()}
                remote_call = getattr(self._workers[i], method_name)
                result.append(remote_call.remote(*sliced_args, **sliced_kwargs))
            return result
    
    return [getattr(worker, method_name).remote(*args, **kwargs) for worker in self._workers]

execute_all_async("generate_sequences")  # rollout 阶段
execute_all_async("compute_log_prob")    # 计算 logprob 阶段
execute_all_async("update_actor")        # 训练阶段

def func_generator(self, method_name, dispatch_fn, collect_fn, execute_fn, blocking):

    def func(*args, **kwargs):
        # Step 1: 切分数据
        args, kwargs = dispatch_fn(self, *args, **kwargs)
        # Step 2: 分发给所有 Worker 异步执行
        output = execute_fn(method_name, *args, **kwargs)
        # Step 3: 阻塞等待（可选）
        if blocking:
            output = ray.get(output)
        # Step 4: 聚合结果
        output = collect_fn(self, output)
        return output

    return func

Driver（单进程）
  └── single-controller：顺序调用各 WorkerGroup 的方法
        └── WorkerGroup（multi-controller）：SPMD 调度
              └── dispatch → execute（多 Worker 并行）→ collect

actor_wg.update_actor(data)

{
	"data_source": used to chose reward function,
    "prompt": [{"role": ..., "content": ...}],
    "reward_model": {
        "style": "rule" or "reward",
        "ground_truth": ...
    },
    "extra_info": a dict containing extra information
}

train_dataset.to_parquet(os.path.join(local_save_dir, "train.parquet"))
test_dataset.to_parquet(os.path.join(local_save_dir, "test.parquet"))

python3 -m verl.trainer.main_ppo \
    algorithm.adv_estimator=grpo \
    data.train_files=$HOME/data/geo3k/train.parquet \
    data.val_files=$HOME/data/geo3k/test.parquet \
    //...

custom_reward_function:
  path: /path/to/my_reward.py   # 你的文件路径
  name: my_reward_fn            # 函数名，如果叫 compute_score 可以不填
reward_model:
  reward_manager: naive

def compute_score(data_source, solution_str, ground_truth, extra_info=None):
    if data_source == "openai/gsm8k":
        return gsm8k_score(solution_str, ground_truth)
    elif data_source == "lighteval/MATH":
        return math_score(solution_str, ground_truth)
    else:
        raise NotImplementedError(f"Unknown data_source: {data_source}")

class NaiveRewardManager:
    def __call__(self, data: DataProto):
        for i in range(len(data)):
            data_item = data[i]
            sequences_str = self.tokenizer.decode(
                torch.cat((valid_prompt_ids, valid_response_ids))
            )
            score = self.compute_score(
                data_source=data_item.non_tensor_batch['data_source'],
                solution_str=sequences_str,
                ground_truth=data_item.non_tensor_batch['reward_model']['ground_truth'],
            )
            reward_tensor[i, valid_response_length - 1] = score
        return reward_tensor

class PrimeRewardManager:
    def __call__(self, data: DataProto):
        try:
            scores = asyncio.run(
                parallel_compute_score_async(
                    self.compute_score,
                    sequences_str, ground_truth, data_sources,
                    num_processes=64,
                )
            )
        except asyncio.TimeoutError:
            scores = [0.0] * len(sequences_str)
        # 写回 reward tensor ...
        return reward_tensor

def compute_score(data_sources, solution_strs, ground_truths, extra_infos=None):
	return [0.0] * len(solution_strs)

def compute_score(data_source, solution_str, ground_truth, extra_info=None):
    format_score      = check_format(solution_str)
    correctness_score = check_answer(solution_str, ground_truth)
    length_penalty    = -0.01 * max(0, len(solution_str) - 1000)
    return 0.2 * format_score + 0.7 * correctness_score + 0.1 * length_penalty

class MultiRewardManager:
    def __init__(self, tokenizer, num_examine, compute_score=None):
        self.tokenizer = tokenizer
        self.reward_fns = [
            (compute_format_score,      0.2),
            (compute_correctness_score, 0.7),
            (compute_length_score,      0.1),
        ]

    def __call__(self, data: DataProto):
        # decode、提取 ground_truth ...
        total = sum(
            w * fn(data_source, solution_str, ground_truth, extra_info)
            for fn, w in self.reward_fns
        )
        # 写回 reward tensor ...
        return reward_tensor

def compute_policy_loss_xxx(
    old_log_prob,      # (bsz, response_len) rollout 时采样的 log prob
    log_prob,          # (bsz, response_len) 当前 policy 的 log prob  
    advantages,        # (bsz, response_len) GAE 优势估计
    response_mask,     # (bsz, response_len) 哪些位置是有效 response
    loss_agg_mode,     # "token_mean" / "seq_mean" / ...
    config,
    **kwargs
) -> (loss_scalar, metrics_dict)

actor_rollout_ref:
  actor:
    policy_loss:
      loss_mode: compute_policy_loss_xxx
    entropy_coeff: 0.001
    use_kl_loss: false

def compute_custom_advantage(  
    token_level_rewards: torch.Tensor,  
    response_mask: torch.Tensor,  
    index: np.ndarray = None,  
    config: Optional[AlgoConfig] = None,  
    **kwargs  
) -> tuple[torch.Tensor, torch.Tensor]:  
    """  
    计算自定义优势  
      
    Returns:  
        advantages: 优势值 (bs, response_length)  
        returns: 回报值 (bs, response_length)  
    """  
    # 实现你的优势计算逻辑  
    return advantages, returns

from verl.trainer.ppo.core_algos import register_adv_est, AdvantageEstimator  
  
@register_adv_est("custom_advantage")
def compute_custom_advantage(...):
    pass

algorithm:  
  adv_estimator: custom_advantage

select_keys = [
    "responses", "response_mask", "input_ids",
    "attention_mask", "position_ids",
    "old_log_probs", "advantages",
    "token_level_rewards",   # ← 新增字段
]

# 然后在 forward 循环里取出来传给 loss fn
token_rewards = model_inputs.get("token_level_rewards", None)

pg_loss, pg_metrics = policy_loss_fn(
    old_log_prob=old_log_prob,
    log_prob=log_prob,
    advantages=advantages,
    response_mask=response_mask,
    loss_agg_mode=loss_agg_mode,
    config=self.config,
    token_level_rewards=token_rewards,  # ← 透传
)

def train(self):
    for epochs in range(self.config.epochs):
        for batch in self.dataloader: 
            exp = self.get_experience(batch["prompts"])
            
            for _ in range(self.config.ppo_epochs):
                for mini_batch in split(exp, self.config.ppo_mini_batch_size):
	                for micro_batch in split(mini_batch, micro_batch_size_per_gpu * dp_size):
	                    self.optimizer.zero_grad()
	                    loss = self.compute_loss(micro_batch)
	                    loss.backward()
	                    self.optimizer.step()

                    ┌─────────────────────────────────────────────┐
                    │              ToolAgentLoop.run()            │
                    └─────────────────────────────────────────────┘
                                          │
                              ┌───────────▼───────────┐
                              │    AgentState.PENDING │  ← 准备 prompt
                              └───────────┬───────────┘
                                          │
                           ┌──────────────▼──────────────┐
                           │   AgentState.GENERATING     │  ← LLM 推理
                           └──────────────┬──────────────┘
                                          │
                        ┌─────────────────┼───────────────────┐
                        │                 │                   │
               有 tool_calls        有 interaction        无任何后续
                        │                 │                   │
         ┌──────────────▼──────┐  ┌───────▼────────┐  ┌──────▼──────────┐
         │ PROCESSING_TOOLS    │  │  INTERACTING   │  │   TERMINATED    │
         │ (执行工具, 拼 token) │  │ (获取用户输入)  │  │   (输出结果)     │
         └──────────────┬──────┘  └───────┬────────┘  └─────────────────┘
                        │                 │
                        └────────┬────────┘
                                 │ 回到 GENERATING
                                 ▼
                      (直到 TERMINATED 退出循环)

class AgentData:
    # ── 输入数据 ──────────────────────────────
    messages: list[dict]        # 对话历史（role/content 格式）
    image_data: list[Image]     # 多模态图像
    video_data: list[tuple]     # 多模态视频
    tools_kwargs: dict          # 工具初始化参数
 
    # ── 训练关键字段 ───────────────────────────
    prompt_ids: list[int]       # 完整的 token id 序列（prompt + 所有 response）
    response_ids: list[int]     # 当前轮次的 response token ids
    response_mask: list[int]    # 1=模型生成的 token, 0=工具/用户输入的 token
    response_logprobs: list[float]  # 对应 logprob，用于 PPO 等算法
 
    # ── 统计与奖励 ─────────────────────────────
    turn_scores: list[float]    # 每轮 interaction 产生的奖励
    tool_rewards: list[float]   # 每次工具调用产生的奖励
    user_turns: int
    assistant_turns: int
 
    # ── 临时状态 ───────────────────────────────
    tool_calls: list[FunctionCall]  # 当前轮次解析出的工具调用
    extra_fields: dict              # 自定义扩展字段

while state != AgentState.TERMINATED:
    if state == AgentState.PENDING:
        state = await self._handle_pending_state(agent_data, sampling_params)
    elif state == AgentState.GENERATING:
        state = await self._handle_generating_state(agent_data, sampling_params)
    elif state == AgentState.PROCESSING_TOOLS:
        state = await self._handle_processing_tools_state(agent_data)
    elif state == AgentState.INTERACTING:
        state = await self._handle_interacting_state(agent_data)
    else:
        logger.error(f"Invalid state: {state}")
        state = AgentState.TERMINATED

async def _handle_pending_state(self, agent_data, sampling_params):
    prompt_ids = await self.apply_chat_template(
        agent_data.messages,
        tools=schemas,      # 把工具 schema 嵌入 prompt
        images=...,
        videos=...,
    )
    agent_data.prompt_ids = prompt_ids
    return AgentState.GENERATING

async def _handle_generating_state(self, agent_data, sampling_params):
    # 1. 调用 LLM 生成
    output = await self.server_manager.generate(
        prompt_ids=agent_data.prompt_ids, ...
    )
 
    # 2. 累积 token 序列
    agent_data.response_ids = output.token_ids
    agent_data.prompt_ids += agent_data.response_ids    # prompt 不断增长
    agent_data.response_mask += [1] * len(agent_data.response_ids)  # 模型输出 mask=1
 
    # 3. 检查终止条件
    if len(agent_data.response_mask) >= self.response_length: return TERMINATED
    if agent_data.assistant_turns >= self.max_assistant_turns: return TERMINATED
 
    # 4. 解析 tool_calls
    _, agent_data.tool_calls = await self.tool_parser.extract_tool_calls(...)
 
    # 5. 决定下一状态
    if agent_data.tool_calls:      return PROCESSING_TOOLS
    elif self.interaction_config:  return INTERACTING
    else:                          return TERMINATED

tool_calls (并行执行，上限 max_parallel_calls)
    │
    ▼
asyncio.gather(*tasks)    ← 并行调用所有工具
    │
    ▼
构造 tool 消息（text / 多模态）
    │
    ▼
apply_chat_template(tool_messages)  ← 转为 token ids
    │
    ▼
prompt_ids += tool_response_ids
response_mask += [0] * len(tool_response_ids)   ← 工具输出 mask=0
    │
    ▼
return GENERATING   ← 继续让模型生成

async def _call_tool(self, tool_call, tools_kwargs, agent_data):
    try:
        instance_id, _ = await tool.create(...)
        response, reward, res = await tool.execute(instance_id, tool_args, agent_data=agent_data)
    except Exception as e:
        return ToolResponse(text=f"Error when executing tool: {e}"), 0.0, {}
    finally:
        await tool.release(instance_id)   # 确保资源释放

if len(tool_response_text) > self.max_tool_response_length:
    # 三种截断模式：left / right / middle

async def _handle_interacting_state(self, agent_data):
    should_terminate, response, reward, metrics = \
        await agent_data.interaction.generate_response(...)
 
    # 将用户回复追加进去，mask=0（非模型生成）
    response_ids = await self.apply_chat_template([{"role": "user", "content": response}])
    agent_data.prompt_ids += response_ids
    agent_data.response_mask += [0] * len(response_ids)
 
    if should_terminate: return TERMINATED
    else: return GENERATING

tool_list = initialize_tools_from_config(tool_config_path)
self.tools = {tool.name: tool for tool in tool_list}
self.tool_schemas = [tool.tool_schema.model_dump(...) for tool in tool_list]

tool_selection = extra_info.get("tool_selection")
if tool_selection:
    agent_data._active_tools = {name: self.tools[name] for name in tool_selection}
    agent_data._active_tool_schemas = [...]
else:
    agent_data._active_tools = self.tools   # 使用全部工具

actor_rollout_ref:
  rollout:
    multi_turn:
      format: hermes

# verl/experimental/agent_loop/tool_parser.py
import regex, json
from verl.experimental.agent_loop.tool_parser import ToolParser, FunctionCall

@ToolParser.register("my_format")
class MyModelToolParser(ToolParser):
    def __init__(self, tokenizer):
        self.tokenizer = tokenizer
        # 根据你模型实际输出的格式定义正则
        self.pattern = regex.compile(
            r"<tool_call>(.*?)</tool_call>", regex.DOTALL
        )

    async def extract_tool_calls(
        self, response_ids: list[int]
    ) -> list[FunctionCall]:
        text = self.tokenizer.decode(response_ids, skip_special_tokens=False)
        matches = self.pattern.findall(text)
        calls = []
        for m in matches:
            try:
                obj = json.loads(m)
                calls.append(FunctionCall(
                    name=obj["name"],
                    arguments=json.dumps(obj["arguments"], ensure_ascii=False)
                ))
            except Exception:
                pass   # 格式解析失败，忽略该次调用
        return calls

format: my_format   # 对应 @ToolParser.register 的 key

AgentLoopOutput(
    prompt_ids   = prompt_ids,                         # 纯 prompt 部分
    response_ids = response_ids[:self.response_length], # 截断到最大长度
    response_mask= response_mask[:self.response_length],
    response_logprobs = ...,
    num_turns    = user_turns + assistant_turns + 1,
    metrics      = {...},
    extra_fields = {
        "turn_scores": [...],   # 每轮 interaction 奖励
        "tool_rewards": [...]   # 每次工具调用奖励
    }
)

# search_tool_config.yaml
tools:
  - class_name: verl.tools.search_tool.SearchTool  # 或你的自定义类
    config:
      retrieval_service_url: http://127.0.0.1:8000/retrieve
      num_workers: 120
      rate_limit: 120
      timeout: 30

actor_rollout_ref:  
  rollout:
	name: vllm
    mode: async  # 异步模式，避免GPU空闲  
    multi_turn:  
      enable: True  
      max_user_turns: 3  # 最大用户轮次  
      max_assistant_turns: 3  # 最大助手轮次 
    tool_kwargs:
      tools_config_file: ./config/tool_config/search_tool_config.yaml

{
	"data_source": used to chose reward function,
    "prompt": [{"role": ..., "content": ...}],
    "reward_model": {
        "style": "rule" or "reward",
        "ground_truth": ...
    },
    "extra_info": a dict containing extra information,
    "agent_name": "tool_agent"
}