训练一个 4B 模型以生成 81% 更快的 PostgreSQL 查询计划

TL;DR

一个 40 亿参数的开源权重模型，在从 GPT‑6 Astra 进行离策略蒸馏并经过多轮强化学习后，能够生成 pg_hint_plan 提示，使 PostgreSQL 在平均情况下将 Join Order Benchmark 的执行速度提升 81%，总工作负载延迟降低 44.7%。

为什么查询优化器仍然落后

Leis 等人（2015, 2025）表明，PostgreSQL 的优化器在性能方面仍有巨大提升空间，尤其是在连接顺序方面，而连接顺序问题是 NP 难题。
优化器依赖粗略的统计信息和均匀分布假设；一个错误的基数估计可能导致计划严重次优。
验证计划质量非常简单——执行时间是一个单一、可观测的指标——这使得该问题非常适合强化学习。

问题定义

目标： 训练一个小型模型，输出 pg_hint_plan 注释，引导 PostgreSQL 生成比其默认基于成本的规划器更快的执行计划。

关键洞察： 对于重复运行相同查询的分析型工作负载，更优计划的摊销收益远超提示生成的一次性成本。

实验平台

组件	详情
数据库	PostgreSQL 15，使用 IMDb 数据集的一个子集（磁盘占用 8.5 GB）。
硬件	开发机器 "FLOPper" – 2 × RTX 3090，16 个 CPU 核心，64 GB 内存。
训练 GPU	从 Lambda 租用的 2 × H100 SXM（每个 80 GB VRAM），约 95 小时。
基准测试	Join Order Benchmark（JOB – 113 个查询，33 种连接图拓扑）。
辅助基准测试	Cardinality Estimation Benchmark（CEB – 约 13.6k 个查询），用于训练数据。

减少测量噪声

校准装置： 四个 Docker 化的 PostgreSQL 容器，每个绑定到 4 个 CPU 核心和 8 GB 内存，从共享队列中拉取查询。
预热策略： 运行查询直到 shared_buffers 的命中块和读取块计数器在连续两次运行中稳定在 2% 以内。
噪声指标： 模拟的 "无操作" 奖励（候选 = 默认）在 shared_buffers = 128 MB 时显示 5% 的平均误差率，最坏情况（p90）误差为 13% – 20%。
调优： 将 shared_buffers 提高到 2 GB 后，大部分噪声被消除（平均误差 ≈ 1.5%，p90 ≈ 0%），并将 JOB 总运行时间从 95 秒降至 60 秒。

模型与工具链

基础模型： empero-ai/Qwen3.8-4B-Distill（从 2.4T 教师模型蒸馏的 Qwen 3.8 4B）。
代理工具链（qo-agent）： 提供六个工具（inspect_relation、get_column_stats、get_plan、evaluate_candidate、keep_default、finish）。
交互格式： 模型输出 PlanAction JSON；工具链将其转换为 pg_hint_plan 注释并测量执行时间。

离策略蒸馏（监督微调）

教师轨迹： 在随机 CEB 查询上进行 120 次 GPT‑6 Astra 滚动（每次 5 个候选），附带推理摘要。
渲染与损失掩码： 将 OpenAI 响应 JSON 转换为 Qwen 令牌流；掩码系统/用户提示和工具输出令牌。
LoRA 适配器： 42.5 MB（2120 万可训练参数），冻结 4B 基础模型权重，使单个 RTX 3090 上即可训练。
训练调度： 在 100 条 Astra 轨迹上训练 1 个周期 → 轻微提升（几何平均 0.72×）。2 个周期 → 几何平均加速 1.08×；3 个周期出现过拟合并退化。
数据扩展： 增加 300 条 Astra 轨迹（过滤掉简单的 "保留默认" 运行）。再训练 2 个周期后，几何平均提升至 1.10×，总工作负载加速达 1.05×。

代理强化学习

奖励设计：
- 加速比 = 中位数(默认) / 中位数(候选)
- 将比率裁剪至 [0.1, 10] 并应用 0.05 的软阈值以忽略噪声。
- 惩罚无效计划（‑0.1）、重复默认计划（‑0.02），以及无有效候选的轨迹（‑0.1）。
优势估计器： 自定义的 "锚定" GRPO，减去组平均值并应用每轮次符号，确保只有真正更快的计划才能获得正优势。
训练超参数： 学习率 = 1e‑5，批量 = 16，每个查询 8 次滚动，每阶段 600 次优化器更新。
并发技巧： 仅在测量阶段租用 PostgreSQL 容器，允许 20 个并发滚动，保持 vLLM 推理完全饱和。

