Skip to content

Commit b529b08

Browse files
JuliaEdomclaude
JuliaEdom
and
claude
committed
feat(autotune): QLoRA training track — local_vllm provider, dataset builder, GPU runners
Distillation track from plan_autotune S1-S3: BIRD train gold -> QLoRA fine-tune of Qwen2.5-Coder-7B, served back through the existing eval harness. Free GPU only (Kaggle T4), so the runners are notebooks rather than a rented box. Verified on hardware: the smoke run completed 5/5 optimizer steps at 138 s/step, which also measured the T4 throughput the checkpoint interval is now sized against. Two hardware-only failure classes are fixed here and pinned so they cannot return: - trl renames config fields between releases (max_seq_length -> max_length), and a stale name is a TypeError that kills the session minutes in. accepted_kwargs() reconciles the config against the real signature, drops unknown options loudly, and refuses to start if a required one (fp16/bf16/seq/max_steps) went missing — silently training in the wrong dtype is worse than not training. Four tests cover every signature shape trl has shipped; both mutations of the helper kill them. - Eval gathered full logits (1 x 4096 x 152k vocab) and upcast them to fp32, a 6.22 GiB allocation that OOMs a 16GB T4 *after* training has already succeeded. prediction_loss_only plus a capped val slice: the verdict for this track is BIRD EA from the harness, eval loss is only a smoke signal. The runners smoke-test the whole path (config -> step -> evaluate -> adapter save) on five steps before committing a 6-12h session. Both earlier launches died in their first minutes and it cost a day each, because nobody checked the status. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 parent 6870e93 commit b529b08

16 files changed

Lines changed: 1142 additions & 2 deletions

‎pyproject.toml‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -125,3 +125,9 @@ ignore_missing_imports = true
125125
[[tool.mypy.overrides]]
126126
module = ["perplexity"]
127127
ignore_missing_imports = true
128+
129+
# GPU-only training deps for scripts/autotune/train_qlora.py: installed on the
130+
# rented GPU box only (requirements_gpu.txt), never on the dev machine or CI.
131+
[[tool.mypy.overrides]]
132+
module = ["unsloth.*", "trl.*", "datasets.*"]
133+
ignore_missing_imports = true

‎scripts/autotune/__init__.py‎

Whitespace-only changes.

‎scripts/autotune/build_dataset.py‎

