Het trainen van een 4B-model om 81% snellere queryplannen te produceren dan Postgres

Inleiding

Kan een klein model met open gewichten, via supervised fine-tuning (SFT) en agentic reinforcement learning (RL), worden getraind om Postgres-queryplannen te produceren die beter zijn dan de standaardplannen van Postgres?

Het antwoord is een volmondig ja. De belangrijkste resultaten zijn:

• Een reductie van de latentie met 44,7% over 113 join-intensieve queries met een 4B-model dat aanvankelijk voor 99 van deze queries geen bruikbaar plan kon produceren.
• De constructie van een Postgres-meetopstelling die ruis door Linux page cache contention minimaliseert over concurrerende containers.
• Het ontwerpen van een aangepaste GRPO-variant voor het scoren van RL-rollouts in een inherent ruisgevoelige omgeving.
• De splitsing van RL over twee machines: vLLM en de trainer op een gehuurde 2x H100-node, en vier Postgres-containers op een lokale desktop.
• Het uitvoeren van off-policy distillation over ruim vijfhonderd GPT-6 Astra agent-trajecten.

Hoe een query-optimizer werkt

De basis

Stel dat we een deel van de IMDb-dataset hebben met de volgende tabellen:

• title: Films, series, etc. (~1M rijen)
• movie_companies: Koppeltelling tussen films en bedrijven (~2M rijen)
• company_name: Namen en oorsprong van bedrijven (~100k rijen)
• company_type: Rollen van bedrijven (4 rijen)
• kind_type: Type titel (7 rijen)

Wanneer we een query uitvoeren, zoals: "Welke Japanse bedrijven brachten in de jaren 2000 de meeste titels uit?", moet Postgres bepalen in welke volgorde de tabellen worden gecombineerd (joins). Dit is cruciaal omdat de volgorde de hoeveelheid te verwerken data bepaalt.

Combinatorische explosie

Het bepalen van de optimale join-volgorde is een NP-hard probleem. Voor een eenvoudige query met drie tabellen zijn er al verschillende mogelijkheden. De complexiteit explodeert echter zodra er meer tabellen bij komen:

• 2 tabellen: 96 mogelijke plannen.
• 3 tabellen: 4.608 mogelijke plannen.
• 4 tabellen: 442.368 mogelijke plannen.
• 5 tabellen: 33.177.600 mogelijke plannen.
• 10 tabellen: ~2,5 sextiljoen (2,5 $\times 10^{18}$) mogelijke plannen.
• 12 tabellen: ~72 sextiljoen (7,2 $\times 10^{22}$) mogelijke plannen.

Postgres gebruikt dynamisch programmeren (en een genetisch algoritme voor 12+ joins) om deze zoekruimte in te perken.

Schatten in plaats van tellen

Postgres kan niet simpelweg alle plannen uitvoeren om te zien welke het snelst is; dat zou het doel van een optimizer verslaan. In plaats daarvan gebruikt Postgres statistieken uit de pg_statistic-tabel om de cardinaliteit (het aantal rijen) te schatten.

Een groot probleem hierbij is de aanname van een uniforme distributie. Postgres gaat ervan uit dat waarden gelijkmatig verdeeld zijn. Als deze aanname echter niet klopt—bijvoorbeeld als een kleine groep Japanse bedrijven verantwoordelijk is voor een onevenredig groot deel van de films—kan de optimizer een plan kiezen dat orders van grootte trager is dan het alternatief.

Postgres sturen met pghintplan

Omdat we het kostenmodel van Postgres niet kunnen wijzigen zonder de broncode aan te passen, maken we gebruik van pghintplan. Dit is een extensie waarmee we via gestructureerde "hints" in SQL-commentaren Postgres kunnen dwingen specifieke plannen te kiezen.

Voorbeeld van een hint:

/*+ HashJoin(a b) SeqScan(a) */
SELECT * FROM pgbench_branches b JOIN pgbench_accounts a ON b.bid = a.bid ORDER BY a.aid;

Hier dwingen we een HashJoin af en een SeqScan op tabel a.

Probleemformulering

De centrale vraag is: Kan een taalmodel leren om hints te produceren die resulteren in betere queryplannen?

