用 4B 模型生成比 Postgres 更快的查询计划
Leis 等人在2015年就提出了这个问题,十年后他们再次发起了同样的研究。
尽管自首次探索以来的十年间相关研究成果颇丰,但他们发现查询优化器的表现仍有很大改进空间。
我刚得知这个结论时十分惊讶:Postgres数据库理应完全了解自身存储的数据,优化查询能有多难呢?
事实证明,难度超乎想象。查询优化器的核心任务之一——连接顺序优化,已被证实属于NP难问题。
既然优化器本身难以改进,那验证优化器选择的查询计划优劣难度相对较低。简单来说,好的优化器生成的计划执行速度快,差的则相反。语言模型尤其擅长学习输出可验证的任务,而查询优化恰好有明确的优化目标——执行时间,因此可以通过强化学习引导模型生成更快的查询计划。
下文将详细介绍我开展的一项实验,旨在探索以下问题:一个小型开源权重模型能否通过监督微调(SFT)和智能体强化学习(RL)进行后训练,生成比Postgres默认计划更优的查询计划?
答案是肯定的,主要成果包括:
- 针对113个多连接查询,将延迟降低44.7%;初始时该4B模型无法为其中99个查询生成有效计划
- 搭建Postgres测试平台,减少并发容器间Linux页缓存竞争带来的噪声
- 设计定制化GRPO变体,在高噪声环境中为RL轨迹评分
- 跨两台机器部署RL流程:租用的2x H100节点运行vLLM和训练器,本地桌面运行4个Postgres容器
- 基于近500条GPT-6 Astra智能体轨迹进行离线策略蒸馏
接下来从头说起。
查询优化器的工作原理
我们先来看IMDb数据集的部分表结构:
-- IMDb影视条目(电影、剧集、单集等)[约100万行]
title (
id integer PRIMARY KEY,
title text,
production_year integer,
kind_id integer -- 外键 -> kind_type
)
-- 影视与制作公司关联表[约200万行]
movie_companies (
id integer PRIMARY KEY,
movie_id integer, -- 外键 -> title.id
company_id integer, -- 外键 -> company_name.id
company_type_id integer, -- 外键 -> company_type.id
note text
)
-- 制作公司信息表[约10万行]
company_name (
id integer PRIMARY KEY,
name text,
country_code text -- '[us]', '[jp]'等
)
-- 公司角色 lookup表[4行]
company_type (
id integer PRIMARY KEY,
kind text -- 'production companies', 'distributors'等
)
-- 影视类型 lookup表[7行]
kind_type (
id integer PRIMARY KEY,
kind text -- 'movie', 'tv series', 'episode'等
)
假设我们要回答:"2000年代日本哪家公司出品的影视条目最多?"可以编写如下查询:
SELECT cn.name,
COUNT(*) AS titles
FROM title AS t,
movie_companies AS mc,
company_name AS cn
WHERE t.id = mc.movie_id
AND mc.company_id = cn.id
AND cn.country_code = '[jp]'
AND t.production_year BETWEEN 2000 AND 2009
GROUP BY cn.name
ORDER BY titles DESC
LIMIT 10;
执行该查询会返回2000-2009年间出品影视条目数量最多的10家日本公司,按数量从高到低排序。
但Postgres是如何得到这个结果的呢?
Postgres获取数据的路径并非固定,而是取决于查询中的选择性谓词(即WHERE子句中的过滤条件)。
为说明这一点,我们先移除日本公司和时间范围的过滤条件:
SELECT cn.name,
COUNT(*) AS titles
FROM title AS t,
movie_companies AS mc,
company_name AS cn
WHERE t.id = mc.movie_id
AND mc.company_id = cn.id
GROUP BY cn.name
ORDER BY titles DESC
LIMIT 10;
此时mc只能通过mc.company_id = cn.id与cn连接,t只能通过t.id = mc.movie_id与mc连接。
这些约束产生两种若考虑交换律,实际有8种连接树,但此处交换律不影响连接结果的行数,因此不予考虑。有效连接树:
表或查询结果的基数指其包含的行数。假设相关表的基数如下:
结合连接关系,我们得到以下连接结果基数:
无论这三个表以何种顺序连接,第二次连接时都会传入200万行数据。
现在我们重新添加选择性谓词:
- 假设10万家公司中有5%是日本公司
movie_companies基数不变(200万行)- 假设100万条影视条目中有20%制作于2000年代
第一种连接顺序先将200万条movie_companies数据过滤为5%的日本公司条目,假设分布均匀,连接结果约为10万行;再与过滤后的title表连接,仅保留其中20%的2000年代条目。
第二种连接顺序先将200万条movie_companies数据过滤为20%的2000年代条目,同样假设分布均匀,第一次连接结果为40万行,即第二次连接会传入40万行数据。
选择第二种连接顺序的工作量是第一种的4倍。
但问题远不止于此。
组合爆炸
每个连接可选择以下算法之一:
- 哈希连接(Hash join)
- 归并连接(Merge join)
- 嵌套循环连接(Nested-loop join)
现在重新考虑交换律虽然交换律不影响结果行数,但会影响连接算法的性能,因此必须考虑。,会产生4种不同的内外连接方向,共8种可能的组合:
此外,每个表还可以通过不同方式扫描,仅考虑以下4种扫描类型:
- 顺序扫描(Sequential)
- 索引扫描(Index)
- 仅索引扫描(Index-only)
- 位图扫描(Bitmap)
执行该查询共有4608种不同方式这还只是保守估计,实际还可考虑并行执行、哈希聚合或排序聚合等情况。
值得注意的是,Postgres并不会评估所有这些计划,它使用动态规划(对包含12个以上连接的查询则使用遗传算法)来剪枝搜索空间。
更糟糕的是,每增加一个连接,搜索空间就会呈组合式爆炸增长:
SELECT cn.name,
COUNT(*) AS titles
FROM movie_companies AS mc,
company_name AS cn
WHERE mc.company_id = cn.id
AND cn.country_code = '[jp]'
GROUP BY cn.name
ORDER BY titles DESC
LIMIT 10;
SELECT cn.name,
COUNT(*) AS titles
FROM title AS t,
movie_companies AS mc,
company_name AS cn
WHERE t.id = mc.movie_id
AND mc.company_id = cn.id
AND cn.country_code = '[jp]'
AND t.production_year BETWEEN 2000 AND 2009
GROUP BY cn.name
ORDER BY titles DESC
LIMIT 10;
SELECT MIN(t.title) AS movie_title
FROM keyword AS k,
movie_info AS mi,
movie_keyword AS mk,
title AS t
WHERE k.keyword LIKE '%sequel%'
AND mi.info IN ('Bulgaria')
AND t.production_year > 2010
AND t.id = mi.movie_id
AND t.id = mk.movie_id
AND mk.movie_id = mi.movie_id
AND k.id = mk.keyword_id;
SELECT MIN(t.title) AS movie_title
FROM company_name AS cn,
keyword AS k,
movie_companies AS mc,
movie_keyword AS mk,
title AS t
WHERE cn.country_code ='[de]'
AND k.keyword ='character-name-in-title'
AND cn.id = mc.company_id
AND mc.movie_id = t.id
AND t.id = mk.movie_id
AND mk.keyword_id = k.id
AND mc.movie_id = mk.movie_id;
SELECT MIN(lt.link) AS link_type,
MIN(t1.title) AS first_movie,
MIN(t2.title) AS second_movie
FROM keyword AS k,
link_type AS lt,
movie_keyword AS mk,
movie_link AS ml,
title AS t1,
title AS t2
WHERE k.keyword ='10,000-mile-club'
AND mk.keyword_id = k.id
AND t1.id = mk.movie_id
AND ml.movie_id = t1.id
AND ml.linked_movie_id = t2.id
AND lt.id = ml.link_type_id
AND mk.movie_id = t1.id;
SELECT MIN(a1.name) AS writer_pseudo_name,
MIN(t.title) AS movie_title
FROM aka_name AS a1,
cast_info AS ci,
company_name AS cn,
movie_companies AS mc,
name AS n1,
role_type AS rt,
title AS t
WHERE cn.country_code ='[us]'
AND rt.role ='writer'
AND a1.person_id = n1.id
AND n1.id = ci.person_id
AND ci.movie_id = t.id
AND t.id = mc.movie_id
AND mc.company_id = cn.id
AND ci.role_id = rt.id
AND a1.person_id = ci.person_id
AND ci.movie_id = mc.movie_id;
SELECT MIN(an.name) AS cool_actor_pseudonym,
MIN(t.title) AS series_named_after_char
FROM aka_name AS an,
cast_info AS ci,
company_name AS cn,
keyword AS k,
movie_companies AS mc,
movie_keyword AS mk,
name AS n,
title AS t
WHERE cn.country_code ='[us]'
AND k.keyword ='character-name-in-title'
AND an.person_id = n.id
AND n.id = ci.person_id
AND ci.movie_id = t.id
AND t.id = mk.movie_id
AND mk.keyword_id = k.id
AND t.id = mc.movie_id
AND mc.company_id = cn.id
AND an.person_id = ci.person_id
AND ci.movie_id = mc.movie_id
AND ci.movie_id = mk.movie_id
AND mc.movie_id = mk.movie_id;
SELECT MIN(cn.name) AS producing_company,
MIN(miidx.info) AS rating,
MIN(t.title) AS movie
FROM company_name AS cn,
company_type AS ct,
info_type AS it,
info_type AS it2,
kind_type AS kt,
movie_companies AS mc,
movie_info AS mi,
movie_info_idx AS miidx,
title AS t
WHERE cn.country_code ='[us]'
AND ct.kind ='production companies'
AND it.info ='rating'
AND it2.info ='release dates'
AND kt.kind ='movie'
AND mi.movie_id = t.id
AND it2.id = mi.info_type_id
AND kt.id = t.kind_id
AND mc.movie_id = t.id
AND cn.id = mc.company_id
AND ct.id = mc.company_type_id
AND miidx.movie_id = t.id
AND it.id = miidx.info_type_id
AND mi.movie_id = miidx.movie_id
AND mi.movie_id = mc.movie_id
AND miidx.movie_id = mc.movie_id;
SELECT MIN(n.name) AS voicing_actress,
MIN(t.title) AS jap_engl_voiced_movie
FROM aka_name AS an,
char_name AS chn,
cast_info AS ci,
company_name AS cn,
info_type AS it,
movie_companies AS mc,
movie_info AS mi,
name AS n,
role_type AS rt,
title AS t
WHERE ci.note IN ('(voice)',
'(voice: Japanese version)',
'(voice) (uncredited)',
'(voice: English version)')
AND cn.country_code ='[us]'
AND it.info = 'release dates'
AND n.gender ='f'
AND rt.role ='actress'
AND t.production_year > 2000
AND t.id = mi.movie_id
AND t.id = mc.movie_id
AND t.id = ci.movie_id
AND mc.movie_id = ci.movie_id
AND mc.movie_id = mi.movie_id
AND mi.movie_id = ci.movie_id
AND cn.id = mc.company_id
AND it.id = mi.info_type_id
AND n.id = ci.person_id
AND rt.id = ci.role_id
AND n.id = an.person_id
AND ci.person_id = an.person_id
AND chn.id = ci.person_role_id;
SELECT MIN(kt.kind) AS movie_kind,
MIN(t.title) AS complete_us_internet_movie
FROM complete_cast AS cc,
comp_cast_type AS cct1,
company_name AS cn,
company_type AS ct,
info_type AS it1,
keyword AS k,
kind_type AS kt,
movie_companies AS mc,
movie_info AS mi,
movie_keyword AS mk,
title AS t
WHERE cct1.kind = 'complete+verified'
AND cn.country_code = '[us]'
AND it1.info = 'release dates'
AND kt.kind IN ('movie')
AND mi.note LIKE '%internet%'
AND mi.info IS NOT NULL
AND (mi.info LIKE 'USA:% 199%'
OR mi.info LIKE 'USA:% 200%')
AND t.production_year > 2000
AND kt.id = t.kind_id
AND t.id = mi.movie_id
AND t.id = mk.movie_id
AND t.id = mc.movie_id
AND t.id = cc.movie_id
AND mk.movie_id = mi.movie_id
AND mk.movie_id = mc.movie_id
AND mk.movie_id = cc.movie_id
AND mi.movie_id = mc.movie_id
AND mi.movie_id = cc.movie_id
AND mc.movie_id = cc.movie_id
AND k.id = mk.keyword_id
AND it1.id = mi.info_type_id
AND cn.id = mc.company_id
AND ct.id = mc.company_type_id
AND cct1.id = cc.status_id;
SELECT MIN(chn.name) AS character_name,
MIN(mi_idx.info) AS rating,
MIN(t.title) AS complete_hero_movie
FROM complete_cast AS cc,
comp_cast_type AS cct1,
comp_cast_type AS cct2,
char_name AS chn,
cast_info AS ci,
info_type AS it2,
keyword AS k,
kind_type AS kt,
movie_info_idx AS mi_idx,
movie_keyword AS mk,
name AS n,
title AS t
WHERE cct1.kind = 'cast'
AND cct2.kind LIKE '%complete%'
AND chn.name IS NOT NULL
AND (chn.name LIKE '%man%'
OR chn.name LIKE '%Man%')
AND it2.info = 'rating'
AND k.keyword IN ('superhero',
'marvel-comics',
'based-on-comic',
'fight')
AND kt.kind = 'movie'
AND mi_idx.info > '8.0'
AND t.production_year > 2005
AND kt.id = t.kind_id
AND t.id = mk.movie_id
AND t.id = ci.movie_id
AND t.id = cc.movie_id
AND t.id = mi_idx.movie_id
AND mk.movie_id = ci.movie_id
AND mk.movie_id = cc.movie_id
AND mk.movie_id = mi_idx.movie_id
AND ci.movie_id = cc.movie_id
AND ci.movie_id = mi_idx.movie_id
AND cc.movie_id = mi_idx.movie_id
AND chn.id = ci.person_role_id
AND n.id = ci.person_id
AND k.id = mk.keyword_id
AND cct1.id = cc.subject_id
AND cct2.id = cc.status_id
AND it2.id = mi_idx.info_type_id;
估算而非计数
Postgres面临一个困境:理论上它可以通过计算基数来选择能最小化后续连接数据量的计划,但这意味着它需要在查询规划阶段实际执行每个连接并统计结果行数,这完全违背了快速查询优化的初衷。查询优化器追求的并非精确的成本最小化,而是在各类查询中表现足够好。
因此,Postgres使用统计数据来估算基数。规划器会查询pg_statistic表,获取各列的常见值及其频率,以及其余数据的直方图。当涉及连接时,情况会更复杂:Postgres不知道一张表的行如何分布到另一张表中,因此它假设某一值在第一张表中的频率可以直接应用到第二张表,这就是前文提到的均匀分布假设。
均匀分布假设作为启发式方法尚可接受,但一旦失效,后果会很严重。回到之前的连接顺序例子,我们假设5%的公司是日本公司,从而过滤200万条movie_companies条目。但如果这5%的日本公司实际上参与了**50%**的影视制作,那么第一次连接结果会是100万行!此时成本模型会选择第一种连接顺序,但实际上第二种顺序更好,因为它仅向第二次连接传入40万行数据。
早期连接中的一次错误估算会传导至整个连接树,导致所有后续估算失效。
如何引导Postgres
Postgres总是选择成本最低的计划,而我们无法在不修改源代码的情况下改变其成本模型,那该如何引导它选择成本更高但实际更优的计划呢?
答案是pg_hint_plan。
pg_hint_plan是一个设计简洁的第三方扩展:只需在SQL语句上方添加结构化的"提示"注释,就能引导Postgres采用提示中指定的计划。例如:
/*+
HashJoin(a b) -- 指定a和b使用哈希连接
SeqScan(a) -- 指定对表a进行顺序扫描而非索引扫描
*/
EXPLAIN -- EXPLAIN仅输出Postgres将使用的计划,不实际执行查询
SELECT * FROM pgbench_branches b JOIN pgbench_accounts a ON b.bid = a.bid ORDER BY a.aid;
QUERY PLAN
---------------------------------------------------------------------------------
Sort -- 计划的根节点,数据向上流动,因此最后执行:按a.aid排序连接结果
(cost=31465.84..31715.84 rows=100000 width=197) -- Postgres对该节点的估算:启动成本..总成本、预估输出行数、平均行宽(字节)
Sort Key: a.aid
-> Hash Join -- 符合HashJoin(a b)提示要求的连接方式
(cost=1.02..4016.02 rows=100000 width=197)
Hash Cond: (a.bid = b.bid) -- 连接条件,来自ON子句
-> Seq Scan on pgbench_accounts a -- 符合SeqScan(a)提示要求的扫描方式,这是探测端:每行在哈希表中查找匹配项
(cost=0.00..2640.00 rows=100000 width=97)
-> Hash -- 构建端:先读取小表并加载到内存哈希表中,以bid为键
(cost=1.01..1.01 rows=1 width=100)
-> Seq Scan on pgbench_branches b
(cost=0.00..1.01 rows=1 width=100)
(7 rows)
示例来自pg_hint_plan的官方文档。
提示要求对pgbench_accounts和pgbench_branches使用HashJoin,并对pgbench_accounts进行顺序扫描;实际查询计划完全遵循了这些要求。
问题建模
既然我们可以通过pg_hint_plan提示影响Postgres选择不同(甚至更优)的查询计划,那么我们的核心问题就变成:
语言模型能否学习生成能产生更优查询计划的提示?
相关研究
这个问题的研究价值何在?
我的第一个想法是给模型提供查询语句,以及与Postgres规划器完全相同的信息,相当于尝试构建一个更优的基数估算器。但我很快意识到这并非值得探索的方向:我们将与数十年的基数估算研究成果竞争,而且仅推理延迟一项就会远超Postgres超快查询优化器的优势。
第二个想法——也是我认为正确的建模方式——聚焦于一种特定的数据库使用模式:密集型分析负载。如果某个查询使用Postgres默认的非最优计划运行了数千次,那么效率提升的空间就被白白浪费了。此时,模型可以通过训练找到运行该特定查询的更优方式。训练过程可能需要预先执行该查询数十到数百次,但分摊到所有查询运行次数上的成本会大幅降低。
我们的目标并非在单次查询的时间/效率帕累托前沿上击败Postgres,而是在重复运行的查询上超越它。
模型与工具链
我选择从一个小型4B模型开始,因为它最容易在我家用的2x RTX 3090设备(昵称FLOPper)上进行训练和推理。
我启动项目时,Qwen 3.8系列模型已发布,但遗憾的是没有4B版本。不过我发现德国一家小型实验室Empero基于Qwen 3.8 2.4T模型蒸馏出了Qwen 3.8 4B版本,即empero-ai/Qwen3.8-4B-Distill。这个蒸馏模型并非全面优于其基础模型Qwen 3.5 4B:它在MMLU任务上表现更好,但在GSM8K任务上略逊一筹。换句话说,该蒸馏模型在通用知识广度评估中表现更优,但在多步数学推理上稍差。至于哪种模型更适合我们的任务,我并不确定,但最终还是选择了这个蒸馏模型。
确定模型后,我构建了一个轻量级智能体工具链qo-agent,用于协调提示生成。它提供以下6种工具:
inspect_relation——列出表的列(含类型和可空性)、索引定义以及预估行数和字节数get_column_stats——获取Postgres规划器对一个表中1-8列的统计信息get_plan——获取默认计划的估算结果,或已提交候选计划的存储计划evaluate_candidate——验证提议的计划操作,然后执行以获取时间和计划诊断信息keep_default——将Postgres的默认计划作为候选计划返回,并结束搜索finish——接受提交的候选计划ID或默认计划作为输入,结束搜索
为利用结构化输出,我要求智能体生成PlanAction JSON对象。调用evaluate_candidate时,会将PlanAction对象编译为提示,并添加到原始查询的开头。
以下是一个智能体轨迹示例:
智能体→ get_plan("default") // 向工具链请求Postgres自行选择的计划,仅返回规划器的估算结果,不实际执行 工具← 默认计划:(t ⋈ mc) ⋈ cn,哈希连接,各节点预估行数 智能体→ evaluate_candidate({ "leading": { "left": { "left": "cn", "right": "mc" }, "right": "t" }, // 连接树:先连接cn和mc,再与t连接,即前文第一种连接树 "joins": [{ "relations": ["cn", "mc"], "force": "hash" }] // 强制cn ⋈ mc步骤使用哈希连接 }) 工具← c1 // 工具链为该候选计划分配的ID,get_plan和finish通过ID引用候选计划 · 有效 · 全新计划 · 118毫秒 · 0.91× 默认计划 // 执行反馈:候选计划的测量时间与默认计划中位数的初步比值,小于1表示比Postgres慢 · 剩余2次尝试机会 // 每次提交(无论有效与否)都会消耗一次候选尝试机会 智能体→ evaluate_candidate({ "leading": { "left": { "left": "t", "right": "mc" }, "right": "cn" }, // 尝试另一种连接树:先连接t和mc,再与cn连接 "joins": [{ "relations": ["t", "mc"], "force": "nestloop" }], // 第一次连接使用嵌套循环,过滤后的t行逐一探测mc "scans": [{ "relation": "mc", "force": "index" }] // 循环的内表mc使用索引扫描 }) 工具← c2 · 有效 · 全新计划 · 87毫秒 · 1.24× 默认计划 // 比默认计划快:比值大于1表示加速 · 剩余1次尝试机会 智能体→ evaluate_candidate({ "leading": { "left": { "left": "t", "right": "mc" }, "right": "cn" }, "joins": [{ "relations": ["t", "mc"], "force": "hash" }] // 与c2连接树相同,但第一次连接使用哈希连接,验证嵌套循环是否是性能提升的关键 }) 工具← c3 · 有效 · 全新计划 · 100毫秒 · 1.08× 默认计划 · 剩余0次尝试机会 // 候选计划预算耗尽,此后模型只能调用finish工具 智能体→ finish({ "selected_candidate_id": "c2" }) // 通过指定候选ID结束搜索,可选择之前的候选计划 工具← 已完成 · 选中c2 // 工具链使用最终配对协议将其与默认计划进行对比测量
基准测试
智能体需要基准测试才能体现价值,幸运的是相关基准已有人构建完成。
连接顺序基准测试(JOB)
Leis等人在论文《查询优化器到底有多好?》中提出了连接顺序基准测试(Join Order Benchmark, JOB),用于评估基于IMDb数据集的基数估算和连接顺序优化能力。
它包含113个查询,基于33个查询模板生成。查询模板的区别在于关系骨架,即引用的表和连接谓词不同,可以视为一类结构相似的问题。从模板生成的查询保留了表和连接图拓扑,但修改了选择性谓词。
示例如下:
SELECT MIN(t.title) AS movie_title
FROM company_name AS cn
JOIN movie_companies AS mc ON mc.company_id = cn.id
JOIN title AS t ON t.id = mc.movie_id
JOIN movie_keyword AS mk ON mk.movie_id = t.id
JOIN keyword AS k ON k.id = mk.keyword_id
WHERE cn.country_code = :country_code
AND k.keyword = 'character-name-in-title';
查询模板2——"与X国公司合作且带有'character-name-in-title'标签的电影中,字母排序最靠前的标题是什么?"
以下是JOB中基于该模板生成的两个实际查询:
SELECT MIN(t.title) AS movie_title
FROM company_name AS cn,
keyword AS k,
movie_companies AS mc,
movie_keyword AS mk,
title AS t
WHERE cn.country_code = '[de]'
AND k.keyword = 'character-name-in-title'
AND cn.id = mc.company_id
AND mc.movie_id = t.id
AND t.id = mk.movie_id
AND mk.keyword_id = k.id
AND mc.movie_id = mk.movie_id;
查询2a——"与德国公司合作的此类电影中,字母排序最靠前的标题是什么?"
SELECT MIN(t.title) AS movie_title
FROM company_name AS cn,
keyword AS k,
movie_companies AS mc,
movie_keyword AS mk,
title AS t
WHERE cn.country_code = '[us]'
AND k.keyword = 'character-name-in-title'
AND cn.id = mc.company_id
AND mc.movie_id = t.id
AND t.id = mk.movie_id
AND mk.keyword_id = k.id
AND mc.movie_id = mk.movie_id;
查询2d——"与美国公司合作的此类电影中,字母排序最靠前的标题是什么?"
基数估算基准测试(CEB)
另一个相关基准是基数估算基准测试(Cardinality Estimation Benchmark, CEB),由论文《Flow-loss: Learning Cardinality Estimates That Matter》提出。它使用相同的IMDb数据库,规模更大,包含约13600个合成生成的查询,基于16个查询模板构建CEB对模板的定义比JOB更宽松,两个CEB模板可能共享相同的连接图,仅选择性谓词不同;而JOB中每个模板的连接图都是唯一的。。
训练与测试划分
由于规模较大,CEB适合用于模型训练,而JOB用于验证模型性能。
你可能会问,同时在IMDb数据集上训练和测试是否合理?如果模型表现良好,是否只是因为它对这个特定数据库非常熟悉?
我认为这恰恰是关键所在。我们希望模型深入学习IMDb数据集。根据我们的问题建模,如果这个智能体持续用于某公司特定数据库的分析负载,我们无需将其泛化到所有数据库。
真正需要注意的是,确保训练过程中不会因使用CEB而过度拟合JOB的查询模板。模型应该以一种通用的方式学习IMDb,使其能够为任何查询生成良好的计划,即使是从未见过的结构化查询类型。为此,我们需要从训练集中剔除与JOB查询拓扑相同的CEB查询。
查询拓扑映射
我们将查询的"拓扑"定义为其结构化连接图(以去别名的表名为节点,连接关系为边)。连接图不包含任何选择性谓词,仅关注连接关系。
与JOB查询拓扑相同的CEB查询将被从训练集中移除。我编写了一个小脚本,将所有JOB和CEB查询转换为拓扑结构并检查是否存在重叠,结果发现没有重叠,因此无需过滤。
JOB:113个查询,33个模板,33种拓扑
- 1
- 2
- 3
- 4
- 5
- 6
- 7
- 8
- 9
- 10
- 11
- 12
- 13
- 14
- 15
- 16
- 17
- 18
- 19
- 20
- 21
- 22
- 23
- 24
- 25
- 26
- 27
- 28
- 29
- 30
- 31
- 32
- 33
CEB:13646个查询,16个模板,12种拓扑
- 1a
- 2a
- 2b
- 2c
- 3a
- 3b
- 4a
- 5a
- 6a
- 7a
- 8a
- 9a
- 9b
- 10a
- 11a
- 11b
JOB和CEB模板以拓扑结构的识别图标形式展示。
如何降低Postgres的噪声
在对智能体进行基准测试和训练之前,我们必须先讨论Postgres的运行方式,因为这直接影响训练过程。
首先列出一些基本情况:
- FLOPper的CPU有16个物理核心,64GB内存,2TB NVMe固态硬盘
- 我们使用的IMDb数据集切片磁盘占用8.5GB
- Postgres会将查询执行过程中检索到的数据页缓存到缓冲区中
- 操作系统也有自己的文件系统缓存,在更底层做同样的事情
如果连续20次运行同一个Postgres查询,每次的执行时间不会完全相同。在日常工作中这不算大问题,但我们的核心论点和训练过程都依赖于判断一种查询执行方式是否比Postgres默认方式更快,因此我们需要尽一切努力降低Postgres的噪声。
首先,我需要了解Postgres查询执行的噪声程度。
我在实验流程中加入了"校准"功能。校准过程很简单:运行基于Postgres镜像构建的Docker容器,每个容器分配固定的CPU核心和内存。初始设置为每个容器分配4个核心,内存上限8GB——设置过低会导致后续训练过慢,过高则可能加剧CPU竞争,进而增加噪声。
启动时,每个容器会用相同的配置初始化Postgres并加载IMDb数据。校准过程会启动一个大小为4的线程池,将所有113个查询放入共享队列。每当一个容器完成一个查询的测量,就会从队列中取下一个查询。
实际测量过程分为两个阶段:
先运行几次查询进行"预热"
再运行20次查询并记录每次的执行时间
容器0 — 预热 测量 空闲
容器1 — 预热 测量 空闲
容器2 — 预热 测量 空闲
容器3 — 预热 测量 空闲
那么,查询预热到底是什么意思?我们需要从操作系统基础原理入手理解。
每当Postgres执行查询时,会向操作系统(此处为Linux)请求数据页。Linux首先检查自身的文件系统缓存(页缓存),如果数据页已存在,则直接返回;否则从磁盘读取,存入缓存后再返回给Postgres。Postgres会将收到的数据页存入自己的shared_buffers缓存,以便重复使用。当shared_buffers满时,Postgres会淘汰部分页,若后续需要这些页,必须再次向Linux请求。
每次shared_buffers中命中数据页时,Postgres会增加一个名为"shared hit blocks"(SHBs)的计数器;若需要向Linux请求,则增加"shared read blocks"(SRBs)计数器。
如果运行带BUFFERS选项的EXPLAIN,Postgres会方便地返回这两个计数器的值。例如,运行EXPLAIN (ANALYZE, TIMING OFF, BUFFERS, FORMAT JSON)会输出类似如下内容:
{
"Plan": {
"Node Type": "Aggregate",
"Shared Hit Blocks": 1800786,
"Shared Read Blocks": 52990,
...
},
"Execution Time": 189.2,
...,
}
这些计数器让我们能够判断查询的"预热程度"。每次预热运行后,我们将其命中和读取计数与上一次运行进行比较。如果两者相差不超过2%(且计划未改变),则认为查询已预热完成,开始正式测量。查询至少需要两次预热才能进行比较,最多不超过五次。其核心思想是:如果计数器不再变化,说明数据已稳定,在后续20次测量中缓存波动将最小化。
查询A运行时,shared_buffers通过Linux调用加载数据。查询B需要不同的数据页,会淘汰shared_buffers中查询A的部分页。当查询A再次运行时,被淘汰的页会被计入读取次数。我将shared_buffers设置为保守的128MB并进行首次校准:
一半的查询仅运行两次就完成了预热。看起来不错……但深入分析后发现,SRB计数并未降至零,而是稳定在一个较大的数值。由于128MB的shared_buffers面对的是8.5GB的数据库,Postgres每次执行都会频繁缓存未命中,需要向Linux请求更多数据页。"稳定"并不意味着"数据已驻留"。
Linux的页缓存速度很快,这本身不是大问题,但当我查看不同查询的20次测量结果时,新的问题出现了。以查询job-13b为例:
job-13b 128 MB shared_buffers
job-13b在shared_buffers为128MB时的20次测量结果,每个点代表一次运行。切换按钮可按运行顺序查看结果,或按x轴分布查看,结果聚成两簇:14次运行耗时在186-204毫秒之间,另外6次在227-253毫秒之间,比前者慢14%-26%。查询的噪声并非均匀分布,而是存在两种不同的速度,且有三分之一的概率运行在较慢的速度下。
我最初想用变异系数来量化噪声:
变异系数衡量测量结果的"波动程度"。如果一个查询耗时100毫秒,变异系数为5%,则可认为波动约为5毫秒。对于job-13b,变异系数为10.3%,不算理想。而且变异系数也不是一个完美的衡量指标,因为它基于平均值,容易受到少数异常值的影响。
我们其实并不太关心20次运行的分布情况,真正关心的是这种噪声在训练过程中会多频繁地误导我们的测量标准。
误导智能体
请容我先跳过一部分内容,以便更清晰地说明我们需要测量的内容。
为降低智能体运行时的噪声,不能仅运行一次智能体提出的计划,而是要连续运行三组交错的(候选计划, 默认计划)配对。选择三组是为了在控制变量和避免智能体评估耗时过多之间取得平衡。得到三组候选/默认执行时间后,分别取两者的中位数并计算比值,以确定最终的加速或减速效果。如果两个中位数相差小于5%(人为设定的阈值),则视为平局;超出该范围则判定候选计划加速或减速。
回到job-13b的例子,我们有14次运行在一个簇,6次在较慢的簇。取三次运行中位数的策略看似不错,但要意识到:如果三次测量中至少两次落在"较慢"簇,中位数就会偏向出现频率较低的较慢簇。
假设有一个候选计划与默认计划执行速度完全相同,那么正确的奖励应为零。从观察到的20次运行中为"候选计划"和"默认计划"各抽取三次时间。从20次中抽取三次共有多种组合,对于job-13b,其中230种组合包含至少两次慢运行,因此约20%的概率某一方的中位数会落在慢簇。
这意味着约20%的概率会出现虚假的14%-26%加速或减速,并将其作为信号反馈给模型,十分危险!
job-13b 128 MB shared_buffers · 无操作候选计划(即与默认计划完全相同)
— 模拟抽取中…
0轮 · 平局0次 · 虚假胜利0次 · 虚假失败0次 · 误导率0%
从job-13b的20次运行中抽取三次作为候选计划、三次作为默认计划,取双方中位数并应用5%平局阈值。在所有可能的抽取组合中,约40%的情况会误导奖励判断。因此我们不能仅依赖变异系数来最小化噪声,因为两个变异系数相同的查询,其误导奖励的概率可能不同——取决于噪声是均匀分布还是分为两个相差超过5%的簇。我们需要最小化的实际指标就是这种误导率。
我编写了一个小脚本,直接从原始校准数据计算误导率。它通过在20次运行中滑动一个大小为6的窗口来实现,每个窗口取交错的两两组合作为(候选计划, 默认计划)配对。大小为6的窗口会产生(t1, t2), (t3, t4), (t5, t6)这样的配对。在任意一对中,两者可以互换候选计划和默认计划的角色,因此每个窗口有多种可能的组合。20次测量意味着窗口会滑动15次,因此每个查询共有多种可能的组合每种组合是一个二进制值,表示该候选/默认配对的中位数比值是否超出5%平局阈值。。
我们从这些原始数据中得出两个指标:首先,计算每个查询的"无操作错误率",即120种模拟组合中比值超出5%的比例;然后将所有113个JOB查询的该比例求和并除以113,得到"平均无操作错误率",即采用三组配对测量策略时,任意JOB查询的奖励被误导的概率。其次,将所有113个查询的无操作错误率从低到高排序,取第90百分位的值作为"p90查询"指标,衡量最易受影响的查询的误导率。
当shared_buffers为128MB且有四个并发容器时,"误导率"脚本得出以下平均无操作错误率和p90查询数值我对每种配置运行了两次校准,以体现两次运行的差异。:
| 运行次数 | 平均无操作错误率 | p90查询 | 中位数变异系数 |
|---|---|---|---|
| 1 | 5.0% | 13% | 2.3% |
| 2 | 5.4% | 20% | 2.4% |
结果并不理想:每20个无操作计划中就有一个会被错误奖励,约十分之一的查询误导率超过13%。
我们可以做得更好。
调优Postgres
我重点调整了Postgres的两个内存相关设置:
shared_buffers:决定Postgres可在自身缓存中保留的数据量work_mem:决定单个排序/哈希操作可使用的内存量,超出则溢出到磁盘
我进行了四次校准:
shared_buffers |
work_mem |
无操作错误率(运行1/运行2) | p90查询 | 中位数变异系数 | 总运行时间 |
|---|---|---|---|---|---|
| 128 MB | 4 MB | 5.0% / 5.4% | 13% / 20% | 2.3% | 95秒 |
| 2 GB | 4 MB | 1.8% / 1.2% | 1.3% / 0% | 1.1% | 60秒 |
| 128 MB | 32 MB | 7.0% / 6.6% | 20% / 23% | 2.6% | 94秒 |
| 2 GB | 32 MB | 1.7% / 1.3% | 0% / 0% | 1.2% | 60秒 |
令人惊讶的是,work_mem对噪声没有任何影响,而shared_buffers起了决定性作用!
当shared_buffers为2GB时,中位数查询完成预热后SRB计数器恰好为零:其工作集完全驻留在Postgres自身缓存中。无操作错误率降低了约4倍,第90百分位查询的误导率从13%-20%降至几乎为零。之前出现两簇结果的查询job-13b,变异系数从10.3%降至0.9%,20次运行的耗时相差不超过7毫秒。
还出现了一个我最初并未追求的好处:默认计划本身也变快了。所有113个JOB查询的总运行时间从95秒降至60秒,这仅仅是因为数据驻留在缓存中。换句话说,现在测量候选计划和默认计划的速度会显著加快,训练过程也会更省时。
我将shared_buffers设为2GB、work_mem设为4MB,并在项目后续阶段保持该配置。
基线与指标
我使用两个指标来衡量智能体的性能。
几何平均加速比
几何平均加速比赋予所有查询同等权重。例如,在一个包含两个查询的样本中,如果查询1比基线快2倍,查询2比基线慢一半(即0.5倍),则几何平均加速比为。无论查询1的基线耗时5分钟、候选耗时2.5分钟,而查询2仅从25秒退化到50秒,两者的权重都是相等的。
总负载加速比
总负载加速比将整个查询集视为一个批次,只需将所有基线时间求和,再除以所有候选计划时间的总和。在上述例子中,总负载加速比为。
两个指标各有侧重:总负载加速比衡量实际实用性,数据分析师构建分析查询套件时希望减少整个批次的总运行时间;但从模型训练角度看,总负载加速比可能完全受智能体偶然发现的单个查询计划影响,其余批次可能表现不佳,这意味着模型并未真正学到有用的知识,只是运气好。而几何平均加速比不关注绝对时间,能衡量模型在整个批次中的实际学习效果:大于1表示平均查询执行速度更快。
前沿模型对照实验
在让未训练的4B模型通过qo-agent工具链运行之前,我想验证这个问题是否真的能被当前的前沿模型解决。如果像GPT-6 Astra或Qwen 3.8 2.4T这样的模型都无法改进Postgres的默认查询计划,那么我也不能指望4B模型能做到。
我选取了10个JOB查询作为样本,测试Astra和Qwen 3.8 2.4T通过qo-agent工具链的表现:
| 模型 | 候选计划数量 | 完成测试的查询数 | 几何平均加速比 | 总负载加速比 | 性能退化查询数 |
|---|---|---|---|---|---|
| Astra [m](中等推理能力) | 1 | 9/10 | 0.85x | 1.00x | 3 |
| Astra [m](中等推理能力) | 5 | 10/10 | 2.54x | 2.12x | 0 |
| Astra [m, r](中等推理能力+推理摘要) | 5 | 10/10 | 2.39x | 1.57x | 1 |
| Qwen 3.8 2.4T [m](中等推理能力) | 1 | 7/10 | 2.02x | 1.30x | 1 |
| Qwen 3.8 2.4T [m](中等推理能力) | 5 | 10/10 | 2.26x | 1.35x | 1 |
Astra和Qwen 3.8 2.4T在相同的10个JOB查询样本上的测试结果。前沿模型在不同候选计划数量(即智能体轨迹中允许生成的候选计划数:1个或5个)下进行测试;对于Astra,还测试了是否启用推理摘要我有点惊讶地发现,启用推理摘要后Astra的性能比之前的5候选计划测试有所下降,但这些评估仅在10个JOB查询样本上运行了一次,因此我将性能下降归因于随机波动。。Astra通过OpenAI API推理,Qwen 3.8 2.4T通过Modal和OpenRouter推理。
从单候选计划和5候选计划的得分差异可以看出,智能体显然能够在候选计划的连续执行中进行上下文学习。这让我有信心坚持采用智能体多轮对话的方式,而非试图训练4B模型一次性生成最优计划。
在智能体运行过程中,每个候选计划先预热一次,再测量一次。用完候选尝试机会后,模型只能调用finish工具,并被要求选择得分最高的候选计划(或保留默认计划)。选中候选计划后,运行三组交错的(候选计划, 默认计划)配对,并通过限幅器处理结果:
限幅器将两个中位数的比值限制在之间。这些限幅数值是人为设定的,我发现它们可以防止几何平均加速比被极端加速或极端退化的情况过度影响。
结论:前沿模型能够通过智能体方式进行查询优化。
未训练4B模型基线测试
现在我们可以评估未训练的4B模型在JOB上的表现了!
同样采用每个智能体轨迹生成5个候选计划的配置,结果却不尽人意:
| 轨迹结束方式 | 查询数 |
|---|---|
| 无有效候选计划 | 81 |
| 选择失败 | 16 |
| 超时 | 1 |
| 候选计划与默认计划重复 | 7 |
| 保留默认计划 | 2 |
| 候选计划与默认计划对比测量 | 6 |
只有最后三类轨迹会被计入得分,113个查询中仅15个轨迹有效。其中9个轨迹的得分必然为1.00x(因为计划与默认计划相同如果去除成本和行数估算的EXPLAIN输出相同,则认为候选计划与默认计划等效。,或模型选择保留默认计划)。仅剩6个候选计划满足以下条件:
- 结构完整(模型生成的
PlanAction成功编译为提示注释) - 有效(生成的提示符合 Schema)
- 全新(生成的计划与Postgres默认计划不同)
其中5个计划实现了1.02x-1.30x的加速,1个计划的速度仅为默认的0.05x。
最常见的问题包括:
PlanAction不是有效的JSON对象- 连接树未包含所有关系恰好一次
- 操作被额外的
action键包裹 Leading树的子树与查询连接图实际不相连- 模型在错误的时机调用特定工具,或调用不存在的工具
这个模型不仅在任务上表现糟糕,甚至无法理解其所处的工具链。
基于监督微调的离线策略蒸馏
我首先需要让4B模型掌握qo-agent工具链的"语言",之后再让它擅长查询优化。
我们可以使用监督微调(SFT)来实现这一点,具体来说是离线策略蒸馏。
离线策略蒸馏是一种训练方法,学生模型(有时称为策略)通过模仿教师模型的输出来学习。之所以称为"离线",是因为训练数据并非由学生模型/策略自身生成。这种方法非常简单:向学生模型展示完整的教师轨迹(有时称为演示),对于轨迹中的每个 token,计算学生模型生成该 token 的概率对应的逐 token 损失,将这些损失平均得到演示级损失,然后通过标准的反向传播链式法则计算学生模型所有可训练参数的梯度,优化器据此调整参数值以最小化损失。
离线策略蒸馏的一大优势是不需要大量数据就能起效,因为每个演示能提供数千到数万个 token 预测样本,学生模型的权重可以快速适应。
但我们如何获取这些演示呢?手动编写耗时太长,次优方案是编写工具随机生成看似有效的轨迹有趣的是,我最初尝试过这种方法,效果尚可,能够教会4B模型使用工具链,但也带来了其他问题,主要是降低了模型生成全新候选计划的可能性。。不过我们已经有了最好的教师:更智能、更大的模型。
策略很简单:让智能模型通过qo-agent工具链生成大量轨迹,然后将这些轨迹蒸馏到我们的4B学生模型中。
渲染、损失掩码与展开
这里有一些细节需要说明。
假设我们选择GPT-6 Astra作为教师模型,它以OpenAI的Responses API格式生成轨迹,但Qwen模型无法理解这种格式,因此需要将人类友好的Responses API JSON格式转换为模型友好的token格式,这就是"渲染"的概念。渲染库可以将轨迹文本转换为特定模型能理解的原始token序列。
另一个需要考虑的问题是,确定模型实际需要预测哪些token。模型不会生成轨迹中的某些token,比如系统提示、用户提示或工具链执行后的工具调用结果。但模型在预测下一个token时仍需要看到这些token(它们是上下文的一部分),我们只需对模型负责生成的token计算损失。这时可以采用"损失掩码"策略:损失掩码库可以标记教师轨迹中哪些部分仅作为上下文,哪些部分需要学生模型预测。
最后,在针对智能体轨迹进行微调时,我们通常不会将整个轨迹作为单个训练单元,而是将其拆分为一组(上下文, 回复)对。这些对中的"上下文"是累加的,包含之前的模型回复。
轨迹会被展开为:
低秩适配(LoRA)
我们的4B模型有46.6亿个可训练参数,如果在SFT中每次更新所有参数,需要至少64GB显存的GPU。我希望先在消费级RTX 3090(单卡24GB显存)上验证我的假设,因此需要更高效的参数更新方法。
低秩适配器(LoRA)是实现这一目标的标准方法。从高层次来说,它冻结模型原有权重,转而训练一对更小的矩阵,两者相乘的结果可调整原始模型中选定的权重。
46.6亿个权重在bf16格式下占用9.32GB,而我最终训练的小型LoRA仅42.5MB,包含2120万个可训练参数。
教师模型选择
我们已经知道前沿模型可以在qo-agent中良好运行,接下来需要在GPT-6 Astra和Qwen 3.8 2.4T之间做出选择。
回顾两者在10个JOB查询上的测试结果,这次重点关注token级别的上下文长度:
| 每个轨迹5个候选计划 | Astra(带摘要) | Qwen 3.8 2.4T |
|---|---|---|
| 平均模型轮次 | 7.6 | 12.2 |
| 平均最终上下文长度(token) | 22009 | 45144 |
| 最大最终上下文长度(token) | 28663 | 73608 |
| 每轮平均输出token数 | 161 | 1795 |
| 每轮平均推理token数 | 65 | 1417 |
| 10个查询总耗时 | 3分13秒 | 15分51秒 |
前文提到的10个JOB查询的5候选计划评估,此处测量每个轨迹消耗的上下文长度。
从表面上看,Astra在各方面都更胜一筹,但它有一个重大缺陷:通过API运行时,不会返回推理token,而是用简短的摘要代替。而Qwen 3.8 2.4T是开源权重模型,可以完整返回推理token。
我担心基于仅包含推理摘要的轨迹进行训练会影响效果,论文《How to Steal Reasoning Without Reasoning Traces》恰好提到了这个问题:基于推理摘要训练会导致学生模型性能下降。为解决这个问题,作者提出了一种新方法——轨迹反转,即把推理摘要人工扩展为原始推理token可能的形式。虽然这并非Astra推理时实际生成的token,但更长的推理块能提升小模型的迁移性能。这让我稍感安心:如果选择Astra后性能不佳,我可以尝试轨迹反转。
另一个考虑因素是上下文长度。鉴于我最初在单张RTX 3090上进行SFT,需要确保任何轨迹的token序列长度不超过约50000,否则3090有限的24GB显存可能导致内存不足(OOM)。所有Astra轨迹都符合这个限制,但部分Qwen 3.8 2.4T轨迹超出了该范围。
我决定先尝试基于Astra轨迹进行SFT,如果性能不佳,再探索轨迹反转;如果仍不行,就租用更强大的GPU进行训练,提升全局上下文限制,改用Qwen 3.8 2.4T轨迹。
我首先从CEB中随机选取查询,生成120条带推理摘要的Astra轨迹,其中100条用于训练,20条作为保留验证集。所有轨迹通过Prime Intellect的renderers库转换为Qwen格式,进行适当的损失掩码处理,然后展开并打包为可用的训练演示数据。接着我使用Prime Intellect的prime-rl库进行SFT,训练约2100万个参数的LoRA。100条Astra轨迹展开打包后变为382个训练样本。我仅训练一个epoch(即每个样本恰好训练一次),批量大小设为1(每个演示更新一次权重)。最后在JOB上评估得到的LoRA:
| 检查点 | 生成至少一个有效候选计划的查询数 | 完成测试的查询数 | 几何平均加速比 | 总负载加速比 | 性能提升查询数 | 性能退化查询数 |
|---|---|---|---|---|---|---|
| 未训练4B模型 | 14/113 | 15/113 | 0.85x | 0.85x | 3 | 1 |
| 1个epoch | 48/113 | 44/113 | 0.72x | 0.76x | 5 | 16 |
单epoch适配器与未训练模型在所有113个JOB查询上的对比。当轨迹最终生成可测量的候选计划、与默认计划重复的计划或保留默认计划时,该查询被视为完成测试。性能提升/退化指完成测试的查询比默认计划快/慢超过5%。
结果令人鼓舞!适配器学会了使用工具链。
我可以选择生成更多新轨迹,或在现有数据集上训练更多epoch。鉴于后者成本更低,我选择了后者。
就在这时,我急于更快得到训练结果,于是在Lambda上租用了一个2x H100节点。
现在在H100上训练切换到H100意味着我们拥有80GB显存,而不是24GB。我们本可以大幅提升上下文token限制,但决定仍保持约50000的限制。,而非单张RTX 3090,我在现有LoRA基础上又进行了两次训练:第二个epoch和第三个epochH100上的训练速度快得多:第一个epoch在RTX 3090上耗时4小时,第二个epoch在H100上仅用45分钟。:
| 检查点 | 生成至少一个有效候选计划的查询数 | 完成测试的查询数 | 几何平均加速比 | 总负载加速比 | 性能提升查询数 | 性能退化查询数 |
|---|---|---|---|---|---|---|
| 1个epoch | 48/113 | 44/113 | 0.72x | 0.76x | 5 | 16 |
| 2个epoch | 85/113 | 108/113 | 1.08x | 1.04x | 12 | 8 |
| 3个epoch | 57/113 | 99/113 | 0.82x | 0.89x | 9 | 15 |
在相同的100条Astra轨迹上训练2个和3个epoch后,在JOB上的评估结果。在第二个和第三个epoch之前,qo-agent工具链进行了升级,允许模型在搜索后保留默认计划,因此更多查询被计入完成测试。2个和3个epoch的评估设置完全相同。
两个epoch提升了结果,但三个epoch导致性能退化!这特别值得注意,因为第二个epoch期间验证损失完全没有变化:
20条保留轨迹的验证损失
优化器更新次数,每个打包训练样本对应一次更新
这给我们上了很好的一课:验证损失平稳并不一定意味着模型停止学习有用的行为。
不过我认为在进行RL之前,SFT还有更多潜力可挖。我之前对训练轨迹数据集没有做任何过滤,也没有仔细检查是否遗漏了某些能力。事实证明确实如此,主要是模型构建有效Leading树的能力不足。
我又生成了320条Astra轨迹,其中300条用于训练,20条用于验证。我过滤掉了6条Astra未尝试任何候选计划直接保留默认计划的训练轨迹。我分两次运行,每次训练两个epoch:
| 检查点 | 生成至少一个有效候选计划的查询数 | 完成测试的查询数 | 几何平均加速比 | 总负载加速比 | 性能提升查询数 | 性能退化查询数 |
|---|---|---|---|---|---|---|
| 前100条轨迹训练2个epoch | 85/113 | 108/113 | 1.08x | 1.04x | 12 | 8 |
| + 新300条轨迹训练1个epoch | 77/113 | 101/113 | 1.10x | 1.05x | 29 | 13 |
| + 新300条轨迹训练2个epoch | 71/113 | 107/113 | 1.16x | 1.06x | 20 | 5 |
在2个epoch适配器的基础上继续训练300条新轨迹后,在JOB上的评估结果。
我们的4B模型不仅学会了使用工具链,还能在多个JOB查询上做出真正有效的决策!幸运的是,基于推理摘要训练并未损害性能。
让4B模型精通查询优化
模型现在已经掌握了qo-agent工具链的"语言",并且从SFT中获得了一些额外的性能提升。接下来要让它真正擅长查询优化。
智能体强化学习
智能体RL与SFT的不同之处在于,我们会在智能体工具链中用当前策略对训练查询运行多次。每次运行(也称为轨迹)会产生最终输出,并根据可验证的标准评分。最后,每条轨迹的得分会相对于同一查询的其他轨迹进行加权:正"优势"会强化模型,使其未来更可能生成该轨迹;负优势则会惩罚模型,降低未来生成该轨迹的概率。
设计每条轨迹的奖励与相对优势
最初的奖励算法很简单:
加速比计算为三次默认计划测量的中位数除以三次候选计划测量的中位数。
对于每个无效的候选计划,我们从加速比的自然对数中减去0.1;如果轨迹生成的计划与Postgres默认计划的指纹相同,再减去0.05;如果轨迹最终没有有效候选计划,则直接从奖励中减去3,替代之前的0.1或0.05扣除。
使用这个奖励算法进行的前几次RL训练,导致模型因-3的严厉惩罚而害怕生成无效计划,转而采取保守策略,反复返回Postgres默认计划,接受较小的0.05奖励损失。
GRPO算法加剧了这个问题。普通GRPO算法将多条轨迹的奖励转换为相对"优势":
prime-rl实现的GRPO:简单地将每条轨迹的奖励减去该组的平均奖励。GRPO论文还会除以该组的标准差。
假设我们对某个查询执行四条轨迹,得到以下计划、执行速度和奖励:
prime-rl的GRPO。该组中没有计划击败Postgres,但仍有两条轨迹获得正优势,因为它们超过了组平均奖励。这会强化不良行为,让模型认为生成与Postgres默认计划等效的计划是可以接受的!
有缺陷的奖励算法和GRPO都需要替换为能相对于默认计划执行时间评分优势的方法。
我对奖励算法进行了如下更新:
修改内容包括:限制加速比的范围以约束奖励标量值,设置0.05的软阈值以应对测量噪声,若智能体调用keep_default或finish(default)则得分为0;轨迹最终无有效候选计划时扣除0.1(而非之前的3.0),与默认计划指纹相同的候选计划扣除0.02。
我将prime-rl风格的GRPO替换为定制的"锚定"变体:
让我们看看修改后的GRPO在上述四条轨迹中的表现:
修改后的GRPO成功地为这四条表现不佳的轨迹都赋予了负优势,全面降低了它们未来被生成的可能性。
开始训练
确定奖励和相对优势模型后,我就可以开始训练了……
不过首先需要搭建一个机制,让Postgres测量在FLOPper上进行,训练和推理在2x H100节点上进行。我本想在Lambda节点上进行Postgres测量,但在他们的机器上运行了多次Postgres校准后,我有理由相信自己与其他用户共享了机器的非GPU部分,因为噪声比FLOPper上大得多。
我通过Tailscale将FLOPper与Lambda节点连接起来,流程如下:
- 查询轨迹在FLOPper上启动,并在整个生命周期内占用一个可用的Postgres容器
- 轨迹通过第一个H100上运行的vLLM进行推理,生成轨迹
- 所有测量在占用的Postgres容器上进行
- 最后,轨迹结果发送到第二个H100,计算相对优势并更新权重
现在终于可以开始训练了。我首先将最终的SFT适配器合并到基础模型中,创建一个新的基础模型,然后为RL训练初始化一个新的适配器。
我从一个极其保守的学习率1e-06开始,批量大小设为8,每个查询四条轨迹,仅进行120次优化器更新以验证机制是否可行。
在JOB上评估训练后的LoRA,结果完全失败:
| 检查点 | 生成至少一个有效候选计划的查询数 | 完成测试的查询数 | 几何平均加速比 | 总负载加速比 | 性能提升查询数 | 性能退化查询数 |
|---|---|---|---|---|---|---|
| SFT起点 | 71/113 | 107/113 | 1.16x | 1.06x | 20 | 5 |
| RL 120次更新 | 71/113 | 106/113 | 1.14x | 0.99x | 21 | 7 |
第一个RL检查点与其起点SFT检查点在所有113个JOB查询上的对比。
问题可能出在奖励设计不合理、修改后的GRPO无效,或者我们的训练不够激进。
后者最容易测试。我将学习率提高一个数量级,从1e-06改为1e-05,批量大小增加到16,每个查询的轨迹数从4条增加到8条,并决定进行600次优化器更新,而非仅120次。
这时我发现,一条轨迹92%的时间都花在了vLLM推理上!鉴于优化器更新次数增加,我需要更高效地利用资源。
我重构了训练流程,让轨迹仅在需要测量时才租用可用的Postgres工作节点。这种异步方法让我能够更充分地利用vLLM和训练器,同时运行20条并发轨迹,共享四个Postgres工作节点。
整个轨迹周期占用工作节点 · 4条并发
完成轨迹数0 工作节点测量时间占比0%
仅在测量时租用工作节点 · 20条并发
完成轨迹数0 工作节点测量时间占比0%
要求轨迹全程占用工作节点会将最大并发轨迹数限制为可用工作节点数;而仅在需要Postgres时租用工作节点,可运行20条并发轨迹,提高Postgres工作节点的实际利用率。
我连续进行了两次600次优化器更新的RL训练,获得正优势的轨迹比例稳步上升:
获得正优势的轨迹占比
优化器更新次数
两次600次更新训练中,每个到达训练器的轨迹获得的锚定奖励信号。击败Postgres的轨迹占比和获得正优势的轨迹占比在训练过程中稳步上升。
两个检查点在JOB上的评估结果如下:
| 检查