跳转至

main_eval.py — 这是一个离线评估脚本

文件概述

模块路径: verl.trainer.main_eval

这是一个离线评估脚本,用于在训练完成后,对已经生成好的回复文件进行奖励评估。它读取包含生成结果的 parquet 文件,使用自定义的奖励函数计算每条回复的分数,并按数据源分组统计平均分。

在训练流程中的位置

这个文件不在训练循环中,而是一个独立的评估工具。通常在训练完成后、或需要对特定 checkpoint 的生成结果进行离线打分时使用。

关键代码讲解

1. Ray Remote 评分函数

@ray.remote
def process_item(config, data_source, response_lst, reward_data):
    reward_fn = get_custom_reward_fn(config)
    ground_truth = reward_data["ground_truth"]
    score_lst = [reward_fn(data_source, r, ground_truth) for r in response_lst]
    return data_source, np.mean(score_lst)

这个函数被 @ray.remote 装饰,意味着它会被调度到 Ray 集群的某个 Worker 上并行执行。对于每个数据项: - 加载自定义奖励函数 - 获取标准答案 (ground_truth) - 对所有回复逐一打分 - 返回数据源标识和平均分

2. Hydra 评估主入口

@hydra.main(config_path="config", config_name="evaluation", version_base=None)
def main(config):
    # 1. 加载数据
    local_path = copy_to_local(config.data.path, ...)
    dataset = pd.read_parquet(local_path)
    responses = dataset[config.data.response_key]
    data_sources = dataset[config.data.data_source_key]
    reward_model_data = dataset[config.data.reward_model_key]

    # 2. 初始化 Ray
    if not ray.is_initialized():
        ray.init(...)

    # 3. 创建并行评估任务
    remote_tasks = [
        process_item.remote(config, data_sources[i], responses[i], reward_model_data[i])
        for i in range(total)
    ]

    # 4. 收集结果
    data_source_reward = defaultdict(list)
    with tqdm(total=total) as pbar:
        while len(remote_tasks) > 0:
            done_ids, remote_tasks = ray.wait(remote_tasks)
            for result_id in done_ids:
                data_source, score = ray.get(result_id)
                data_source_reward[data_source].append(score)
                pbar.update(1)

    # 5. 计算并打印按数据源分组的平均分
    metric_dict = {}
    for data_source, rewards in data_source_reward.items():
        metric_dict[f"test_score/{data_source}"] = np.mean(rewards)
    print(metric_dict)

使用 ray.wait() 实现了动态结果收集:不需要等待所有任务完成,而是逐批处理已完成的任务,配合 tqdm 进度条显示。

核心类/函数列表

名称 类型 作用
process_item Ray remote function 对单个数据项的所有回复进行打分
main(config) function Hydra 入口,协调离线评估流程

数据流和调用关系

parquet 文件 (包含 prompt + 生成的回复 + ground truth)
      |
      v
main() --> pd.read_parquet()
      |
      +-- ray.init()
      |
      +-- process_item.remote() x N  (并行评估)
      |       |
      |       +-- get_custom_reward_fn()  (reward.py)
      |       +-- reward_fn(data_source, response, ground_truth)
      |
      +-- ray.wait() / ray.get()  (收集结果)
      |
      +-- 按 data_source 分组统计 --> print(metric_dict)

小结

main_eval.py 是一个简洁但实用的离线评估工具,利用 Ray 的分布式能力进行并行评分。它与训练流程解耦,适合在训练后独立运行来评估模型生成质量。