结果

阶段	有效候选查询	评分任务	几何平均加速比	总工作负载加速比	胜利（≥ 5% 更快）	回退（≥ 5% 更慢）
未训练的 4B	14/113	15/113	0.85×	0.85×	3	1
SFT 后（1 个周期）	48/113	44/113	0.72×	0.76×	5	16
SFT 后（2 个周期，额外 300 条轨迹）	77/113	101/113	1.10×	1.05×	29	13
RL 600 次更新	99/113	113/113	1.35×	1.16×	34	0
RL 1200 次更新	101/113	112/113	1.41×	1.29×	38	2
RL 1200 次更新 + 每查询 3 次滚动（最佳 15 选 1）	113/113	339/339	1.81×	1.81×	68	0

最终检查点在每查询最多采样 15 个计划中选择最佳候选时，实现了 1.81× 的几何平均加速比（总延迟降低 44.7%）。

模型实际学到的内容

高频操作： Leading 树提示（917 次使用）、扫描提示（1141 次）和 Parallel 提示（572 次）。
连接方法偏好： 即使哈希连接是默认设置，也常强制使用嵌套循环连接。
扫描偏好： 优先选择索引扫描而非位图或顺序扫描。
配置调整： 频繁设置 enable_sort=off 和 random_page_cost=1.1，表明模型利用了成本模型的敏感性。
成功模式： 通过 Leading 重新排序连接、添加一个更选择性的扫描，以及启用并行性，共同构成了大部分加速效果。

成本分解

项目	成本
Lambda 2 × H100 租用（约 95 小时）	~ $800
OpenAI API 用于 Astra 轨迹（约 40 万 token）	~ $400
FLOPper 电力消耗（持续 GPU 使用）	~ $9/天（可忽略）
总计	≈ $1,200

社区反应（精选 HN 评论）

"在 8 GB 数据集上实现 81% 更快的查询计划，而该数据集完全可装入内存……" – refibrillator（对过拟合和可扩展性的担忧）。 "前沿智能非常强大；我从 Astra 轨迹中进行的蒸馏已充分证明大型模型不会消失" – devsda（指出开源权重蒸馏的广泛相关性）。 "要超过 Postgres 3 倍以上并不难；你不需要模型" – huahaiy（指出手工调优提示也能带来巨大收益）。 "模型经常使用 enable_sort=off 和 random_page_cost=1.1 – 可能解释了性能提升" – zacmps（强调可能的混淆配置变更）。 "我希望能有一个产品能自动为重复的分析查询生成提示" – ashley95（建议一个实用的 SaaS 方向）。

启示

LLM 可以学会生成有效的优化器提示。 即使是 4B 模型，在从更大教师模型蒸馏并结合真实执行反馈进行强化学习后，也能在真实的分析基准测试中超越 PostgreSQL 的原生规划器。
噪声感知测量至关重要。 适当的缓存大小（shared_buffers = 2 GB）和三对中位数的中位数评估将虚假奖励信号从 <5% 降低至接近零。
离策略蒸馏加速了 SFT。 数百条教师轨迹提供了足够的令牌级监督，使模型学会工具链语言和基本提示语法。
奖励设计比算法新颖性更重要。 早期的 GRPO 变体强化了重复默认计划；自定义锚定 GRPO 和裁剪加速奖励对稳定强化学习进展至关重要。
小模型在特定领域任务中具有成本效益。 训练成本 ≈ $1.2k，单个 RTX 3090 上推理成本可忽略，但模型在可重复工作负载上实现了 >80% 的加速。

未来工作

结构化提示扫荡（如 Bao）与 LLM 生成提示的对比。
在策略蒸馏，用于比较采样效率。
轨迹反演，以获得更丰富的推理令牌监督。
扩展到更大数据集（TB 级）和混合 OLTP/OLAP 工作负载。
自动化提示缓存服务，持续基于生产查询日志重新训练。

代码与引用

所有代码在 https://github.com/polyphilz/qorl 开源。

@article{bansal2026qorl,
  title = {Training a 4B model to produce 81% faster query plans than Postgres},
  author = {Bansal, Rohan},
  journal = {rohanbansal.com},
  year = {2026},
  month = {September},
  url = {https://rohanbansal.com/qorl}
}