跳到主要内容

数据截至 (上游 commit 932e1f2f4c5a)

04 · 与训练管线的集成

这一章讲什么: lighteval 怎么嵌进训练工作流——训到一半的 checkpoint 直接评、分数写到训练曲线旁边;以及输出层(details、hash、离线重算、inspect-ai 新入口)。读完你能判断「训练中评测」适不适合照抄进自己的栈。

1. 它要解决的小问题​

离线评测(训完再评)有两个痛点:

  • 反馈太慢:训练几天后发现第 2 天就过拟合了,白费算力。
  • 分数不在一处:评测分数在评测日志里,loss 曲线在 TensorBoard 里,对不齐步数。

HF 的做法是让评测器直接认识训练框架的产物:nanotron 的 checkpoint 目录里本来就有模型配置、分词器配置、并行配置——lighteval 读过来,原样重建模型并加载权重,评完把分数按 checkpoint 的训练步数写回 TensorBoard。

2. 直觉:checkpoint 即配置,评测即训练的一环​

两个设计决定让这件事成立:

  1. nanotron checkpoint 的配置 YAML 就是模型定义。 main_nanotron 入口(src/lighteval/main_nanotron.py:42-121)吃两个 YAML:checkpoint 自带的 config(模型结构/分词器/训练元信息)+ 一份 lighteval 专属配置(评哪些任务、推送到哪)。前者本来就躺在 checkpoint 旁边,训到哪存到哪。
  2. LightEvalConfig 设计为「可覆盖」。 它的 docstring 明说:所有字段可选,因为可以「训完后用新参数覆盖存下来的配置」(src/lighteval/models/nanotron/nanotron_model.py:133-144)。训练脚本里嵌一份默认评测配置,训中自动评、训后想换任务再手动评,共用同一条代码路径。

直觉一句话:传统做法是「训练完,导出模型,再配一套评测环境」;lighteval 是「评测配置跟着训练配置走,checkpoint 落盘即评测就绪」。

3. 图示:训练中评测的数据流​

nanotron 训练 (SLURM 集群)
   │  存 checkpoint(含 config.yaml)
   ▼
lighteval nanotron <checkpoint.yaml> <lighteval.yaml>
   │  ① 读 checkpoint 配置 → 重建模型结构
   │  ② load_weights 加载权重 → model.eval()
   │  ③ 走标准 Pipeline(第 1~3 章)
   ▼
EvaluationTracker
   ├─► 本地:results JSON + details parquet
   ├─► TensorBoard:按 checkpoint 的 step 写点
   └─► HF Hub:results/details 推到 org 仓库

怎么读这张图:左边是训练侧的产物,右边是评测侧的输出;中间的 lighteval nanotron 命令是唯一胶水。job 调度上它认 SLURM_JOB_ID(main_nanotron.py:99),天然为集群场景设计。

4. 原理演示:两份配置的合流​

# 示意,非源码
# checkpoint 自带的 config.yaml(训练框架写的):
nanotron_yaml = {
    "model": {...},      # 模型结构
    "tokenizer": {...},  # 分词器
    "general": {...},    # run 名、当前 step —— 后面写 TensorBoard 要用
}

# lighteval_config.yaml(人写的,见 examples/nanotron/):
lighteval_yaml = {
    "tasks": {"tasks": "lighteval|gsm8k|5", "max_samples": 10},
    "parallelism": {"dp": 1, "pp": 1, "tp": 1},   # 评测时的并行方式
    "logging": {"output_dir": "outputs", "push_to_hub": False},
}

# 合流:训练侧配置 + 评测侧配置 → 一个 FullNanotronConfig
config = FullNanotronConfig(lighteval_config, model_args, tokenizer_args, general_args)

重点:模型结构信息不重复声明——评测用的是 checkpoint 里训练框架自己写的那份,杜绝「训练用 A 配置、评测用 B 配置」的漂移。

