Text-to-SQL (BIRD Benchmark)

Text-to-SQL (BIRD Benchmark)

GRPO training for Qwen3-32B on the BIRD SQL benchmark, served by Arctic RL with the ZoRRo trainer. Mirrors the hyperparameters of the stock-verl baseline at verl_opensource/examples/bird_sql/run_qwen3_32b_bird_grpo.sh so the two backends can be compared apples-to-apples on wall-clock speed. Pure GRPO, without a frozen reference model (no KL anchoring).

Overview

• Model: Qwen/Qwen3-32B
• Topology: 4 nodes × 8 H200 GPUs (32 GPUs total), colocate=True (training + sampling share each GPU bundle), DeepSpeed ZeRO stage-3 with CPU optimizer offload, vLLM rollout (TP=2). Without KL there is no frozen reference model, so the ref log-prob pool is disabled (log_prob_gpus=0); under ZoRRo log-probs are recomputed through the training engine itself.
• Data: BIRD-SQL only — train on BIRD train.json, validate on BIRD dev.json. Spider / GretelAI are not used in this recipe.
• Reward: SQLite execution-match against the gold SQL (bird_reward.py)

Step 1: Ray and Multi-Node Hostfile

When using a multi-node training environment, Ray and DeepSpeed need a special file called hostfile (from MPI) to discover participating nodes and GPU counts.

Most CSPs provide this file. If you don't have one, create it with the following format:

10.1.1.1 slots=8
10.1.1.2 slots=8
10.1.1.3 slots=8
10.1.1.4 slots=8

The first column is node IPs; the second is the GPU count on each node.

Export the path so the install and launcher scripts can use it:

export JOB_HOSTFILE=/path/of/hostfile

Alternatively, edit the HOSTFILE setting at the top of the training scripts.

For a single-node run, you can skip the hostfile — the launcher falls back to NNODES=1 and uses the local 8 GPUs.

Step 2: Install Packages

Conda environments are node-local. Each node has its own ~/miniconda3, so the environment must be created and populated on every node in $JOB_HOSTFILE. Use a fresh, recipe-specific env name.

Bootstrap the Environment

Pick an env name and resolve the env's bin/ path (identical on every node):

export CONDA_ENV=txt2sql
CONDA_BASE=$(conda info --base)
ENV=$CONDA_BASE/envs/$CONDA_ENV/bin

Bootstrap the environment on the launching node:

conda create -y -n $CONDA_ENV python=3.12
$ENV/python -m pip install -q uv
$ENV/uv pip install --python $ENV/python deepspeed

Clone the repositories onto shared storage:

git clone https://github.com/Snowflake-AI-Research/Arctic-Platform
git clone -b arctic_rl_share_v0.7.1 --single-branch https://github.com/Snowflake-AI-Research/verl
cd Arctic-Platform/recipes/rl/verl/txt2sql

Distribute Environment to All Nodes

Create the environment on remaining nodes and install pinned dependencies on all nodes. CUDA 12.9 is assumed; adjust the torch index URL and cuda-bindings pin in requirements.txt if using a different CUDA version.

$ENV/ds_ssh -f $JOB_HOSTFILE "[ -x $ENV/python ] || $CONDA_BASE/bin/conda create -y -n $CONDA_ENV python=3.12"
$ENV/ds_ssh -f $JOB_HOSTFILE "$ENV/python -m pip install -q uv"

# Install torch (CUDA 12.9) first, then pinned packages
$ENV/ds_ssh -f $JOB_HOSTFILE "$ENV/uv pip install --python $ENV/python torch==2.10.0 --index-url https://download.pytorch.org/whl/cu129 -U"
$ENV/ds_ssh -f $JOB_HOSTFILE "$ENV/uv pip install --python $ENV/python -r $PWD/requirements.txt --override $PWD/overrides.txt"
$ENV/ds_ssh -f $JOB_HOSTFILE "$ENV/uv pip install --python $ENV/python -U pip wheel packaging setuptools"

Install Flash Attention

The launcher uses flash_attention_2 by default. Build it from source:

$ENV/ds_ssh -f $JOB_HOSTFILE "$ENV/uv pip install --python $ENV/python flash-attn --no-build-isolation"

Alternatively, download a prebuilt FA2 wheel from flash-attention releases.

For flash_attention_3 on Hopper GPUs (faster), install the matching flash_attn_3 wheel from flash-attention3-wheels, then enable it in the launcher by commenting/uncommenting the GPU auto-selection block.

Install verl

Install the Snowflake verl fork editable on all nodes:

cd ../../../../../verl
grep -v flash-attn requirements.txt > requirements-no-fa.txt
$ENV/ds_ssh -f $JOB_HOSTFILE "cd $PWD && $ENV/uv pip install --python $ENV/python -r requirements-no-fa.txt && $ENV/uv pip install --python $ENV/python -e ."
cd -

For a single-node run, skip ds_ssh: create the env, bootstrap uv, and run the $ENV/uv pip install commands directly on the local node.

Step 3: Get Raw BIRD Data

