Small-Model Distillation, Part 6: Training a 9B Model to Fix SQL Queries

TL;DR

Can a small language model learn to fix SQL queries by copying stronger models and practicing on its own mistakes? I trained Qwen3.5-9B on this task. Its best version solved 115 of 220 evaluation tasks, matching GPT-5.5 at medium reasoning effort. The base model solved 78.

Successful teacher examples provided most of the improvement. Correcting the student’s own mistakes improved the score further, but every stage also lost some previously solved tasks. The reported score comes from the best checkpoint selected on these evaluation tasks.

SQL repair task and evaluation setup

The model receives a user issue and a buggy SQL query. It can inspect the schema, run queries against a real SQLite database, read the results, and submit corrected SQL. It has eight turns. A task counts as solved only when the submitted query passes hidden tests the model never sees.

The data comes from birdsql/six-gym-sqlite, filtered to Query tasks across four SQLite databases. There are 879 training tasks and 220 evaluation tasks: books 56, chinook 50, movie_3 55, and netflix 59.

The task split, prompt, tools, parser, database execution, eight-turn limit, and scorer stayed fixed. Evaluation tasks never entered the training data. They did influence checkpoint selection, so they were held out from training but were not an independent final test.

Both the base model and trained checkpoints were evaluated on the Mac. The base used 8-bit MLX with a 1,024-token response limit per turn; the trained checkpoints used BF16, a 16-bit number format, through Hugging Face and PEFT with a 512-token limit.

Learning from teacher demonstrations

The first training stage used successful conversations from Qwen3.5-35B-A3B. Through supervised fine-tuning, or SFT, the student learned to reproduce the teacher’s actions in context. Here, distillation means learning from the teacher’s written responses, rather than matching its internal probability distributions.

Each training example pairs the task and conversation so far with the next assistant action. Earlier actions and tool results provide context; inspecting a table, running SQL, or submitting an answer supplies the target. Training therefore covers decisions throughout the conversation, including how to respond to a tool result.

The first stage took the student to 96 solved evaluation tasks. Training continued from the same adapter using successful DeepSeek V4 Pro conversations, bringing the score to 107. A successful teacher conversation shows one route from the original task to a correct answer. It does not necessarily show how to recover from the particular mistake the student makes halfway through a conversation.

Learning from student failures

The next two rounds focused on the student’s own mistakes. The current student attempted the training tasks, then DeepSeek continued from states where it had failed. Only continuations whose final SQL passed the verifier entered training. This is a DAgger-style correction loop: collect examples at states visited by the student, obtain better actions there, then train on them. The training step was still SFT.

View full size ↗

The failed student action was never a training target. For a correction example, the preceding conversation supplied the context and the verified teacher action supplied the target. Earlier messages and tool output were masked out of the loss, so the model was trained on the new action.

Each correction round also included successful student conversations, called retention data, to give the model practice at tasks it already solved. Round 2 added earlier successes from the 107-point model on training tasks the 112-point model now failed. Retention was intended to preserve capability. There was no comparison run without it, so its effect on forgetting remains unmeasured.

Training data and settings

A conversation can produce several training rows, one for each action. Row counts therefore differ from counts of tasks or successful conversations.

After the initial SFT run, each subsequent stage continued from the preceding selected adapter: the versions scoring 96, 107, and 112, respectively. The training schedules were:

All four stages used BF16 LoRA, which updates a small set of adapter weights while keeping the underlying model weights fixed. LoRA rank and alpha were both 32, dropout was zero, and the batch size was one with gradients accumulated over eight examples. The random seed was 42. The length caps above include context and target; they are separate from the generated-token limits during evaluation.

SQL example: returning all tied results

A task from the netflix database shows why a correction can help. The user wanted the items with the most view summaries and their count. The original SQL grouped movie and season ids together and ended with ORDER BY COUNT(a.f) DESC LIMIT 1.

The student fixed one problem: null ids were being counted as a group. After inspecting the schema, it added WHERE movie_id IS NOT NULL to one half of the union and WHERE season_id IS NOT NULL to the other. Its submitted query still ended in LIMIT 1.

That left a second problem. LIMIT 1 returns one row even when several items share the highest count, while the question requires all tied items.

The successful teacher continuation ran a query before submitting, replacing LIMIT 1 with a subquery that computes the maximum count and returns every row matching it. The database came back with two rows: item 11189 with 21 views, and item 11220 with 21 views. The teacher then submitted the corrected query, and the hidden tests passed.