5. 真实实现​

5.1 nanotron 入口:两个 YAML 进,一条 Pipeline 出​

nanotron()(src/lighteval/main_nanotron.py:42-121)的流程:

  1. 读 checkpoint YAML,用 nanotron 的 get_config_from_dict 折成 ModelArgs/TokenizerArgs/GeneralArgs(main_nanotron.py:68-79);
  2. 读 lighteval YAML 成 LightEvalConfig(含 parallelism dp/pp/tp,模板见 examples/nanotron/lighteval_config_override_template.yaml);
  3. 合成 FullNanotronConfig(nanotron_model.py:148-153),建 EvaluationTracker 时把 nanotron_run_info=nanotron_config.nanotron_general 传进去(main_nanotron.py:85-95)——训练元信息由此进入输出层;
  4. 走标准 Pipeline,launcher_type=ParallelismManager.NANOTRON(main_nanotron.py:97-115)。

并行初始化在 Pipeline._init_parallelism_manager:nanotron 路会建 ParallelContext(tp/pp/dp)并先做 test_all_gather 冒烟测通信(src/lighteval/pipeline.py:157-175)。

5.2 NanotronLightevalModel:按训练时的分片方式重建模型​

NanotronLightevalModel.__init__(nanotron_model.py:163-308)与 transformers 后端的差别在:模型不是 from_pretrained 来的,而是用 nanotron 的 build_model 按训练时的并行配置重新搭骨架,再 load_weights 填权重(nanotron_model.py:263-295)——因为 checkpoint 里的权重是按 tp/pp 分片存的,只有同样的并行拓扑才能读回来。加载完 model.eval()(nanotron_model.py:295)。

限制也直接:PP > 1 抛错「PP parallelism is not supported yet」(nanotron_model.py:206-208)——流水线并行下「没有输出的 PP rank 怎么同步」还没实现(代码注释原话)。

请求侧的 greedy_until/loglikelihood 结构与 transformers 后端同构(动态 batch、GenerativeTaskDatasetNanotron),入口在 nanotron_model.py:946 与 :488。

5.3 输出层:EvaluationTracker 的四个目的地​

EvaluationTracker(src/lighteval/logging/evaluation_tracker.py:95)内部是五个专职 logger(details / metrics / versions / general config / task config,evaluation_tracker.py:155-160)。save()(evaluation_tracker.py:251-295)按开关分发:

目的地内容函数
本地磁盘results_*.json + 每任务一个 details parquetsave_results(evaluation_tracker.py:306)、save_details(:358)
HF Hubresults + details 推到 org 的 dataset 仓库push_to_hub(evaluation_tracker.py:391)
TensorBoard聚合指标写 event,同步到 Hub 仓库push_to_tensorboard(evaluation_tracker.py:666)
wandb/trackio聚合指标push_to_wandb(evaluation_tracker.py:298-304)

最关键的一行在 TensorBoard 路:global_step = self.nanotron_run_info.step(evaluation_tracker.py:679-684)——nanotron 场景下,评测分数写进 event 文件时用的 x 轴是** checkpoint 的训练步数**。于是 eval 曲线和 loss 曲线天然对齐,这就是「评测嵌进训练监控」的落点。非 nanotron 场景 step 为 0。

writer 是 HFSummaryWriter(evaluation_tracker.py:37 import,:690-696 使用)——TensorBoard event 直接落在 Hub 仓库里,训练队友打开 Hub 就能看。

5.4 details 与 hash:分数之外的复现性​

DetailsLogger(src/lighteval/logging/info_loggers.py:138)逐样本记录题面、完整 prompt、模型原始输出、分数。聚合时(info_loggers.py:277-305)对每个任务的四样东西分别做 xxhash:

  • hash_examples(题面)、hash_full_prompts(完整 prompt)、hash_input_tokens、hash_cont_tokens;