BIRD-SQL provides train.json, dev.json, SQLite files, and per-table database_description/*.csv files. Since the official endpoint is slow/unreliable, pull from the Sudnya/bird-sql HuggingFace mirror (train: ~9,428 rows; dev: ~1,534 rows).

Fetch onto shared storage visible to all nodes. The train_databases.zip archive is ~20 GB:

export BIRD_DIR=/data/bird
$ENV/python - <<'PY'
import os, glob, shutil, tempfile, zipfile
import pandas as pd
from huggingface_hub import hf_hub_download
import json

repo, bird = "Sudnya/bird-sql", os.environ["BIRD_DIR"]

def dump_questions(member, out_json):
    df = pd.read_parquet(hf_hub_download(repo, member, repo_type="dataset"))
    rows = [{"db_id": r["db_id"], "question": r["question"],
             "evidence": ("" if pd.isna(r.get("evidence")) else r["evidence"]), "SQL": r["SQL"]}
            for _, r in df.iterrows()]
    os.makedirs(os.path.dirname(out_json), exist_ok=True)
    json.dump(rows, open(out_json, "w"))
    print(f"wrote {len(rows)} rows -> {out_json}")

def fetch_dbs(member, dest_parent, dbname):
    z = hf_hub_download(repo, member, repo_type="dataset")
    tmp = tempfile.mkdtemp(dir=dest_parent)
    with zipfile.ZipFile(z) as zf: zf.extractall(tmp)
    root = os.path.dirname(os.path.dirname(glob.glob(f"{tmp}/**/*.sqlite", recursive=True)[0]))
    os.replace(root, os.path.join(dest_parent, dbname))
    print(f"ready -> {os.path.join(dest_parent, dbname)}")

dump_questions("data/train-00000-of-00001.parquet", f"{bird}/train/train.json")
dump_questions("data/validation-00000-of-00001.parquet", f"{bird}/dev/dev.json")
fetch_dbs("databases/train_databases.zip", f"{bird}/train", "train_databases")
fetch_dbs("databases/dev_databases.zip", f"{bird}/dev", "dev_databases")
PY

You should end up with:

/data/bird/
├── train/
│   ├── train.json
│   └── train_databases/<db_id>/<db_id>.sqlite     (+ database_description/*.csv)
└── dev/
    ├── dev.json
    └── dev_databases/<db_id>/<db_id>.sqlite       (+ database_description/*.csv)

Step 4: Preprocess to verl Parquets

preprocess_bird.py converts raw BIRD data into verl-compatible parquets with augmented training prompts and clean validation prompts.

Train: Extended Prompts

For each BIRD train.json row, a prompt is built in the arctic_text_to_sql_r1 format (system prompt + user message with <think> and <answer> tags, answer containing fenced sql block). The user message contains the database schema as CREATE TABLE DDL extracted from the SQLite file, augmented with:

1. Foreign-key summary line at the top of each table block
2. Sample-rows block as a markdown-style pipe table (10 rows by default)
3. Per-column example values as inline -- example: [...] comments (10 distinct values per column)
4. BIRD database_description/*.csv enrichment appended to column comments (name, desc, values)
5. Evidence text from the BIRD row appended to the question when present

A Qwen3-tokenizer length filter drops samples whose tokenized prompt exceeds the cap (default 32,768 tokens). This yields approximately ~8.6 k train rows out of ~9.4 k raw BIRD rows.

Val: Clean Prompts

For BIRD dev.json, augmentations are deliberately disabled to avoid leaking auxiliary context:

• No database_description enrichment (use_descriptions=False)
• No sample-rows table (sample_rows=0)
• No FK summary line (include_fk_summary=False)
• Only 3 inline -- example: [...] values per column (vs. 10 in train)
• No token-length filter

This produces approximately ~1.5 k val rows.

Run Preprocessing

From the recipe directory, using the installed environment:

$ENV/python preprocess_bird.py \
    --bird_dir /data/bird \
    --output_dir /data/snowflakesql/txt2sql \
    --max_tokens 32768 \
    --tokenizer Qwen/Qwen3-1.7B \
    --num_examples 10 \
    --sample_rows 10

Output structure:

/data/snowflakesql/txt2sql/
├── train.parquet     # BIRD train, augmented, token-filtered  (~8.6k rows)
└── val.parquet       # BIRD dev, clean, no token filter        (~1.5k rows)

Each row follows verl's standard schema: data_source, prompt, ability ("sql"), reward_model (with ground_truth SQL), and extra_info (including db_id, db_path, question, split, index). The extra_info.db_path is opened by bird_reward.py at training time, so SQLite files must remain at their absolute paths on all nodes.

Optional flags for ablating augmentations:

• --no_descriptions - disable BIRD enrichment
• --no_fk_summary - disable foreign-key lines
• --num_examples 0 - disable per-column examples
• --sample_rows 0 - disable sample rows
• --sources bird spider gretelai - include additional sources (not used by this recipe)

Step 5: Train

Start Ray Cluster

Start the Ray cluster across all nodes (environment is now installed everywhere):

bash ./restart_multi_ray.sh

Configure and Launch Training

Edit environment variables in run_qwen3_32b_bird_grpo_arl_zorro_yes.sh to match your setup:

• HF_HOME - HuggingFace hub cache location (optional)
• VLLM_CACHE_ROOT - vLLM cache directory

The script defaults DATA_DIR to /data/snowflakesql/txt2sql. If you kept that path, launch with no overrides:

bash run_qwen3_32b_bird_grpo_arl_zorro_yes.sh

Otherwise, override the data paths inline:

bash run_qwen3_32b_bird_grpo_arl_zorro_yes.sh \
    data.train_files=/data/snowflakesql/txt2sql/train.parquet \
    data.val_files=/data/snowflakesql/txt2sql/val.parquet

Or set DATA_DIR in the environment before running the script.

Reward Scoring

The SQL reward is scored by bird_reward.py (supplied with this recipe and auto-wired via custom_reward_function). It executes the model's predicted SQL against the row's SQLite database and compares the result set to the gold query. No additional setup is required.

Files