robojudge-finetune

RoboJudge Finetune — Qwen3.6-35B-A3B 具身世界模型评判器

把 Qwen3.6-35B-A3B 微调成一个具身操作视频的自动评判器，在 MLLM-as-Embodied-World-Judge 上给两个维度打分：

本仓库包含完整的训练与评测代码，以及 1 个 epoch 的 checkpoint（checkpoint-epoch1/）。

---

一、结果

在 bench/test_clean_v2.jsonl（855 条人工标注）上，指标为 Pearson 相关系数， 全部用 benchmark 官方的 bench/eval_metrics.py 计算。

排行榜（按 PA-Pearson 排序）

表中所有 judge 都是在同一份 test_clean_v2 上重新计算的。官方 LEADERBOARD.md 用的是 875 条的 gold_875，两者数值不可直接混用 （例如 gemini-3.5-flash 在 gold_875 上是 0.530，在这里是 0.5612）。

微调带来的提升

逐 epoch 结果：训 1 个 epoch 就够

训练 loss 从 2.78 一路降到 0.12，但判别能力从第 2 个 epoch 起就在退化。 后期模型逐渐记住"这类视频大概几分"的绝对刻度，却丢失了跨样本的相对判别力， 而 Pearson 衡量的正是后者。第 4、5 个 epoch 甚至开始破坏 JSON 输出格式 （解析失败从 0 涨到 8）。

结论：用 checkpoint-epoch1/，别多训。

未微调模型不能直接当评判器

未微调的 35B 在 855 条里有 662 条给不出可解析的分数——即使把生成长度放宽到 2048 token（微调模型的 4 倍），它仍然写长篇散文而不给结论。微调后 855/855 全部可解析。

SFT 最大的收益是可用性，其次才是准确率。

---

二、环境

硬件

单节点 8 × NVIDIA B300（268 GB）。全量微调 35B 的显存账：

FSDP2 分片后每卡约 70 GB，单节点 2144 GB 显存绰绰有余。

实测速度：2.15 秒/步，1620 步（5 epoch）共 1 小时 22 分。

软件

容器 nvcr.io/nvidia/pytorch:25.09-py3，在其中建 venv：

sbatch slurm/setup_env.slurm

该脚本做的事（--system-site-packages 继承容器的 torch，不另装）：

python -m venv --system-site-packages $P/venv_lf
source $P/venv_lf/bin/activate

# transformers >= 5.9 是 Qwen3_5MoeForConditionalGeneration 的下限
python -m pip install "transformers>=5.9,<6" accelerate deepspeed

# --no-deps：否则 LLaMA-Factory 的 pin 会把 transformers 降回去
python -m pip install --no-deps -e <LLaMA-Factory 路径>

python -m pip install av decord peft trl datasets sentencepiece \
    tiktoken einops pydantic fire omegaconf tensorboard matplotlib

⚠️ 三个必须注意的坑

1. 版本 pin 要绕过，但依赖版本要压回去

LLaMA-Factory 的 pyproject.toml 写着 transformers<=5.8.0，这个上限早于 Qwen3_5MoeForConditionalGeneration 这个架构，照着装模型根本加载不了。所以：

• 装 transformers>=5.9，并用 --no-deps 安装 llamafactory
• 运行时设 export DISABLE_VERSION_CHECK=1（代码里还有第二道断言）
• 但 trl / peft / datasets 必须压回 pin 范围内，否则 AutoModelForCausalLMWithValueHead 之类的符号会消失：

python -m pip install --no-deps "trl==0.24.0" "peft==0.18.1" "datasets==4.0.0"

2. 不要装 torchaudio / causal_conv1d / flash-linear-attention

这三个包在 NVIDIA 容器里都会因 ABI 不匹配而炸：

--no-deps 只解决依赖解析，解决不了二进制兼容。带 C/CUDA 扩展的包在这个容器里 要么现场编译，要么别装。

torchaudio 是 mm_plugin.py 的硬依赖，用 patches/mm_plugin_lazy_torchaudio.patch 改成惰性导入即可（我们的数据没有音频）。

3. enable_liger_kernel 必须关闭

Liger 的算子接收普通 tensor，而 FSDP2 的参数是 DTensor，混用直接报 aten.mm.default: got mixed torch.Tensor and DTensor。

---

三、数据准备

1. 下载

sbatch slurm/download.slurm     # 模型 67 GB + 数据集

数据集 data/ 下有 273,110 个文件，但 SFT 只用到其中 10,364 个（3.8%）。 scripts/fetch_needed.py 按精确文件名列表拉取，比 snapshot_download 的 allow_patterns 快得多——后者要拿 10,364 条 pattern 对 273,110 个文件做 fnmatch，约 28 亿次匹配，实测跑十分钟一个请求都发不出去。

