12 KiB

Raw Blame History

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 推荐字段：

{
  "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 转成统一格式：

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：

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

构建后先校验字段和引用合法性：

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 对话格式：

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：

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 强化训练

使用业务奖励函数训练短答案、引用和格式：

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：重复惩罚

离线评测

先生成预测文件：

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 之一：

{"id": "case_001", "prediction": "{\"answer\":\"原路退回通常需要1-3个工作日\", \"citations\":[\"D2\"]}"}

评测：

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 业务演示页：

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	大模型对照，不写成自研底座

查看推荐配置和估算参数：

python scripts/estimate_searchlm_scale.py
python scripts/estimate_searchlm_scale.py --json

profile 配置位于 configs/searchlm_profiles.json。训练入口支持：

--model_profile SearchLM-100M
--model_profile SearchLM-300M
--model_profile SearchLM-500M

DeepSpeed 服务器环境安装：

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 中断后的损失：

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

中断后用同样的核心参数恢复：

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 数。

恢复成功时日志会出现：

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。

推荐完整训练顺序：

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，并评估其对检索长上下文问答的显存、吞吐和答案质量影响。

12 KiB Raw Blame History Unescape Escape