De focus ligt hierbij op zware analytische workloads. Voor een eenmalige query is de overhead van een model te groot, maar voor queries die duizenden keren worden uitgevoerd, is de geamortiseerde winst enorm.

Model en Harness

Ik heb gekozen voor een klein 4B-model (empero-ai/Qwen3.8-4B-Distill) om training en inferentie op eigen hardware (2x RTX 3090) mogelijk te maken.

Er is een lichtgewicht agent-harness gebouwd, qo-agent, met de volgende tools:

• inspect_relation: Lijst kolommen, types, indexen en geschatte rijen.
• getcolumnstats: Haalt planner-statistieken op voor 1-8 kolommen.
• get_plan: Haalt de standaardplanning of een voorgestelde kandidaat op.
• evaluate_candidate: Valideert en voert een voorgesteld plan uit voor timing-diagnostiek.
• keep_default: Beëindigt de zoektocht en kiest het standaardplan van Postgres.
• finish: Selecteert de beste kandidaat of het standaardplan.

Benchmarks

Join Order Benchmark (JOB)

De JOB bestaat uit 113 queries verdeeld over 33 templates, gebaseerd op de IMDb-dataset. Deze benchmark evalueert cardinaliteitsschattingen en join-order optimalisatie.

Cardinality Estimation Benchmark (CEB)

De CEB is veel groter (~13,6k synthetische queries over 16 templates). Vanwege de omvang is deze benchmark ideaal voor het trainen van het model.

Om overfitting te voorkomen, zijn queries uit de CEB die dezelfde topologie (join-graaf) hebben als JOB-queries, uit de trainingsset verwijderd.

Ruis minimaliseren: "De olifant temmen"

Het meten van query-snelheid in Postgres is ruisgevoelig door de buffer cache van Postgres en de page cache van Linux.

De "Fooling Rate"

Om te voorkomen dat het model wordt beloond voor toevallige variaties in snelheid (ruis), is er een meetprotocol ontwikkeld. In plaats van één run, worden drie interleaved (kandidaat, default) paren uitgevoerd. De mediaan van deze runs wordt genomen.

Om de ruis te kwantificeren, is de no-op error rate berekend: hoe vaak een plan dat identiek is aan het standaardplan toch als "sneller" of "trager" wordt gemeten.

Tuning van Postgres

Door te experimenteren met geheugeninstellingen bleek dat shared_buffers de grootste impact had op de ruis:

• 128 MB shared_buffers: Hoge ruis; veel reads uit de Linux page cache.
• 2 GB shared_buffers: De ruis daalde drastisch omdat de working set volledig in het Postgres-geheugen paste (resident).

Voor de rest van het project zijn sharedbuffers = 2 GB en workmem = 4 MB gebruikt.

Baselines en statistieken

Twee metrieken worden gebruikt om de prestaties te meten:

1. Geometric Mean Speedup ($S_{geo}$):

$$S{geo} = \left( \prod{i=1}^{N} \frac{bi}{ci} \right)^{1/N}$$ (Waar $bi$ de default tijd is en $ci$ de kandidaat tijd. Dit geeft elke query een gelijk gewicht.)

1. Total Workload Speedup ($S_{workload}$):

$$S{workload} = \frac{\sum bi}{\sum c_i}$$ (Dit meet de totale tijdwinst over de hele batch queries.)

Controle met frontier-intelligentie

Om te valideren of het probleem oplosbaar is, zijn GPT-6 Astra en Qwen 3.8 2.4T getest. Beide modellen konden via de agent-harness aanzienlijk betere plannen vinden dan Postgres standaard, wat bewees dat in-context learning over opeenvolgende executies werkt.

De vanilla 4B baseline

Het ongetrainde 4B-model presteerde zeer slecht. Slechts 15 van de 113 trajecten resulteerden in een gescoord plan, en de meeste daarvan waren identiek aan het standaardplan. Het model begreep de syntax van de harness niet en produceerde vaak ongeldige PlanAction objecten.

Off-policy distillatie via Supervised Fine-Tuning (SFT)

Om het model de "taal" van de harness aan te leren, is off-policy distillation gebruikt. Hierbij imiteert het studentmodel (4B) de trajecten van een teacher-model (GPT-6 Astra).