Lines changed: 169 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,169 @@
1+
"""Build the QLoRA train/val set for plan_autotune S1.
2+
3+
Renders each BIRD-train row through the SAME template the product pipeline
4+
uses (`generate_sql.txt` via `load_prompt`), with the few-shot block rendered
5+
exactly as config E renders it at `--fewshot-top-k 0` — the literal string
6+
"(none)". The completion is the strict JSON contract the pipeline parses
7+
(`parse_generate_sql_output`), not bare SQL: a student trained on bare SQL
8+
would fail the product parser on every question.
9+
10+
Known train/inference gap (accepted for v1, see plan_autotune.md S6): here
11+
`schema_block` is the raw DDL shipped in `bird_train.parquet`, while at eval
12+
time config E renders retrieval-selected table cards with sampled values.
13+
Closing that gap needs the 33GB BIRD train databases — gated.
14+
15+
Usage (from repo root):
16+
.venv/Scripts/python.exe scripts/autotune/build_dataset.py
17+
"""
18+
19+
from __future__ import annotations
20+
21+
import json
22+
import random
23+
from pathlib import Path
24+
from typing import Any
25+
26+
import pandas as pd
27+
import sqlglot
28+
from sqlglot import expressions as sqlglot_exp
29+
30+
from nl_sql.agent.prompts import load_prompt
31+
32+
ROOT = Path(__file__).resolve().parents[2]
33+
PARQUET = ROOT / "data" / "bird_train.parquet"
34+
OUT_DIR = ROOT / "data" / "autotune"
35+
DEV_DB_DIR = ROOT / "data" / "bird_mini_dev" / "MINIDEV" / "dev_databases"
36+
37+
VAL_SIZE = 500
38+
SEED = 0
39+
MAX_EXAMPLE_CHARS = 32_000
40+
DIALECT = "sqlite"
41+
# The pipeline never scores rationale/confidence; fixed filler keeps the
42+
# completion shaped like the contract without inventing per-row prose.
43+
FIXED_CONFIDENCE = 0.9
44+
45+
46+
def compose_question(question: str, evidence: str) -> str:
47+
"""Mirror eval.runner's evidence embedding: `<question>\\n\\nHint: <evidence>`."""
48+
if not evidence or not evidence.strip():
49+
return question
50+
return f"{question}\n\nHint: {evidence}"
51+
52+
53+
def extract_tables(sql: str) -> list[str]:
54+
"""Distinct table names in order of appearance; [] when unparseable."""
55+
try:
56+
tree = sqlglot.parse_one(sql, dialect=DIALECT)
57+
except sqlglot.errors.ParseError:
58+
return []
59+
seen: list[str] = []
60+
for table in tree.find_all(sqlglot_exp.Table):
61+
name = table.name
62+
if name and name not in seen:
63+
seen.append(name)
64+
return seen
65+
66+
67+
def build_example(row: Any) -> dict[str, str]:
68+
question = compose_question(str(row.question), str(row.evidence or ""))
69+
prompt = load_prompt(
70+
"generate_sql",
71+
dialect=DIALECT,
72+
schema_block=str(row.schema),
73+
fewshot_block="(none)", # exact render of an empty few-shot context
74+
question=question,
75+
)
76+
sql = str(row.SQL).strip().rstrip(";").strip()
77+
tables = extract_tables(sql)
78+
completion = json.dumps(
79+
{
80+
"sql": sql,
81+
"rationale": f"Uses {', '.join(tables)}." if tables else "Direct query.",
82+
"tables_used": tables,
83+
"confidence": FIXED_CONFIDENCE,
84+
},
85+
ensure_ascii=False,
86+
)
87+
return {
88+
"db_id": str(row.db_id),
89+
"question": str(row.question),
90+
"prompt": prompt,
91+
"completion": completion,
92+
}
93+
94+
95+
def sql_parses(sql: str) -> bool:
96+
try:
97+
sqlglot.parse_one(sql, dialect=DIALECT)
98+
except sqlglot.errors.ParseError:
99+
return False
100+
return True
101+
102+
103+
def write_jsonl(path: Path, rows: list[dict[str, str]]) -> None:
104+
with path.open("w", encoding="utf-8") as fh:
105+
for row in rows:
106+
fh.write(json.dumps(row, ensure_ascii=False) + "\n")
107+
108+
109+
def main() -> None:
110+
df = pd.read_parquet(PARQUET)
111+
total = len(df)
112+
113+
train_db_ids = {str(x) for x in df["db_id"].unique()}
114+
dev_db_ids = {p.name for p in DEV_DB_DIR.iterdir() if p.is_dir()}
115+
overlap = train_db_ids & dev_db_ids
116+
assert not overlap, f"train/dev db_id leakage: {sorted(overlap)}"
117+
print(f"leakage check: {len(train_db_ids)} train dbs vs {len(dev_db_ids)} dev dbs — disjoint")
118+
119+
indices = list(range(total))
120+
random.Random(SEED).shuffle(indices)
121+
val_idx = set(indices[:VAL_SIZE])
122+
123+
val_rows: list[dict[str, str]] = []
124+
train_rows: list[dict[str, str]] = []
125+
dropped_parse = 0
126+
dropped_dupe = 0
127+
dropped_long = 0
128+
seen_pairs: set[tuple[str, str]] = set()
129+
130+
for pos, row in enumerate(df.itertuples(index=False)):
131+
example = build_example(row)
132+
if pos in val_idx: # cut BEFORE filters, per plan
133+
val_rows.append(example)
134+
continue
135+
sql = str(row.SQL)
136+
if not sql_parses(sql):
137+
dropped_parse += 1
138+
continue
139+
pair = (str(row.question), sql)
140+
if pair in seen_pairs:
141+
dropped_dupe += 1
142+
continue
143+
seen_pairs.add(pair)
144+
if len(example["prompt"]) + len(example["completion"]) > MAX_EXAMPLE_CHARS:
145+
dropped_long += 1
146+
continue
147+
train_rows.append(example)
148+
149+
OUT_DIR.mkdir(parents=True, exist_ok=True)
150+
write_jsonl(OUT_DIR / "train.jsonl", train_rows)
151+
write_jsonl(OUT_DIR / "val.jsonl", val_rows)
152+
153+
print(f"total rows: {total}")
154+
print(f"val (pre-filter): {len(val_rows)}")
155+
print(f"train kept: {len(train_rows)}")
156+
print(f"dropped parse: {dropped_parse}")
157+
print(f"dropped dupes: {dropped_dupe}")
158+
print(f"dropped >chars: {dropped_long} (cap {MAX_EXAMPLE_CHARS})")
159+
160+
sample = random.Random(SEED).sample(train_rows, 3)
161+
for i, ex in enumerate(sample):
162+
print(f"\n--- sample {i} [{ex['db_id']}] prompt head ---")
163+
print(ex["prompt"][:400])
164+
print(f"--- sample {i} completion ---")
165+
print(ex["completion"][:400])
166+
167+
168+
if __name__ == "__main__":
169+
main()

