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 的分布式能力进行并行评分。它与训练流程解耦,适合在训练后独立运行来评估模型生成质量。