Skip to content

feat: format reward experiment + RTX 4060 support - #190

Open
Tassel51 wants to merge 1 commit into
PeterGriffinJin:mainfrom
Tassel51:main
Open

feat: format reward experiment + RTX 4060 support#190
Tassel51 wants to merge 1 commit into
PeterGriffinJin:mainfrom
Tassel51:main

Conversation

@Tassel51

@Tassel51 Tassel51 commented Sep 1, 2026

Copy link
Copy Markdown

大三学生复现 Search-R1 的时候做的一些小改进,分享给大家~

改了啥

1. 格式奖励对比实验 (train_compare.py)

复现的时候发现原始方案只有 EM 奖励,中间过程完全没有反馈,就加了个格式奖励试试。

奖励公式:EM(0.7) + think(0.1) + search(0.1) + answer(0.1)

实测结果(Qwen2.5-3B, 40道NQ题):

  • EM 准确率:60% → 65% (+5%)
  • think 标签使用率:40% → 65%
  • search 标签使用率:62% → 72%

主要是想验证格式奖励能不能让模型更规范地用标签,没想到准确率也涨了一点哈哈。

2. RTX 4060 适配 (infer_4060.py)

我显卡是 4060 8GB 的,原版推理脚本跑起来有问题,就改了一个简化版:

  • 用内联检索代替 threading server(避免 segfault)
  • 支持 chat template
  • 显存占用优化到 5.75GB

3. 其他小工具

  • reward_format.py: 格式奖励函数
  • plot_training.py: 训练曲线可视化
  • download_with_retry.py: 带重试的下载脚本(网不好时候用的)

测试环境

  • GPU: RTX 4060 Laptop (8GB)
  • 模型: Qwen2.5-3B
  • 数据: 40道 NQ 问题 + 75篇文档

代码都在 train_compare.py 里,直接 python train_compare.py 就能跑,不需要额外配置。

如果有什么问题欢迎提 issue~

- Add train_compare.py: Format reward vs baseline comparison experiment
  - Pure sampling-based comparison (no SFT training needed)
  - Works on RTX 4060 with 8GB VRAM
  - Tests EM accuracy and format compliance
- Add reward_format.py: Format reward functions
  - think/search/answer tag detection
  - Combined reward scoring
- Add infer_4060.py: Simplified inference for RTX 4060
  - Inline search (no threading issues)
  - Chat template support
- Add plot_training.py: Training visualization
- Add download_with_retry.py: Download utility with retry logic

Tested on Qwen2.5-3B with 40 NQ questions.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant