trl-lab

qwen3.5-2b-grpo-sqale

README

License: apache-2.0

Quickstart

The model is an agent, so it needs a loop that executes its tool calls against a database and feeds the results back. sqale_agent.py in this repository is a self-contained reference loop that reproduces the training environment: the same prompt, tools, turn budget and final-answer fallback. It opens your database read-only.

bash

pip install vllm huggingface_hubhf download trl-lab/qwen3.5-2b-grpo-sqale sqale_agent.py --local-dir . python sqale_agent.py --model trl-lab/qwen3.5-2b-grpo-sqale \  --db path/to/database.sqlite \  --question "Which supplier shipped the most back-ordered items last quarter?" \  --trace

From Python, one agent can answer many questions with a single loaded engine:

python

from sqale_agent import SQaLeAgent, vllm_generate agent = SQaLeAgent(*vllm_generate("trl-lab/qwen3.5-2b-grpo-sqale"))episode = agent.run("How many customers placed more than three orders in 2024?", "shop.sqlite")print(episode["sql"])       # the submitted queryprint(episode["ended_by"])  # submit_sql, final_turn, fallback_to_last_working_query, ...

To build your own loop, follow the protocol in The agent loop below.

Highlights

+33.0 points on BIRD dev over the untrained base model (19.3% → 52.3%), and +27.6 points on the SQaLe test set (38.7% → 66.3%).
Strong for its size. On 100 moderate SQaLe test questions it reaches 52% at 224 TFLOPs per question. Every general-purpose model in the 9–12B range except Qwen3.5-9B scores lower, and Qwen3.5-27B is 11 points ahead at 1,710 TFLOPs per question.
Finds the right tables in large schemas. On full-size SQaLe test schemas it opens every gold table in 89.7% of episodes.
Reads values from the data. Before submitting, it has looked up 88% of the string literals the gold answer depends on. The untrained base model looks up 35%.

The SQaLe model family

The paper trains Qwen3.5-2B three times with one recipe, once per training corpus, to measure how much of a small model's ability comes from the data it is trained on. Base model, agentic environment, tool set, reward, GRPO objective, step count, batch size and evaluation protocol are identical across the three. Only the training corpus and the strength of the length curriculum differ, so the gaps between the models come from the data.

Table with columns: Model, Training corpus, Length curriculum
Model	Training corpus	Length curriculum
MSQaLe (this model)	SQaLe	full, four length buckets
MBIRD	BIRD train	none
MSynSQL	SynSQL-2.5M

Evaluation

Execution accuracy (%) on 300 questions per benchmark, as reported in the paper. Every model runs in the agentic environment used for training, with the schema withheld and no test-time scaffolding. The SQaLe test split shares no schema with the training split and is evaluated over the full schema. EHRSQL is a single-domain benchmark over electronic health records.

Table with columns: Model, SQaLe test, BIRD dev, EHRSQL
Model	SQaLe test	BIRD dev	EHRSQL
MSQaLe (this model)	66.3	52.3	23.7
MBIRD	54.0	54.7	23.7
MSynSQL	50.7	44.3	13.3
Qwen3.5-2B (untrained)	38.7	19.3	8.2

On the SQaLe test set, training on SQaLe adds 27.6 points to the base model. The same recipe adds 15.3 points when trained on BIRD and 12.0 points when trained on SynSQL.

Schema size

The SQaLe test questions are also evaluated at three database sizes: the gold tables only, the gold tables plus 32 distractor tables, and the full schema. Brackets give the share of episodes in which the model opened every gold table.

Table with columns: Tables in the database, MSQaLe, MBIRD, MSynSQL
Tables in the database	MSQaLe	MBIRD	MSynSQL
Gold tables only	72.7 (95.3)	64.7 (98.7)	59.3 (98.3)
+ 32 distractor tables	68.0 (96.7)	61.0 (90.7)	52.3 (92.7)
Full schema	66.3 (89.7)	54.0 (85.0)	50.7 (83.3)

Accuracy falls for every model as the schema grows, and MSQaLe leads at every size. When only the gold tables are present, all three models find them. At full size, MSQaLe opens all gold tables in 89.7% of episodes, against 85.0% for MBIRD and 83.3% for MSynSQL. BIRD and SynSQL schemas hold 5–10 tables, so the models trained on them never had to learn to search for the relevant tables. MSQaLe learns this search during training.

Looking up values

Many questions depend on literals that exist only in the table values. Before submitting, MSQaLe has looked up 88% of the string literals the gold answer depends on, against 55% for MSynSQL, 40% for MBIRD and 35% for the base model. The difference is largest for literals the question does not state. MSQaLe searches the data for 61% of them and uses the correct value in 55% of cases, against 41% for MBIRD and 39% for MSynSQL.

The agent loop

What an episode looks like

text

Question   Which supplier shipped the most back-ordered items last quarter?Turn 1     {"tool": "list_tables"}           → suppliers, shipments, order_items, … (113 tables)Turn 2     {"tool": "join_path", "args": {"table_a": "suppliers", "table_b": "order_items"}}           → suppliers → shipments → order_items   ⋮Turn 7     {"tool": "submit_sql", "args": {"sql": "SELECT s.name …"}}

Each turn opens with a <think> block in which the model plans its next call, and a turn may hold several calls.

Using your own harness

The following details from training matter for accuracy if you build your own loop.

Use the training prompt verbatim. SYSTEM_PROMPT and USER_TEMPLATE in sqale_agent.py are the exact strings used in training. The system prompt holds the task, the tool list and an execution-plan scaffold. The user turn states that the schema is not shown and gives the question.
Tool calls are bare JSON. After its <think> block the model writes one or more {"tool": "<name>", "args": {...}} objects with no prose and no code fences. Keep thinking enabled.
Return results as the next user turn, formatted as Tool result: followed by the result and the remaining budget, for example [4 tool turns left before you must answer].
Keep earlier reasoning in context. During training every earlier <think> block stayed in the context. The Qwen chat template drops those blocks when a conversation is rebuilt from a message list, so append each completion to the prompt string instead, as sqale_agent.py does.

Table with columns: Tool, Returns
Tool	Returns
`list_tables`	every table name
`describe_table(table_name)`	the table's `CREATE TABLE` statement
`foreign_keys(table_name)`	declared foreign keys as `child.col -> parent.col`; `table_name` is optional
`join_path(table_a, table_b)`	the shortest chain of foreign-key joins between two tables

Training

Environment. Each training question becomes a tool-use episode over a live SQLite database built from its schema and rows, and the schema is never placed in the prompt. A graph-aware sampler decides which tables the episode's database holds. It draws a target size between 4 and 64 tables, starts from the gold tables and grows along foreign keys until the target is reached, then shuffles the table order. The model therefore sees both small and large databases, and list_tables does not reveal which tables matter. Early in training, a submit gate refuses submit_sql until the same query has been run with run_query. The gate's probability decays from 1.0 to 0 over the first 80 steps, and it is never active at evaluation.

Reward. A query whose result set matches the gold result scores 3.0. A query that executes with a wrong result scores 1.0 + 0.9·F1 + b, where F1 is the multiset F1 against the gold result and b adds 0.1 times the share of joins along declared foreign keys and 0.1 for agreeing with the most common result in the group. A query that parses but does not execute scores 0.5, and anything else scores 0. The executable tier tops out at 2.1, so no partial credit ranks a wrong query above a correct one. Correctness is decided by comparing result sets, with no LLM judge.

Length curriculum. Half of SQaLe's gold queries are shorter than 150 characters, so a uniformly sampled batch mostly holds short questions. The curriculum buckets training questions by the length of their masked gold SQL at the 50th, 80th and 95th percentiles. The first 400 steps draw from the short bucket only. The mixture then moves linearly from (0.40, 0.30, 0.20, 0.10) to (0.15, 0.20, 0.25, 0.40) over 1,000 steps and holds there for the remaining 400. After the warm-up no bucket falls below 0.15, so short questions keep contributing gradient for the whole run.

Optimisation. GRPO from the Qwen3.5-2B base checkpoint, with 18 questions × 8 candidates = 144 episodes per update for 1,800 steps. This repository holds the last checkpoint. BIRD dev was monitored during training and not used to select the checkpoint.

Table with columns: Hyperparameter, Value
Hyperparameter	Value
Base model / initial policy	Qwen3.5-2B (base checkpoint)
Reference policy	base model, re-anchored every 111 updates
Training corpus	SQaLe, difficulty bands simple, moderate, hard
Schema in prompt	no (withheld; tools only)
Protocol	execution-plan scaffold, 7 tools + `submit_sql`
Tool turns / final turn	6 / 1
Submit gate probability	1.0 → 0.0 over the first 80 steps
Consecutive query-error cap

Intended use and limitations

MSQaLe answers natural-language questions over SQLite databases through tool access, including databases with hundreds of tables. It is a research model released with the paper.
It expects the agent loop described above. Putting the DDL in a single prompt and asking for SQL is a different setting from the one it was trained in.
It was trained and evaluated on SQLite only.
At 2B parameters it still trails much larger models. The untrained Qwen3.6-27B scores 76.0 / 69.3 / 55.0 on the three benchmarks above, and EHRSQL's clinical questions remain hard for every 2B model in the paper.
The model's SQL runs against your data, so give it read-only access, as sqale_agent.py does.
The checkpoint is a full Qwen3_5ForConditionalGeneration export and carries Qwen3.5-2B's vision tower unchanged. Only the language model was trained, and only text was evaluated.

Citation

If you use this model or the SQaLe dataset, please cite:

bibtex

@misc{wolff2026sqale,  title  = {{SQaLe}: A Large Realistic Dataset to Empower Small Specialised Text-to-{SQL} Models},  author = {Wolff, Cornelius and Gomm, Daniel and Hulsebos, Madelon},  year   = {2026}}

Authors: Cornelius Wolff and Daniel Gomm (University of Amsterdam, Centrum Wiskunde & Informatica), Madelon Hulsebos (Centrum Wiskunde & Informatica). Questions and feedback are welcome in the Community tab of this repository.

The model builds on Qwen3.5-2B. SQaLe's schemas come from SchemaPile, and the evaluation uses BIRD and EHRSQL. We thank their authors for making them available.

Available on FriendliAI

Dedicated Endpoints

Run this model inference on single tenant GPU with unmatched speed and reliability at scale.

Learn more

Model Details

Model Provider

trl-lab

Model Tree

Base

Qwen/Qwen3.5-2B

Fine-tuned

this model

Input Modalities

TextImageVideo

Output Modalities

Text

Supported Functionality