‎scripts/autotune/colab/colab_autotune.ipynb‎

Lines changed: 184 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,184 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "markdown",
5+
"metadata": {},
6+
"source": [
7+
"# NL_SQL autotune - Colab runner (duplicate of the Kaggle S3 run)\n",
8+
"\n",
9+
"Second shot at the same QLoRA training, so a lost Kaggle session does not cost another day.\n",
10+
"Same `train_qlora.py`, same dataset - both are pulled from the Kaggle dataset\n",
11+
"`liovinajo/nlsql-autotune-data`, which is the single source of truth for this track.\n",
12+
"\n",
13+
"**Colab differs from Kaggle in one way that matters:** `/content` is wiped on every\n",
14+
"disconnect, and the free tier disconnects often. So checkpoints live on Google Drive and the\n",
15+
"run resumes from them. Re-running this notebook after a drop continues where it stopped.\n",
16+
"\n",
17+
"Steps: Runtime -> Change runtime type -> **T4 GPU**, then Run all. Two prompts appear:\n",
18+
"Drive authorisation, and an upload box for `kaggle.json` (Kaggle -> Settings -> Create New\n",
19+
"API Token). Nothing else is interactive."
20+
]
21+
},
22+
{
23+
"cell_type": "code",
24+
"metadata": {},
25+
"source": [
26+
"BASE_MODEL = \"Qwen/Qwen2.5-Coder-7B-Instruct\"\n",
27+
"KAGGLE_DATASET = \"liovinajo/nlsql-autotune-data\"\n",
28+
"\n",
29+
"DRIVE_DIR = \"/content/drive/MyDrive/nlsql_autotune\" # checkpoints + adapter survive here\n",
30+
"DATA_DIR = \"/content/nlsql_data\" # re-downloaded per session, never on Drive\n",
31+
"OUT = DRIVE_DIR + \"/qlora_out\"\n",
32+
"\n",
33+
"EPOCHS = 1.0\n",
34+
"MAX_SEQ = 4096\n",
35+
"SMOKE_STEPS = 5 # >0: prove the whole path on N steps first (0 = straight into the real run)"
36+
],
37+
"execution_count": null,
38+
"outputs": []
39+
},
40+
{
41+
"cell_type": "code",
42+
"metadata": {},
43+
"source": [
44+
"import subprocess\n",
45+
"import sys\n",
46+
"from pathlib import Path\n",
47+
"\n",
48+
"from google.colab import drive, files\n",
49+
"\n",
50+
"print(subprocess.run([\"nvidia-smi\"], capture_output=True, text=True).stdout, flush=True)\n",
51+
"\n",
52+
"drive.mount(\"/content/drive\")\n",
53+
"Path(DRIVE_DIR).mkdir(parents=True, exist_ok=True)\n",
54+
"\n",
55+
"# kaggle.json: reuse the copy on Drive if a previous session left one, else ask once.\n",
56+
"cred = Path(\"/root/.kaggle/kaggle.json\")\n",
57+
"cred.parent.mkdir(parents=True, exist_ok=True)\n",
58+
"drive_cred = Path(DRIVE_DIR) / \"kaggle.json\"\n",
59+
"if drive_cred.exists():\n",
60+
" cred.write_bytes(drive_cred.read_bytes())\n",
61+
" print(\"kaggle.json taken from Drive\", flush=True)\n",
62+
"else:\n",
63+
" print(\"upload kaggle.json (Kaggle -> Settings -> Create New API Token)\", flush=True)\n",
64+
" for name, blob in files.upload().items():\n",
65+
" if name.endswith(\".json\"):\n",
66+
" cred.write_bytes(blob)\n",
67+
" drive_cred.write_bytes(blob) # so the next session after a drop is unattended\n",
68+
"cred.chmod(0o600)\n",
69+
"\n",
70+
"subprocess.run([sys.executable, \"-m\", \"pip\", \"install\", \"-q\", \"kaggle\"], check=True)\n",
71+
"Path(DATA_DIR).mkdir(parents=True, exist_ok=True)\n",
72+
"subprocess.run(\n",
73+
" [\"kaggle\", \"datasets\", \"download\", \"-d\", KAGGLE_DATASET, \"-p\", DATA_DIR, \"--unzip\"],\n",
74+
" check=True,\n",
75+
")\n",
76+
"print(sorted(p.name for p in Path(DATA_DIR).iterdir()), flush=True)"
77+
],
78+
"execution_count": null,
79+
"outputs": []
80+
},
81+
{
82+
"cell_type": "code",
83+
"metadata": {},
84+
"source": [
85+
"import subprocess\n",
86+
"import sys\n",
87+
"\n",
88+
"subprocess.run([sys.executable, \"-m\", \"pip\", \"install\", \"-q\", \"unsloth\"], check=True)\n",
89+
"subprocess.run(\n",
90+
" [\n",
91+
" sys.executable,\n",
92+
" \"-m\",\n",
93+
" \"pip\",\n",
94+
" \"install\",\n",
95+
" \"-q\",\n",
96+
" \"datasets>=3.0\",\n",
97+
" \"trl>=0.12\",\n",
98+
" \"peft>=0.13\",\n",
99+
" \"bitsandbytes>=0.44\",\n",
100+
" \"accelerate>=1.0\",\n",
101+
" ],\n",
102+
" check=True,\n",
103+
")"
104+
],
105+
"execution_count": null,
106+
"outputs": []
107+
},
108+
{
109+
"cell_type": "code",
110+
"metadata": {},
111+
"source": [
112+
"import shutil\n",
113+
"import subprocess\n",
114+
"import sys\n",
115+
"from pathlib import Path\n",
116+
"\n",
117+
"data_dir = Path(DATA_DIR)\n",
118+
"script = data_dir / \"train_qlora.py\"\n",
119+
"if not script.exists():\n",
120+
" raise RuntimeError(\"train_qlora.py missing from the downloaded dataset\")\n",
121+
"\n",
122+
"\n",
123+
"def train_cmd(out, extra):\n",
124+
" return [\n",
125+
" sys.executable,\n",
126+
" str(script),\n",
127+
" \"--model\",\n",
128+
" BASE_MODEL,\n",
129+
" \"--train\",\n",
130+
" str(data_dir / \"train.jsonl\"),\n",
131+
" \"--val\",\n",
132+
" str(data_dir / \"val.jsonl\"),\n",
133+
" \"--out\",\n",
134+
" out,\n",
135+
" \"--max-seq-len\",\n",
136+
" str(MAX_SEQ),\n",
137+
" \"--batch-size\",\n",
138+
" \"1\",\n",
139+
" \"--grad-accum\",\n",
140+
" \"16\",\n",
141+
" \"--no-merge\",\n",
142+
" *extra,\n",
143+
" ]\n",
144+
"\n",
145+
"\n",
146+
"if SMOKE_STEPS and not Path(OUT, \"checkpoints\").exists():\n",
147+
" # Only on a fresh start - after a disconnect the path is already proven and the\n",
148+
" # session time is better spent resuming. Smoke writes to /content, not Drive.\n",
149+
" smoke_out = \"/content/smoke_out\"\n",
150+
" smoke = train_cmd(smoke_out, [\"--epochs\", \"1\", \"--max-steps\", str(SMOKE_STEPS)])\n",
151+
" print(\" \".join(smoke), flush=True)\n",
152+
" subprocess.run(smoke, check=True)\n",
153+
" if not Path(smoke_out, \"adapter\", \"adapter_config.json\").exists():\n",
154+
" raise RuntimeError(\"smoke run finished but saved no adapter\")\n",
155+
" shutil.rmtree(smoke_out)\n",
156+
" print(\"SMOKE OK -- starting the real run\", flush=True)\n",
157+
"\n",
158+
"cmd = train_cmd(OUT, [\"--epochs\", str(EPOCHS), \"--resume\"])\n",
159+
"print(\" \".join(cmd), flush=True)\n",
160+
"subprocess.run(cmd, check=True)\n",
161+
"shutil.make_archive(DRIVE_DIR + \"/adapter\", \"zip\", OUT, \"adapter\")\n",
162+
"print(\"DONE: \" + DRIVE_DIR + \"/adapter.zip\", flush=True)"
163+
],
164+
"execution_count": null,
165+
"outputs": []
166+
}
167+
],
168+
"metadata": {
169+
"accelerator": "GPU",
170+
"colab": {
171+
"provenance": [],
172+
"gpuType": "T4"
173+
},
174+
"kernelspec": {
175+
"display_name": "Python 3",
176+
"name": "python3"
177+
},
178+
"language_info": {
179+
"name": "python"
180+
}
181+
},
182+
"nbformat": 4,
183+
"nbformat_minor": 0
184+
}

0 commit comments

Comments
 (0)