再把各任务的 hash 按任务名排序串起来做一次总 hash(info_loggers.py:300-305)。

用途:两次 run 只要任一环节的 prompt 变了,hash 就对不上——比对人眼 diff 几千条 prompt 快得多。这也是「我的分数能不能和别人那张表比」的第一道自检。

5.5 离线重算:回答存过就不必再跑模型​

Pipeline.evaluate 开头检查 load_responses_from_details_date_id(src/lighteval/pipeline.py:278-287):给了就从之前的 details 里把 ModelResponse 重建出来(_load_responses_from_details,pipeline.py:401-424),跳过模型调用直接进指标计算。场景:换指标定义、修指标 bug 后重出分,不用重新烧 GPU。找不到文件会降级为正常跑模型并打 warning(pipeline.py:282-284)。限制:只支持单一采样方式(pipeline.py:407-410)。

5.6 第二条路:inspect-ai 入口与 baseline 工具​

本 commit 的 CLI 有一条较新的 lighteval eval(src/lighteval/__main__.py:73,README 标注为 preferred):它把 lighteval 的任务定义翻译成 inspect-ai 的 Task——get_inspect_ai_task(src/lighteval/main_inspect.py:49-92)用任务配置里的 sample_fields/solver/scorer 三字段搭桥,没配这三字段的任务会报「不被 inspect_ai 支持」(main_inspect.py:56-59)。第 1 章 gsm8k 例子里那三个字段就是为这条路准备的。它绕过了本 teardown 第 2 章的后端体系,走 inspect-ai 自己的执行器,适合评远程 API 模型。

另有 lighteval baseline(src/lighteval/main_baseline.py:33-52):不加载任何模型,按「多选题随机猜 = 1/选项数,其余指标记 0」算基线分,给任务设计 sanity check 用。

6. 坑​

  • push_to_tensorboard 依赖 nanotron 在场。 没装 nanotron 会直接跳过并打 warning(evaluation_tracker.py:673-675)——即便你只是想给 vLLM 的 run 推 tensorboard,也得装 nanotron(inferred:原因是 step/run 信息从 nanotron 配置读)。
  • nanotron 后端 PP>1 不支持(nanotron_model.py:206-208),见 §5.2。
  • details 数据集强制全字符串列。 存 details 时对嵌套对象「强制 cast 成 str 以避免格式问题」(evaluation_tracker.py:261-262 注释)——读 details 做后处理时记得反序列化。
  • 离线重算只支持单采样方式(pipeline.py:407-410),混合指标的任务不能用。
  • inspect-ai 路只覆盖配了 sample_fields 的任务(main_inspect.py:56-59);内置任务并非全部配齐,跑之前先用 lighteval tasks 查。
  • 推 Hub 需要 org。 push_to_hub/push_to_tensorboard 任一开而 results_org 为空,直接抛异常(evaluation_tracker.py:166-169)。

7. 代码地图(本章 + 全局)​

主题文件路径符号名
nanotron 入口src/lighteval/main_nanotron.pynanotron
nanotron 配置src/lighteval/models/nanotron/nanotron_model.pyLightEvalConfig、FullNanotronConfig、NanotronLightevalModel
并行初始化src/lighteval/pipeline.pyPipeline._init_parallelism_manager、test_all_gather
输出总控src/lighteval/logging/evaluation_tracker.pyEvaluationTracker.save、push_to_hub、push_to_tensorboard
逐样本记录与 hashsrc/lighteval/logging/info_loggers.pyDetailsLogger.log、DetailsLogger.aggregate
离线重算src/lighteval/pipeline.pyPipeline._load_responses_from_details
inspect-ai 桥src/lighteval/main_inspect.pyget_inspect_ai_task、eval、bundle
baselinesrc/lighteval/main_baseline.pybaseline
CLI 注册src/lighteval/__main__.pyapp(typer)
训练侧配置模板examples/nanotron/lighteval_config_override_template.yaml—