訓練 4B 模型以產生快 81% 的 PostgreSQL 查詢計畫

摘要

一個具有 40 億參數的開放權重模型，在經過來自 GPT-6 Astra 的離策略蒸餾（off-policy distillation）以及數個輪次的強化學習後，能夠生成 pg_hint_plan 提示，使 PostgreSQL 在 Join Order Benchmark 上的平均執行速度提升 81%，並將總工作負載延遲降低 44.7%。

為什麼查詢最佳化器仍然落後

  • Leis 等人（2015, 2025）指出，PostgreSQL 的最佳化器在效能上仍有大量提升空間，尤其是在連線排序（join ordering）方面，這屬於 NP-hard 問題。
  • 最佳化器依賴粗略的統計資訊和均勻分佈假設；單一的基數估計錯誤就可能導致計畫嚴重次優。
  • 驗證計畫品質非常簡單——執行時間是一個單一且可觀測的指標——這使得該問題非常適合使用強化學習。

問題定義

目標： 訓練一個小型模型，使其輸出 pg_hint_plan 註解，引導 PostgreSQL 產生比預設成本基礎規劃器更快的計畫。

關鍵洞察： 對於重複執行相同查詢的分析型工作負載，更優計畫帶來的攤提效益大於生成提示的一次性成本。

實驗平台

組件 詳情
資料庫 運行在 IMDb 數據集切片（磁碟佔用 8.5 GB）上的 PostgreSQL 15。
硬體 開發機器 "FLOPper" – 2 × RTX 3090, 16 個 CPU 核心, 64 GB RAM。
訓練 GPU 從 Lambda 租用的 2 × H100 SXM（每張 80 GB VRAM），時長約 95 小時。
基準測試 Join Order Benchmark (JOB – 113 個查詢, 33 種連線圖拓撲)。
輔助基準測試 Cardinality Estimation Benchmark (CEB – 約 13.6 k 個查詢)，用於訓練數據。

減少測量噪聲

  • 校準裝置： 四個 Docker 化的 PostgreSQL 容器，每個固定分配 4 個 CPU 核心和 8 GB RAM，從共享佇列中獲取查詢。
  • 預熱策略： 執行查詢直到 shared_buffers 的命中區塊（hit-block）和讀取區塊（read-block）計數器在連續兩次運行中穩定在 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.4 T 教師模型蒸餾而成的 Qwen 3.8 4 B）。
  • 智能體框架 (qo-agent)： 提供六個工具（inspect_relation, get_column_stats, get_plan, evaluate_candidate, keep_default, finish）。
  • 互動格式： 模型輸出 PlanAction JSON；框架將其轉換為 pg_hint_plan 註解並測量執行時間。

離策略蒸餾（監督式微調）

  1. 教師軌跡： 在隨機 CEB 查詢上進行 120 次 GPT-6 Astra 滾動（每次 5 個候選），並附帶推理摘要。
  2. 渲染與損失遮罩： 將 OpenAI 回應 JSON 轉換為 Qwen 令牌流；遮罩掉系統/用戶提示和工具輸出令牌。
  3. LoRA 適配器： 42.5 MB（21.2 M 可訓練參數），凍結 4 B 基礎模型權重，使其能在單張 RTX 3090 上進行訓練。
  4. 訓練排程： 在 100 條 Astra 軌跡上訓練一個輪次 → 溫和改進（幾何平均 0.72×）。兩個輪次 → 1.08× 幾何平均加速；三個輪次導致過擬合並退步。
  5. 數據擴展： 新增 300 條 Astra 軌跡（過濾掉簡單的「保持預設」運行）。再進行兩個輪次，將幾何平均提升至 1.10×，總工作負載加速至 1.05×。

智能體強化學習

  • 獎勵設計：
    • 加速比 = median(default) / median(candidate)。
    • 將比率裁剪至 [0.1, 10]，並應用 0.05 的軟閾值以忽略噪聲。
    • 對無效計畫（-0.1）、重複預設計畫（-0.02）以及沒有有效候選的軌跡（-0.1）進行懲罰。
  • 優勢估計器： 自定義的「錨定」GRPO，減去組平均值並應用每次滾動符號，確保只有真正更快的計畫獲得正優勢。
  • 訓練超參數： LR = 1e-5, batch = 16, 每個查詢 8 次滾動, 每個階段 600 次最佳化器更新。
  • 並發技巧： 僅在測量階段租用 PostgreSQL 容器，允許 20 個並發滾動，並保持 vLLM 推理完全飽和。

結果

階段 有效候選查詢 計分任務 幾何平均加速 總工作負載加速 勝利（≥ 5% 更快） 退步（≥ 5% 更慢）
未訓練 4 B 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 1 200 次更新 101/113 112/113 1.41× 1.29× 38 2
RL 1 200 + 每查詢 3 次滾動（15 選 1） 113/113 339/339 1.81× 1.81× 68 0

*最終檢查點在每個查詢中從最多 15 個採樣計畫中選擇最佳候選時，實現了 1.81× 的幾何平均加速（總延遲降低 44.7%）。

模型實際學到了什麼

  • 常見動作： Leading 樹提示（使用 917 次），掃描提示（1 141 次），以及 Parallel 提示（572 次）。
  • 連線方法偏好： 即使預設是哈希連線，也經常強制使用巢狀循環連線（Nested-loop joins）。
  • 掃描偏好： 優先使用索引掃描，而非位圖或順序掃描。
  • 配置調整： 經常設置 enable_sort=off 和 random_page_cost=1.1，表明模型利用了成本模型的敏感性。
  • 成功模式： 通過 Leading 重新排序連線、添加一個更具選擇性的掃描，以及啟用並行性，共同佔據了大部分加速來源。

成本細分

項目 成本
Lambda 2 × H100 租賃（≈ 95 小時） ~ $800
用於 Astra 軌跡的 OpenAI API（≈ 400 k 令牌） ~ $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 方向）。

重點總結

  1. LLM 可以學會輸出有效的最佳化器提示。 即使是一個 4 B 模型，一旦從更大的教師模型蒸餾並通過真實執行反饋進行強化，就能在現實的分析基準測試中優於 PostgreSQL 的原生規劃器。
  2. 感知噪聲的測量至關重要。 適當的緩存大小（shared_buffers = 2 GB）和三對中位數的中位數評估將虛假獎勵信號從 < 5% 降低到接近零。
  3. 離策略蒸餾加速了 SFT。 幾百條教師軌跡提供了足夠的令牌級監督，以教會模型框架語言和基本提示語法。
  4. 獎勵設計比算法新穎性更重要。 早期的 GRPO 變體強化了重複預設計畫；自定義的錨定 GRPO 和裁剪加速獎勵對於穩定的 RL 進展至關重要。
  5. 小模型在特定領域任務中具有成本效益。 訓練成本 ≈ $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}
}

Sources

相關