The successful correction came from the second branch. A continuation from the last state before failure did not pass verification, but a retry from one turn earlier did. Both branches stayed inside the original eight-turn budget.

The following illustration isolates the tie-handling change. It is explanatory SQL rather than the full recorded query; item_counts stands for the grouped counts:

-- One row, even when several items share the highest count:
SELECT item_id, views
FROM item_counts
ORDER BY views DESC
LIMIT 1;

-- Every item whose count equals the maximum:
SELECT item_id, views
FROM item_counts
WHERE views = (SELECT MAX(views) FROM item_counts);

The training example paired this earlier conversation with the teacher’s verified next action.

Correction data yield

Round 2 ran the current student across all 879 training tasks. It solved 435, which supplied retention data. The failed attempts produced 870 candidate states, up to two per task. DeepSeek was asked to continue from 803 of them. The other 67 were skipped because another state for that task had already produced a verified correction, or because of an infrastructure error.

Of those 803 requested continuations, 94 passed verification: 11.7%. Only verified continuations entered the correction data. The other 709 requests yielded no accepted correction. The 11.7% measures how often a request produced usable training data. It does not measure teacher accuracy across unique tasks: a task could produce two candidate states, and a request could fail operationally before producing testable SQL.

Model performance

The final student and GPT-5.5 medium each solved 115 tasks. DeepSeek’s reported 120 combines an initial run with a retry pass and still includes seven infrastructure failures. The retries give that reference a different attempt history from a single evaluation pass.

View full size ↗

Round 2 saved five checkpoints, scoring 108, 115, 114, 111, and 112. I selected the best, checkpoint 100. Continuing training did not improve the score. The selected score needs confirmation on new tasks.

Task-by-task comparison

The paired comparison shows where the two models agree and disagree.

Both models solve the same 96 tasks and disagree on 38, split evenly between them. Their union is 134 solved tasks. That is an oracle upper bound: reaching it would require knowing which model’s answer to choose for each task. No router or combined system was tested.

Compared with GPT-5.5 medium, the final 9B model solves more tasks in two databases and fewer in the other two.

Gains and regressions across training stages

View full size ↗

Each row compares the new checkpoint with the one immediately before it. The totals progress from 78 to 96, 107, 112, and 115. Newly solved counts cannot be added across rows to count unique tasks: a task can be lost and regained. Every stage introduced regressions, including the correction rounds that used retention data.

Across the full sequence, the final model solves 43 tasks the base model could not and loses six the base model could solve. The headline scores alone would hide those six regressions.

Failure analysis

View full size ↗

Compared with the base run, the final model submits more often, repeats itself less, and hits the turn limit less. Parse failures stayed near zero.

Of the 105 remaining failures, 72 reached a submission whose SQL failed the hidden tests. The other 33 were repeated actions, a parse failure, or a turn-limit stop. Most remaining failures were wrong answers, but operating the agent loop was still part of the problem. These categories alone do not identify which kinds of joins, grouping, or edge cases caused the wrong answers.

GPT-5.5 medium also fails 105 tasks, but every failure in this run is a submission that failed the tests. It has no recorded loop-completion failures. The same success total therefore hides a different mix of errors.

Experiment costs

The recorded outlay was roughly €50, but the records do not provide an exact cost per training example or solved task.

The records show the credit purchase, but not how much of that balance the experiment consumed. These amounts also exclude engineering time.

GPT-5.5 access came through a subscription, without a metered API bill for its evaluation. Production throughput and serving cost for the 9B model were not measured. These records support a small experimental outlay, but they do not establish serving savings or a break-even point.

Conclusions

A 9B model reached the same selected score as GPT-5.5 medium on this SQL workload. Learning from successful teacher conversations provided most of the gain, and corrections from the student’s own failed states added another eight solved tasks. Every stage also introduced regressions. Comparing results task by task exposed losses that the aggregate scores concealed.

The selected checkpoint matched the reference total on these 220 tasks. Performance on new tasks and the cost of serving the model remain unmeasured.

References

• BIRD’s six-gym-sqlite dataset: source of the SQL repair tasks.
• Ross, Gordon, and Bagnell, “A Reduction of Imitation Learning and Structured Prediction to No-Regret Online Learning”, AISTATS 2011: the DAgger method behind collecting expert actions at student-visited states.