346 lines
12 KiB
Markdown
346 lines
12 KiB
Markdown
# SearchShortQA：搜索增强短问答强化学习系统
|
||
|
||
## 业务场景
|
||
|
||
面向企业知识库、客服问答、运营内搜等场景，用户提出一个短问题后，系统基于检索片段生成简短答案，并必须给出引用来源。核心目标不是开放式聊天，而是降低幻觉、提升引用可追溯性和短答案可用性。
|
||
|
||
简历标题建议使用：
|
||
|
||
> SearchShortQA：面向企业知识库的搜索增强短问答对齐系统
|
||
|
||
模型表述建议使用：
|
||
|
||
> 基于 LLaMA/Qwen-style Decoder-only Transformer 架构，构建 100M/300M 两档小参数搜索增强短问答模型；2B+ 模型仅作为 Qwen/LLaMA 开源模型 LoRA/QLoRA 对照，不表述为从零训练。
|
||
|
||
## 数据来源
|
||
|
||
推荐用公开 QA/RAG 数据做底座，再转换成企业知识库样式数据：
|
||
|
||
| 数据源 | 用途 | 备注 |
|
||
| --- | --- | --- |
|
||
| DuReader / DuReader Robust / DuReader Retrieval | 中文主数据 | 真实搜索问题、证据文档、人工答案，最适合中文搜索增强问答 |
|
||
| Natural Questions | 英文开域 QA | Google 搜索问题 + Wikipedia 答案 |
|
||
| MS MARCO | 英文检索/QA | Bing 搜索问题、passage ranking、QA |
|
||
| KILT | provenance 评测 | 强调答案来源追溯，适合引用合法率指标 |
|
||
|
||
落地时保留 `manifest.json`，记录原始数据路径、构造参数、切分比例和随机种子。简历里只写“基于公开 QA/RAG 数据构造”，不要包装成真实企业私有数据。
|
||
|
||
## 数据格式
|
||
|
||
原始 JSONL 推荐字段：
|
||
|
||
```json
|
||
{
|
||
"id": "case_001",
|
||
"question": "退款多久到账？",
|
||
"answer": "原路退回通常需要1-3个工作日",
|
||
"gold_doc_ids": ["D2"],
|
||
"contexts": [
|
||
{"id": "D1", "title": "会员规则", "text": "..."},
|
||
{"id": "D2", "title": "退款规则", "text": "退款审核通过后，原路退回通常需要1-3个工作日。"}
|
||
]
|
||
}
|
||
```
|
||
|
||
兼容字段别名：
|
||
|
||
- 问题：`question` / `query` / `prompt`
|
||
- 答案：`answer` / `answers`
|
||
- 检索片段：`contexts` / `documents` / `snippets` / `search_results`
|
||
- 引用标签：`gold_doc_ids` / `citations` / `supporting_doc_ids`
|
||
|
||
## 构造 SearchShortQA 数据
|
||
|
||
把本地下载好的 DuReader/NQ/MS MARCO/KILT 风格 JSONL 转成统一格式：
|
||
|
||
```bash
|
||
python scripts/build_search_shortqa_dataset.py \
|
||
--inputs dataset/raw/dureader_train.jsonl dataset/raw/msmarco_train.jsonl \
|
||
--output_dir dataset/search_shortqa \
|
||
--train_size 30000 \
|
||
--dev_size 1000 \
|
||
--test_size 1000 \
|
||
--max_contexts 6 \
|
||
--num_negatives 3 \
|
||
--negative_mining bm25
|
||
```
|
||
|
||
如果服务器已安装 `sentence_transformers`，可以切到 embedding hard negative：
|
||
|
||
```bash
|
||
python scripts/build_search_shortqa_dataset.py \
|
||
--inputs dataset/raw/dureader_train.jsonl dataset/raw/msmarco_train.jsonl \
|
||
--output_dir dataset/search_shortqa \
|
||
--train_size 30000 \
|
||
--dev_size 1000 \
|
||
--test_size 1000 \
|
||
--max_contexts 6 \
|
||
--num_negatives 3 \
|
||
--negative_mining embedding \
|
||
--embedding_model BAAI/bge-small-zh-v1.5
|
||
```
|
||
|
||
构造逻辑：
|
||
|
||
- 正样本：原始证据文档或包含答案的文档作为 `gold_doc_ids`。
|
||
- 负样本：默认从全局语料用 BM25-ish 词面相似度采样 hard negatives，也可使用 embedding 相似度采样。
|
||
- 截断策略：优先保留 `gold_doc_ids` 对应证据片段，再补充 hard negatives，避免标签引用的文档被 `max_contexts` 截掉。
|
||
- 输出文件：
|
||
- `dataset/search_shortqa/search_shortqa_train.jsonl`
|
||
- `dataset/search_shortqa/search_shortqa_dev.jsonl`
|
||
- `dataset/search_shortqa/search_shortqa_test.jsonl`
|
||
- `dataset/search_shortqa/manifest.json`
|
||
|
||
构建后先校验字段和引用合法性：
|
||
|
||
```bash
|
||
python scripts/validate_search_shortqa_dataset.py \
|
||
--data dataset/search_shortqa/search_shortqa_train.jsonl
|
||
|
||
python scripts/validate_search_shortqa_dataset.py \
|
||
--data dataset/search_shortqa/search_shortqa_dev.jsonl
|
||
```
|
||
|
||
推荐规模：
|
||
|
||
| Split | 推荐规模 |
|
||
| --- | ---: |
|
||
| SFT train | 30k-100k |
|
||
| GRPO train | 3k-10k prompts |
|
||
| Dev/Test | 1k-3k each |
|
||
| Bad cases | 每轮 100-300 |
|
||
|
||
## SFT 冷启动
|
||
|
||
先把业务数据转成 SFT 对话格式：
|
||
|
||
```bash
|
||
python scripts/build_search_shortqa_sft.py \
|
||
--input dataset/search_shortqa/search_shortqa_train.jsonl \
|
||
--output dataset/search_shortqa/search_shortqa_sft.jsonl \
|
||
--max_contexts 6
|
||
```
|
||
|
||
然后执行 SFT：
|
||
|
||
```bash
|
||
cd trainer
|
||
deepspeed --num_gpus=7 train_full_sft_deepspeed.py \
|
||
--model_profile SearchLM-300M \
|
||
--data_path ../dataset/search_shortqa/search_shortqa_sft.jsonl \
|
||
--save_weight full_sft_search \
|
||
--from_weight full_sft \
|
||
--dtype float16
|
||
```
|
||
|
||
## GRPO 强化训练
|
||
|
||
使用业务奖励函数训练短答案、引用和格式：
|
||
|
||
```bash
|
||
cd trainer
|
||
deepspeed --num_gpus=7 train_search_grpo_deepspeed.py \
|
||
--model_profile SearchLM-300M \
|
||
--data_path ../dataset/search_shortqa/search_shortqa_train.jsonl \
|
||
--from_weight full_sft_search \
|
||
--save_weight search_grpo \
|
||
--num_generations 6 \
|
||
--batch_size 2 \
|
||
--dtype float16
|
||
```
|
||
|
||
奖励函数位于 `trainer/search_shortqa_reward.py`，包含：
|
||
|
||
- `answer_f1`：预测答案与标准答案 token F1
|
||
- `exact_match`：短答案精确匹配
|
||
- `citation_score`：引用 precision / recall / valid rate
|
||
- `format_score`：是否输出合法 JSON
|
||
- `brevity_score`：短答案长度约束
|
||
- `repetition_penalty`：重复惩罚
|
||
|
||
## 离线评测
|
||
|
||
先生成预测文件：
|
||
|
||
```bash
|
||
python scripts/predict_search_shortqa.py \
|
||
--model_profile SearchLM-300M \
|
||
--data dataset/search_shortqa/search_shortqa_dev.jsonl \
|
||
--output outputs/search_shortqa_dev_pred.jsonl \
|
||
--save_dir out \
|
||
--weight search_grpo \
|
||
--load_from model \
|
||
--dtype float16
|
||
```
|
||
|
||
预测文件每行至少包含 `prediction` / `response` / `output` 之一：
|
||
|
||
```json
|
||
{"id": "case_001", "prediction": "{\"answer\":\"原路退回通常需要1-3个工作日\", \"citations\":[\"D2\"]}"}
|
||
```
|
||
|
||
评测：
|
||
|
||
```bash
|
||
python scripts/eval_search_shortqa.py \
|
||
--data dataset/search_shortqa/search_shortqa_dev.jsonl \
|
||
--pred outputs/search_shortqa_dev_pred.jsonl \
|
||
--by id \
|
||
--experiment searchlm300_grpo_search \
|
||
--json_output reports/searchlm300_grpo_search_eval.json \
|
||
--bad_cases_output reports/searchlm300_grpo_search_bad_cases.jsonl
|
||
```
|
||
|
||
输出指标可直接写入实验报告：
|
||
|
||
- `exact_match`
|
||
- `answer_f1`
|
||
- `citation_precision`
|
||
- `citation_recall`
|
||
- `citation_valid`
|
||
- `format_score`
|
||
- `reward_total`
|
||
|
||
## 交互 Demo
|
||
|
||
启动 SearchShortQA 业务演示页：
|
||
|
||
```bash
|
||
streamlit run scripts/search_shortqa_demo.py -- \
|
||
--data dataset/examples/search_shortqa_raw_demo.jsonl
|
||
```
|
||
|
||
Demo 支持编辑问题、检索片段和模型 JSON 输出，实时展示 reward、F1、引用合法率、格式正确率、奖励分解和引用检查表。它适合面试时展示“模型输出是否基于证据、引用是否合法、bad case 如何定位”。
|
||
|
||
## 模型规模建议
|
||
|
||
不要把主线写成“从零训练 2B”。7×RTX3090 更适合：
|
||
|
||
| 模型 | 用途 |
|
||
| --- | --- |
|
||
| SearchLM-100M | debug、结构消融、快速验证 |
|
||
| SearchLM-300M | 简历主模型，SFT + GRPO |
|
||
| SearchLM-500M | 可选大一档实验 |
|
||
| Qwen2.5-1.5B/3B LoRA | 大模型对照，不写成自研底座 |
|
||
|
||
查看推荐配置和估算参数：
|
||
|
||
```bash
|
||
python scripts/estimate_searchlm_scale.py
|
||
python scripts/estimate_searchlm_scale.py --json
|
||
```
|
||
|
||
profile 配置位于 `configs/searchlm_profiles.json`。训练入口支持：
|
||
|
||
```bash
|
||
--model_profile SearchLM-100M
|
||
--model_profile SearchLM-300M
|
||
--model_profile SearchLM-500M
|
||
```
|
||
|
||
DeepSpeed 服务器环境安装：
|
||
|
||
```bash
|
||
pip install deepspeed
|
||
```
|
||
|
||
默认使用 `trainer/ds_config_zero2.json`，采用 ZeRO-2 切分 optimizer state 和 gradient，适合 7×RTX3090 的 100M/300M/500M 训练链路。完整 DeepSpeed checkpoint 保存在 `../checkpoints/`，同时 rank0 会导出一份轻量 `.pth` 权重到 `../out/` 用于评测和推理。
|
||
|
||
## 断点续训
|
||
|
||
DeepSpeed 训练脚本支持 `--from_resume 1`。恢复的是 `../checkpoints/` 下的完整训练状态，包括模型、optimizer、scheduler、DeepSpeed engine 和已训练到的 `epoch/step`；`../out/` 下的 `.pth` 只是轻量推理权重，不能单独用于连续训练。
|
||
|
||
第一次训练建议降低 `save_interval`，减少在线 GPU 中断后的损失：
|
||
|
||
```bash
|
||
cd trainer
|
||
|
||
deepspeed --num_gpus=7 train_pretrain_deepspeed.py \
|
||
--model_profile SearchLM-300M \
|
||
--data_path ../dataset/pretrain_t2t_mini.jsonl \
|
||
--save_weight pretrain_searchlm \
|
||
--dtype float16 \
|
||
--save_interval 200
|
||
```
|
||
|
||
中断后用同样的核心参数恢复：
|
||
|
||
```bash
|
||
cd trainer
|
||
|
||
deepspeed --num_gpus=7 train_pretrain_deepspeed.py \
|
||
--model_profile SearchLM-300M \
|
||
--data_path ../dataset/pretrain_t2t_mini.jsonl \
|
||
--save_weight pretrain_searchlm \
|
||
--dtype float16 \
|
||
--save_interval 200 \
|
||
--from_resume 1
|
||
```
|
||
|
||
恢复时需要保持以下参数一致：`--model_profile`、模型结构相关参数、`--save_weight`、`--data_path`、`--max_seq_len`、GPU 数量、`--batch_size` 和 `--accumulation_steps`。`--epochs` 表示总目标 epoch 数，不是剩余 epoch 数。
|
||
|
||
恢复成功时日志会出现：
|
||
|
||
```text
|
||
DeepSpeed checkpoint loaded: ...
|
||
Epoch [x/y]: 跳过前N个step，从step N+1开始
|
||
```
|
||
|
||
为了避免中断时覆盖最后一个可用 checkpoint，DeepSpeed checkpoint 会保存成 step 级目录，例如 `search_grpo_1024_mla_epoch0_step50`，并用 `search_grpo_1024_mla_latest` 指针记录最后完整保存的 checkpoint。
|
||
|
||
推荐完整训练顺序：
|
||
|
||
```bash
|
||
cd trainer
|
||
|
||
# 1. 通用预训练：使用通用文本，而不是 SearchShortQA
|
||
deepspeed --num_gpus=7 train_pretrain_deepspeed.py \
|
||
--model_profile SearchLM-300M \
|
||
--data_path ../dataset/pretrain_t2t_mini.jsonl \
|
||
--save_weight pretrain_searchlm \
|
||
--dtype float16
|
||
|
||
# 2. 通用 SFT：保留指令跟随能力
|
||
deepspeed --num_gpus=7 train_full_sft_deepspeed.py \
|
||
--model_profile SearchLM-300M \
|
||
--data_path ../dataset/sft_t2t_mini.jsonl \
|
||
--from_weight pretrain_searchlm \
|
||
--save_weight full_sft_searchlm \
|
||
--dtype float16
|
||
|
||
# 3. SearchShortQA SFT：学习检索片段、短答案和引用格式
|
||
deepspeed --num_gpus=7 train_full_sft_deepspeed.py \
|
||
--model_profile SearchLM-300M \
|
||
--data_path ../dataset/search_shortqa/search_shortqa_sft.jsonl \
|
||
--from_weight full_sft_searchlm \
|
||
--save_weight full_sft_search \
|
||
--dtype float16
|
||
|
||
# 4. SearchShortQA GRPO：用可验证奖励优化效果
|
||
deepspeed --num_gpus=7 train_search_grpo_deepspeed.py \
|
||
--model_profile SearchLM-300M \
|
||
--data_path ../dataset/search_shortqa/search_shortqa_train.jsonl \
|
||
--from_weight full_sft_search \
|
||
--save_weight search_grpo \
|
||
--num_generations 6 \
|
||
--dtype float16
|
||
```
|
||
|
||
2B+ 全参数 DDP 不推荐作为主线，因为 24GB 3090 上 actor/ref/optimizer/activation 会让 PPO/GRPO 显存压力过高；如果要体现 2B+，用 Qwen/LLaMA 开源权重做 LoRA/QLoRA 对照即可。
|
||
|
||
已有 64M 权重的处理建议：
|
||
|
||
- 如果只是把训练后端从 DDP 换成 DeepSpeed，模型结构、tokenizer、hidden/layer/head 配置都不变，则不需要重新预训练，可以继续做通用 SFT、SearchShortQA SFT 和 GRPO。
|
||
- 如果改成 SearchLM-300M/500M，或者切换到 MLA/MoE、改变 `hidden_size`、层数、head 数、词表/tokenizer，则旧 64M 不能作为同构 checkpoint 直接复用，需要重新预训练或只把 64M 当 baseline。
|
||
- 简历主线建议保留 64M 作为 baseline，同时至少补一版 SearchLM-100M 或 SearchLM-300M 的 DeepSpeed 预训练日志；这样可以写“64M baseline -> 300M 主模型”的规模消融和工程优化。
|
||
|
||
## 简历表达
|
||
|
||
可以包装为：
|
||
|
||
> 构建 SearchShortQA 搜索增强短问答数据集，基于 DuReader/NQ/MS MARCO 等公开 QA/RAG 数据统一转换为“问题-检索片段-短答案-引用来源”格式，并通过 BM25-ish hard negative 采样构造干扰片段；基于 DeepSpeed ZeRO-2 在 7×RTX3090 上训练 300M 级 LLaMA/Qwen-style Decoder-only 模型，完成 SFT 冷启动与 GRPO 强化对齐，在固定测试集上评估 EM、F1、引用合法率、JSON 格式正确率和推理吞吐。
|
||
|
||
如果同时使用 MLA：
|
||
|
||
> 在 SearchShortQA 场景下对比 GQA/MQA/MLA 注意力结构，MLA 将 KV Cache 理论占用由 768 降至 176 floats/token/layer，降低 77.1%，压缩约 4.36x，并评估其对检索长上下文问答的显存、吞吐和答案质量影响。
|