Technische implementatie

• Rendering & Masking: Trajecten werden omgezet naar een token-formaat dat Qwen begrijpt. Er is loss-masking toegepast zodat het model alleen wordt getraind op de tokens die het zelf moet genereren, niet op de output van de tools.
• LoRA (Low-Rank Adaptation): Vanwege VRAM-beperkingen op de RTX 3090's is LoRA gebruikt. In plaats van 4,66 miljard parameters te trainen, werden slechts 21,2 miljoen parameters in kleine matrices geoptimaliseerd.

Resultaten van SFT

Na training op 100 Astra-trajecten over meerdere epochs, leerde het model de harness. Het begon valide plannen te produceren en behaalde een $S{geo}$ van 1.16x en een $S{workload}$ van 1.06x.

Het 4B-model optimaliseren via Reinforcement Learning

Na SFT was het model in staat om de tools te gebruiken, maar nog niet optimaal te optimaliseren.

Agentic Reinforcement Learning

Het model voert $N$ rollouts uit per query. De resultaten worden gescoord tegen de default plan-tijd.

Beloningsontwerp en GRPO

Aanvankelijk werd standaard GRPO gebruikt, maar dit leidde ertoe dat het model "veilig" speelde door simpelweg het default plan te kopiëren. Er is daarom een anchored GRPO-variant ontwikkeld.

De beloning ($q_i$) is nu gebaseerd op:

• De logaritme van de speedup (geclipped tussen 0.1x en 10x).
• Een soft-threshold van 0.05 om ruis te negeren.
• Strafpunten voor ongeldige plannen of duplicaten van het default plan.
• Een baseline ($b_i$) die de maximale waarde is van 0 en de gemiddelde kwaliteit van de andere rollouts in de groep.

Het trainingsproces

De infrastructuur was verdeeld:

• FLOPper (Local): Voerde de Postgres-metingen uit in containers.
• Lambda (Cloud): 2x H100 node voor vLLM-inferentie en weight updates.

De training bestond uit twee runs van 600 optimizer-updates. De effectiviteit nam gestaag toe naarmate het model vaker plannen vond die sneller waren dan die van Postgres.

Resultaten

De finale resultaten op de JOB-benchmark zijn als volgt:

Checkpoint	Valide kandidaten	Tasks scored	$S_{geo}$	$S_{workload}$	Wins	Regressions
SFT (Startpunt)	71/113	107/113	1.16x	1.06x	20	5
RL (600 updates)	99/113	113/113	1.35x	1.16x	34	0
RL (1200 updates)	101/113	112/113	1.41x	1.29x	38	2

Bij een best-of-15 evaluatie (waarbij de beste van drie trajecten wordt gekozen), steeg de prestatie naar een $S{geo}$ van 1.81x en een $S{workload}$ van 1.81x. Dit komt neer op een totale latentiereductie van 44,7%.

Wat het model heeft geleerd

Analyse van de trajecten laat zien dat het model:

• Onderzoekt: In 87% van de gevallen inspecteert het eerst de relaties of statistieken.
• Redeneert: Het model identificeert specifieke inefficiënties (bijv. "lossy filters") en dwingt alternatieven af (bijv. bitmap scans).
• Strategieën kiest: Het model geeft de voorkeur aan Nested Loops over Hash Joins in specifieke scenario's en gebruikt vaak enablesort=off en randompage_cost=1.1.

Kosten en Conclusie

Kosten

De totale kosten bedroegen $1.200:

• ~$800 voor de gehuurde H100-node (Lambda).
• ~$400 voor OpenAI API-fees (Astra teacher-trajecten).

Conclusie

Via off-policy distillatie en reinforcement learning is een klein 4B-model getransformeerd van een model dat de toolset niet begreep, naar een specialist die queryplannen produceert die significant sneller zijn dan die van de Postgres-optimizer.

Dit experiment toont aan dat bedrijven met specifieke domeindata relatief goedkoop kleine, open-weights modellen kunnen trainen voor niche-taken, waarbij ze de kracht van frontier-modellen gebruiken voor de initiële distillatie en RL voor de uiteindelijke optimalisatie.