HF 对单仓库限流 1000 次 API 请求 / 5 分钟（≈3.3 次/秒），并发调到 8 以上 只会被限流退避，反而更慢。fetch_needed.py 里 max_workers=3 是贴着这个配额定的。

2. 重写媒体路径

原始 SFT json 里存的是构建机上的绝对路径：

/datapool/pengruotian/robojudge/datasets/
    mllm-as-embodied-world-judge_new_continue/data/agibot_world/.../x.mp4

注意 data/ 上一级目录叫 ..._new_continue，和仓库名对不上，所以只能从 data/ 这一段开始匹配：

python scripts/prepare_sft_data.py

输出到 sft_data/，原始下载保持不动。脚本同时会报告有多少媒体文件缺失， 可当作完整性检查器用（正常应为 0）。

3. 校验

python scripts/check_correspondence.py          # 训练数据 ↔ 视频文件

我们的校验结果：

• 训练媒体 10,364 个，缺失 0
• 测试媒体 1,710 个，缺失 0
• 训练集与测试集重叠 0（无数据泄漏）

---

四、训练

sbatch slurm/train.slurm                       # 默认 HEAD=merged
HEAD=merged SMOKE=1 sbatch slurm/train.slurm   # 16 样本冒烟测试

或在交互式节点上跑（便于调试，省去每次改代码都排一次队）：

salloc --no-shell -N1 -t 04:00:00 -p batch -A <account>
./scripts/run_interactive.sh <jobid> merged 0 0

关键配置（configs/qwen36_moe_merged_full_sft.yaml）

model_name_or_path: <本地 Qwen3.6-35B-A3B 路径>
template: qwen3_5_nothink          # 3.6 的实现族仍是 qwen3_5
dataset: embodied_world_judge_sft_ia,embodied_world_judge_sft_pa   # 合并训练

finetuning_type: full
freeze_vision_tower: true          # 5182 条样本训不动 27 层 ViT
freeze_multi_modal_projector: true

cutoff_len: 4096
video_maxlen: 16
image_max_pixels: 65536
video_max_pixels: 16384

per_device_train_batch_size: 4
gradient_accumulation_steps: 1     # 有效 batch = 8 卡 × 4 × 1 = 32
learning_rate: 1.0e-5
num_train_epochs: 5.0
lr_scheduler_type: cosine
warmup_ratio: 0.1

save_strategy: epoch
save_only_model: true              # 见下方说明
enable_liger_kernel: false
gradient_checkpointing: true
optim: adamw_torch_fused

为什么这么配

模型叫 3.6，代码里是 qwen3_5_moe。 config.json 里 model_type: qwen3_5_moe、architectures: [Qwen3_5MoeForConditionalGeneration]。 3.6 是产品版本号，实现族仍是 qwen3_5，所以 template、FSDP wrap 类名、 kernel 全都用 qwen3_5 那一套。

save_only_model: true 非常重要。 不加这个，每个 checkpoint 会存 435 GB（139 GB 模型 + 139 GB pytorch_model_fsdp.bin + 35 GB+ optimizer.bin）， 写一次要 13.6 分钟，训练全程暂停等它。5 个 checkpoint 就是 2.2 TB / 68 分钟， 比训练本身还久。开了之后：130 GB / 2.4 分钟，快 5.8 倍。 优化器状态只有续训才需要，而我们存 checkpoint 是为了评测。

reshard_after_forward: false（在 configs/fsdp2_qwen36_moe.yaml）。 这个模型 35B 总参数但每 token 只激活 3B，瓶颈是搬运权重而非算术。默认设置 每步要 all-gather 70 GB 两次（前向一次、反向一次），关掉后省掉一整次集合通信， 代价是每卡常驻 70 GB——268 GB 的卡完全吃得下。

每卡 batch 4 而不是 1。 通信受限时，一次 all-gather 服务的样本越多越划算。 per_device 1→4 把通信开销摊薄 4 倍。配合 gradient_accumulation_steps 相应调整， 有效 batch 始终锁在 32。

梯度检查点保持开启。 反直觉但正确：GPU 等通信时本来就空闲，重算前向几乎不额外 花时间，却能腾出显存支撑更大的 batch。

多节点时必须同步调整 accum。 加节点是为了省墙钟时间，不是改变优化过程。 train.slurm 里有自动换算：

GA=$(( 32 / (8 * SLURM_NNODES * PDB) ))

保证 节点数 × 8 × per_device × accum 恒等于 32。否则 2 节点会让有效 batch 翻倍到 64、更新次数减半，那是另一次训练，不是同一次训练跑得更快。

---

五、评测

sbatch slurm/eval.slurm            # 所有 epoch checkpoint
sbatch slurm/eval_baseline.slurm   # 未微调 base，作为 epoch 0
python scripts/report_eval.py      # 出 Pearson 曲线
python scripts/plot_distribution.py --epoch 1   # 出分布四联图
./scripts/report_all.sh            # 一次性全出

设计要点

8 卡数据并行。 每张卡加载一份完整模型（bf16 仅 70 GB）、处理 1/8 样本。 单卡串行跑 8,550 次生成要好几小时，8 路分片后约 24 分钟一个 checkpoint。 分片用交错切分（rows[shard::8]）而非连续切分——测试集样本按 dataset 聚集， 连续切会让每片数据构成不同，单片耗时就无法用来推算整体。

按 item_id 续跑。 分片文件用追加写，启动时扫描该 epoch 的所有分片文件 判断哪些样本已完成。被抢占时最多损失一条，而不是整个 epoch。 （我们实际被抢占过 3 次，这个机制救回了大量时间。）

指标用 benchmark 官方脚本。 report_eval.py 调用 bench/eval_metrics.py，不用自己的实现，这样 pa_pearson / ia_pearson 和排行榜逐位可比。

基线要给足生成长度。 未微调模型写散文不收尾，512 token 下 100% 被截断， 会得出"基线 0% 可解析"的假象——那报的是 token 上限，不是模型能力。 eval_baseline.slurm 里给 2048，并加了从散文中提取分数的兜底解析， 这样能把"输出规范性"和"判断能力"两件事分开看。

---

六、Checkpoint

checkpoint-epoch1/       # 推荐使用，PA-Pearson 0.5855 / IA-Pearson 0.5816

加载方式：

from transformers import AutoModelForImageTextToText, AutoProcessor

processor = AutoProcessor.from_pretrained("yqi19/robojudge-finetune",
                                          subfolder="checkpoint-epoch1")
model = AutoModelForImageTextToText.from_pretrained(
    "yqi19/robojudge-finetune", subfolder="checkpoint-epoch1",
    dtype="bfloat16", device_map="auto")

Prompt 必须和训练时逐字一致（系统提示 + 用户提示）， scripts/evaluate_ckpt.py 里是直接从训练数据的第一条记录里取的， 就是为了避免手抄产生的空格或措辞漂移——那会伪装成质量回退。

---

七、已知不足与后续方向

系统性低估。 epoch 1 的 PA 预测均值 2.28，gold 均值 3.03；855 条里 低估 514 条（60.1%），高估仅 81 条。PA 一次 5 分都没打过，而 gold 里有 136 个 5 分。 原因大概率是训练数据 prompt 里那句 "Be conservative: reserve 5 for full, correct completion" 被学得过了头。

Pearson 测的是相关性，不受均值偏移影响，所以 0.5855 是硬的；但 exact accuracy 只有 30.4% 明显是被这个拖累的。做一次分位数校准（把预测分布映射到 gold 分布） 可以在完全不改变 Pearson 的前提下提升 exact 和 QWK——单调变换不改相关系数。 这是几乎白捡的收益，尚未实施。

IA 仍落后于最强闭源模型。 gpt-5.5 的 IA 是 0.6731，我们是 0.5816。 PA 已经第一，IA 还有空间。

推理成本。 微调后模型输出紧凑 JSON（约 200~400 token 即 EOS）， 比未微调的散文式输出快约 15 倍。这在实际部署中是很实在的收益。

---

八、目录结构

configs/
  qwen36_moe_merged_full_sft.yaml   IA+PA 合并训练（推荐）
  qwen36_moe_ia_full_sft.yaml       仅 IA
  qwen36_moe_pa_full_sft.yaml       仅 PA
  fsdp2_qwen36_moe.yaml             FSDP2 分片配置
slurm/
  setup_env.slurm                   容器内建 venv
  download.slurm                    下载模型与数据
  train.slurm                       训练（含自动续跑链）
  eval.slurm / eval_baseline.slurm  评测
scripts/
  download_all.py / fetch_needed.py       下载
  check_correspondence.py                 数据↔视频校验
  prepare_sft_data.py                     重写媒体路径
  run_interactive.sh                      交互式节点训练
  eval_shard.sh / run_eval_all.sh         分片评测
  evaluate_ckpt.py                        单 checkpoint 推理
  report_eval.py                          Pearson 曲线
  plot_distribution.py                    分布/混淆矩阵/误差分档
  plot_loss.py                            训练 loss 曲线
  make_deck.py                            生成汇报 PPT
patches/
  mm_plugin_lazy_torchaudio.patch   LLaMA-Factory 必需改动
data/
  dataset_info.json                 LLaMA-Factory 数据集注册
checkpoint-epoch1/                  推荐 checkpoint

---

引用

• 数据集与 benchmark：HuggingFriends/mllm-as-embodied-world-judge
• 基础模型：Qwen/Qwen3.6-35B-A3B
• 训练框架：LLaMA-Factory