LLMs & SQL-GRPO — from architecture to RL

Fine-Tuning Llama-2-7B with QLoRA — Interactive Tutorial

This notebook fine-tunes meta-llama/Llama-2-7b-hf on the CodeAlpaca-20k instruction dataset using 4-bit quantization (QLoRA) + LoRA adapters, and demonstrates every stage of the pipeline with live examples:

1. Environment / cache setup
2. Load dataset + show a raw example
3. Tokenizer: vocabulary, byte-pair encoding, encode/decode
4. BitsAndBytesConfig: what 4-bit quantization does
5. Load the model, inspect the Llama-2 transformer architecture
6. Forward pass: text → tokens → logits → probabilities → token
7. Autoregressive generation
8. Base-model inference (BEFORE fine-tuning)
9. LoRA: which weights get updated
10. Prompt formatting + train/validation split
11. SFTConfig + training-metric math
12. SFTTrainer internals
13. Training with train + validation loss each step
14. Overfitting check + save adapter
15. Inference comparison: before vs after fine-tuning

Prerequisites: the Llama-2 weights at /mnt/mydata/huggingface/Llama-2-7b-hf and the CodeAlpaca JSON, a CUDA GPU, and torch, transformers, peft, trl, bitsandbytes, datasets installed.

---

The pipeline at a glance

`` raw text ──tokenizer(BPE)──► token ids ──embedding──► vectors │ │ │ 32 × decoder blocks (self-attn + MLP) ▼ ▼ "### Instruction: ..." hidden states │ lm_head (4096→32000) ▼ logits ──softmax──► probabilities ▼ pick next token ──decode──► text `` Training reverses the last arrow at the loss step: given the true next token, compute cross-entropy on the logits and back-propagate — but with QLoRA only the tiny LoRA adapter matrices receive gradients.

Step 0 — Environment / HuggingFace cache setup

Point the HF libraries at the shared /mnt/mydata/huggingface tree where the Llama-2 weights and CodeAlpaca dataset live (direct HF downloads are blocked, so they were cloned with git lfs).

Step 1 — Load CodeAlpaca-20k and show one raw example

CodeAlpaca is an instruction-tuning dataset. Each row has three text fields: instruction (the task), input (optional extra context) and output (the desired answer).

Step 2 — Tokenizer: vocabulary, byte-pair encoding, encode ↔ decode

Is byte-pair encoding used? Yes. Llama-2 uses a SentencePiece tokenizer trained with the BPE algorithm.

• Vocabulary = 32,000 tokens; each token is a subword piece mapped to

an integer id 0..31999.

• How BPE turns text → tokens: starting from bytes/characters, a learned

ordered list of merge rules greedily glues the most frequent adjacent pairs into larger subwords. A leading space is encoded as the meta-symbol ▁ (U+2581).

• Special tokens: <s> (BOS, id 1), </s> (EOS, id 2), <unk> (id 0).

Llama-2 has no pad token, so we set pad_token = eos_token.

• Encode / decode are inverses: tokenizer(text) runs BPE → ids;

tokenizer.decode(ids) reassembles pieces (turning ▁ back to spaces).

Step 3 — BitsAndBytesConfig: 4-bit quantization (QLoRA)

The base Llama-2-7B weights ship as float16 (config torch_dtype='float16'), 2 bytes each → ~14 GB. BitsAndBytesConfig(load_in_4bit=True) loads each linear weight as 4-bit (~0.5 byte) → ~4 GB, so QLoRA fits on one GPU.

Key arguments & options

How NF4 conversion works: split each weight matrix into small blocks (e.g. 64), compute a per-block absmax scale, normalize to [−1,1], round each value to the nearest of 16 fixed NF4 levels (4 bits), store codes + scale. At compute time weights are de-quantized to bnb_4bit_compute_dtype. Weights are frozen — only LoRA adapters train. That is QLoRA.

Step 4 — Load the model & inspect the Llama-2 architecture

Is it decoder-only? Yes. Llama-2 is a decoder-only transformer (LlamaForCausalLM) — no encoder, no cross-attention, only masked (causal) self-attention, making it a left-to-right language model.

7B config: hidden_size=4096, num_hidden_layers=32, num_attention_heads=32, intermediate_size=11008, vocab_size=32000, max_position_embeddings=4096, hidden_act='silu' (SwiGLU), RMSNorm. Each block (pre-norm):

`` h = x + SelfAttention(RMSNorm(x)) # causal MHA with RoPE y = h + SwiGLU_MLP(RMSNorm(h)) # gate_proj, up_proj, down_proj `` Distinctive Llama choices: RMSNorm (not LayerNorm), RoPE rotary positions, SwiGLU MLP, and no biases.

Step 5 — Forward pass: logits → probabilities → token

model(input_ids) returns logits of shape (batch, seq_len, vocab_size). The last position's logits score every one of the 32,000 tokens as the next token. softmax turns those into a probability distribution; argmax (greedy) or sampling picks the next token; decode maps it back to text. The logit probabilities are available — they are exactly out.logits.

Step 6 — Autoregressive generation

Generation repeats the forward pass: logits → softmax → pick token → append → re-run … until EOS or max_new_tokens. model.generate(...) does this loop.

Step 7 — Base-model inference BEFORE fine-tuning

We define a helper that builds the Alpaca prompt, generates the model's answer, and returns only the text after ### Response:. We capture the base model's answers now (before LoRA) so we can compare after training.

Step 8 — LoRA: which weights get updated?

Does LoRA low-rank-update ALL weights? No. It freezes the original weight W ∈ ℝ^{d_out×d_in} and adds a trainable update only to the targeted linear layers:

$$W_\text{eff} = W\;(\text{frozen}) + \frac{\alpha}{r}\, B A,\quad A\in\mathbb{R}^{r\times d_\text{in}},\; B\in\mathbb{R}^{d_\text{out}\times r},\; r\ll d.$$

Only A and B train. Layers not in target_modules (embeddings, norms, k_proj/o_proj, the MLP) get no update.

LoraConfig args: r (rank, 8), lora_alpha (32, scale=α/r=4×), target_modules (q_proj,v_proj — classic minimal QLoRA), lora_dropout (0.05), bias ("none"), task_type ("CAUSAL_LM"). These are the lightweight QLoRA-paper recipe: trains in minutes, still adapts behavior.

Step 9 — Prompt formatting + train/validation split

Format the 3 raw fields into one Alpaca-style training string, take a small subset for speed, then split 90/10 so we can measure a validation loss.

Step 10 — SFTConfig + the training-metric math

Why SFTConfig? In TRL ≥ 0.12 it subclasses TrainingArguments and adds SFT fields (dataset_text_field, max_length, packing, loss_type), so it carries both the generic HF knobs and the SFT knobs in one object.

What the logged metrics mean — let N = non-masked target tokens, pᵢ(v) = softmax prob of vocab token v at position i, yᵢ = true next token:

• loss — per-token cross-entropy / NLL:

$$\text{loss} = -\frac{1}{N}\sum_i \log p_i(y_i)$$ average surprise of the correct token; loss=0 ⇔ p=1 on the truth; perplexity = exp(loss). This is what backprop minimizes.

• entropy — avg entropy of the model's predicted distribution:

$$\text{entropy} = \frac{1}{N}\sum_i\Big(-\sum_v p_i(v)\log p_i(v)\Big)$$ how uncertain/spread-out the prediction is (independent of the label); falls as training sharpens the distribution.

• mean_token_accuracy — teacher-forced top-1 accuracy:

$$\text{acc} = \frac{1}{N}\sum_i \mathbf{1}[\arg\max_v p_i(v)=y_i]$$

• num_tokens — cumulative count of non-masked training tokens (progress).
• grad_norm (for clipping), learning_rate (scheduler), epoch (fraction

of a pass). Evaluation reports the same with an eval_ prefix.

We set eval_steps = logging_steps = 5 so validation loss prints alongside train loss every step.

Step 11 — SFTTrainer: what it does internally

SFTTrainer subclasses transformers.Trainer. It (1) applies peft_config to wrap the base model with LoRA adapters; (2) tokenizes the text field and builds input_ids + labels (padding → -100, ignored in the loss); (3) builds a torch DataLoader + collator; (4) runs the standard PyTorch loop: model(inputs) → loss.backward() → optimizer.step()/zero_grad() (the paged_adamw_8bit optimizer) with an LR scheduler and bf16 autocast. It does not** reimplement autograd — it orchestrates PyTorch's autograd + optimizer.

Step 12 — Train, logging train loss + validation (eval) loss

Watch loss (train) and eval_loss (validation) each step. After training we print a clean step | train_loss | eval_loss table from log_history.

Step 13 — Overfitting check + save the LoRA adapter

If validation loss flattens/rises while train loss keeps falling → overfitting. Mitigate with more data, higher lora_dropout, lower r, fewer steps, weight decay, or early stopping on eval_loss. We save only the small adapter (a few MB), not the 13 GB base model.

Step 14 — Inference comparison: BEFORE vs AFTER fine-tuning

trainer.model = base + trained LoRA adapters, so generating from it reflects the fine-tuned behavior. We reuse the exact prompts whose base answers we captured in Step 7. Typically the fine-tuned model gives concise, clean, Alpaca-style code that stops cleanly, while the base model rambles/echoes headers.

---

Summary

You loaded Llama-2-7B in 4-bit (QLoRA), inspected the tokenizer/vocabulary/BPE and the decoder-only architecture, watched logits→probabilities→tokens→text, attached rank-8 LoRA adapters to q_proj/v_proj (≈0.12% of params trainable), trained with train+validation loss tracking, checked for overfitting, and saw a clear before/after improvement in instruction-following. See TUTORIAL.md for the extended write-up.

Rotary Position Embedding (RoPE) from Scratch in PyTorch

Llama-2 encodes token positions with Rotary Position Embedding (RoPE) rather than learned or sinusoidal absolute position embeddings. This notebook builds RoPE step by step in PyTorch, verifies its defining mathematical property, and checks our implementation against HuggingFace's Llama implementation.

What you'll do

1. Understand why absolute position embeddings are limiting and what RoPE fixes.
2. Derive the RoPE rotation and the frequency schedule θᵢ = base^(−2i/d).
3. Implement precompute_freqs, apply_rope (both the 2×2-rotation form and the

efficient rotate_half form).

1. Verify the key property: the dot product of a rotated query and key

depends only on their relative distance (m − n), not absolute positions.

1. Plug RoPE into a tiny attention layer.
2. Compare against 🤗 transformers Llama RoPE for numerical equivalence.

1. Why RoPE?

A transformer's self-attention is permutation-invariant — without position information it cannot tell "dog bites man" from "man bites dog." Ways to inject position:

• Absolute (learned or sinusoidal): add a position vector pᵢ to each token

embedding. Simple, but the model sees absolute indices and generalizes poorly to longer sequences; relative distances are only implicit.

• Relative position bias: add a learned scalar per (query, key) distance to

the attention logits. Captures relativity but needs extra parameters/lookups.

• RoPE (rotary): rotate the query and key vectors by an angle proportional

to their position inside the attention dot-product. The magic: after rotation, qₘ · kₙ becomes a function of only the relative offset m − n. No added parameters, plays nicely with the dot-product, and extrapolates well.

Llama uses RoPE, applied to the query and key projections of every attention head (not to the value, and not to the residual stream).

2. The math

Take a head of even dimension d. Split the vector into d/2 pairs of coordinates. Pair i (coordinates 2i, 2i+1) is treated as a 2-D point and rotated by angle mθᵢ, where m is the token's position and

$$\theta_i = \text{base}^{-2i/d}, \qquad i = 0,1,\dots,d/2-1,\quad \text{base}=10000.$$

Low i → large angle per step (fast rotation, encodes fine/local position); high i → tiny angle (slow rotation, encodes coarse/long-range position) — a multi-scale clock, analogous to sinusoidal encodings.

The 2-D rotation of pair i at position m:

$$\begin{pmatrix} x'_{2i}\\ x'_{2i+1}\end{pmatrix} = \begin{pmatrix}\cos m\theta_i & -\sin m\theta_i\\ \sin m\theta_i & \cos m\theta_i\end{pmatrix} \begin{pmatrix} x_{2i}\\ x_{2i+1}\end{pmatrix}.$$

Write R(m) for the block-diagonal matrix of all these 2×2 rotations. RoPE sets q̃ₘ = R(m) qₘ and k̃ₙ = R(n) kₙ. Because rotation matrices satisfy R(m)ᵀ R(n) = R(n − m), the attention score is

$$\tilde q_m \cdot \tilde k_n = q_m^\top R(m)^\top R(n)\, k_n = q_m^\top R(n-m)\, k_n,$$

which depends only on the relative offset n − m. That is the property we will verify numerically in Section 5.

3. Precompute the rotation angles (cos / sin tables)

For each position m ∈ [0, seq_len) and each pair index i ∈ [0, d/2) we need the angle mθᵢ. We build freqs[m, i] = m · θᵢ, then take cos and sin. We duplicate each column so the tables have the full head dim d (this pairs up with the rotate_half trick in Section 4).

4. Apply the rotation

Two equivalent formulations:

(a) Pairwise 2×2 rotation — literal, easy to read: for each pair (x_even, x_odd) compute x'_even = x_even·cosθ − x_odd·sinθ, x'_odd = x_even·sinθ + x_odd·cosθ.

(b) rotate_half (the Llama/HF trick) — vectorized and what real code uses. Split the vector into halves [x1 | x2] and define rotate_half(x) = [−x2 | x1]. Then

$$\text{RoPE}(x) = x \odot \cos + \text{rotate\_half}(x)\odot\sin,$$

with the cos/sin tables laid out as [θ | θ] (which is why we duplicated the columns above). This computes the same rotation for every pair simultaneously.

Sanity: does rotate_half match the literal pairwise rotation?

They use different coordinate pairings (HF pairs coordinate i with i+d/2; the textbook pairs 2i with 2i+1), so the raw vectors differ — but both are valid RoPE rotations that yield the same relative-position property. We confirm the HF form is a proper rotation by checking it preserves vector norm.

5. Verify the defining property: scores depend only on relative distance

This is the whole point of RoPE. Take a fixed query vector q and key vector k (same content), place the query at position m and the key at position n, apply RoPE, and compute the dot product ⟨RoPE(q, m), RoPE(k, n)⟩.

Claim: the score depends only on the offset n − m. So all (m, n) pairs with the same n − m (a diagonal of the score matrix) should be equal.

6. RoPE inside a tiny self-attention layer

Now use RoPE the way Llama does: project the input to Q, K, V per head, apply RoPE to Q and K only, then run scaled-dot-product causal attention.

7. Compare against 🤗 Transformers' Llama RoPE

HuggingFace exposes LlamaRotaryEmbedding (produces the cos/sin tables) and apply_rotary_pos_emb (the rotate_half application). We reproduce their cos/sin and rotation and check our apply_rope matches to floating-point tolerance. (If the import path differs across transformers versions, this cell degrades gracefully.)

8. (Optional) Visualize the rotation angles

Plot mθᵢ mod 2π as a heatmap over (position, pair index): low pair indices rotate fast (fine position), high indices rotate slowly (coarse position) — the multi-scale "clock."

---

Summary

• RoPE rotates each pair of query/key coordinates by mθᵢ, with a geometric

frequency schedule θᵢ = base^(−2i/d).

• The efficient implementation is x·cos + rotate_half(x)·sin.
• Its defining property — verified numerically in Section 5 — is that

⟨RoPE(q,m), RoPE(k,n)⟩ depends only on n − m, giving the transformer relative position awareness with no extra parameters and good length extrapolation.

• Our implementation matches 🤗 Transformers' Llama RoPE.

01 · Llama-2-7B architecture, end to end

This notebook opens the hood of Llama-2-7B and traces a piece of text all the way to a next-token probability, explaining every component you asked about:

1. Tokenizer — SentencePiece BPE, 32 000-token vocabulary.
2. Token embedding — a 32000 × 4096 lookup table.
3. Rotary Position Encoding (RoPE) — no learned/absolute position

embeddings; position is injected by rotating the query/key vectors.

1. Causal self-attention mask — each token may attend only to itself and

earlier tokens (autoregressive).

1. The decoder block — RMSNorm → multi-head self-attention (RoPE) →

residual → RMSNorm → SwiGLU MLP → residual, stacked 32×.

1. lm_head — a 4096 → 32000 linear that turns the final hidden state

into logits; softmax → probabilities → next token.

Key config (from config.json): hidden_size=4096, num_hidden_layers=32, num_attention_heads=32, num_key_value_heads=32 (so full multi-head attention, not grouped-query — that is a 70B/Llama-3 feature), intermediate_size=11008, vocab_size=32000, max_position_embeddings=4096, rms_norm_eps=1e-5, rope_theta=10000.

1 · Tokenizer — SentencePiece BPE

Llama-2 uses a SentencePiece tokenizer trained with the byte-pair encoding (BPE) algorithm. Text → a list of subword pieces → integer ids in 0..31999. A leading space is encoded as the meta symbol ▁ (U+2581). Special tokens: <s> = BOS (id 1), </s> = EOS (id 2), <unk> (id 0); Llama-2 ships with no pad token, so we reuse EOS for padding.

Why subwords? A fixed 32 k vocabulary can encode any string: frequent words become one token (▁users), rare words split into pieces (▁SQL → ▁SQ, L). This bounds the embedding table size while keeping sequences short. BPE learns the merge rules by greedily combining the most frequent adjacent pair, repeatedly, on the training corpus.

2 · Token embedding — the 32000 × 4096 lookup table

The first layer, model.embed_tokens, maps each id to a 4096-dim vector. There are 32000 × 4096 ≈ 131 M parameters here alone. Embedding is literally a row lookup: h0[t] = W_E[input_ids[t]].

3 · Rotary Position Encoding (RoPE)

Llama-2 has no absolute/learned position embeddings. Instead, inside each attention head it rotates the query and key vectors by an angle that depends on the token's position. For a 2-D pair of features (x_{2i}, x_{2i+1}) at position m with frequency θ_i = 10000^{-2i/head_dim}:

$$ \begin{pmatrix} x'_{2i}\\ x'_{2i+1}\end{pmatrix}= \begin{pmatrix}\cos m\theta_i & -\sin m\theta_i\\ \sin m\theta_i & \cos m\theta_i\end{pmatrix} \begin{pmatrix} x_{2i}\\ x_{2i+1}\end{pmatrix} $$

The crucial property: the attention score q_m · k_n becomes a function of the relative position m − n only. That gives translation-invariant, extrapolatable positional information without any extra parameters. Below we implement RoPE from scratch and verify the relative-position property.

4 · Causal self-attention mask

Self-attention computes, for each head, $\text{Attn}(Q,K,V)=\text{softmax}\!\big(\tfrac{QK^\top}{\sqrt{d}}+M\big)V$. For a causal (autoregressive) LM, the mask M is 0 on and below the diagonal and −∞ above it, so position t cannot see the future (t+1, t+2, …). This is the only attention masking Llama-2 uses (plus a padding mask to ignore pad tokens in a batch). There is no bidirectional / encoder-style full attention.

5 · The decoder block (×32)

Each of the 32 identical blocks is (pre-norm residual form):

`` h = h + Attn( RMSNorm(h) ) # RoPE inside Attn; causal mask h = h + SwiGLU( RMSNorm(h) ) # gated MLP ``

• RMSNorm normalizes by the root-mean-square of the activations (no mean

subtraction, no bias) — cheaper than LayerNorm: $\text{RMSNorm}(x)=\dfrac{x}{\sqrt{\text{mean}(x^2)+\epsilon}}\odot g$.

• SwiGLU MLP: down( silu(gate(x)) * up(x) ) with hidden width 11008.

The silu(gate)·up gating is what "SwiGLU" means; it outperforms a plain ReLU MLP at the same parameter budget.

Below we print the actual module tree and count parameters per block.

6 · lm_head: hidden state → logits → probabilities → token

After the last block and a final RMSNorm, lm_head (a 4096 → 32000 linear, untied from the embedding in Llama-2 since tie_word_embeddings=false) produces a logit for every vocabulary token at every position. softmax over the last position gives the next-token distribution; argmax (greedy) or sampling picks the token; decode turns it back into text. Repeating this is autoregressive generation.

Recap — the full path

`` "count users" ─BPE→ ids ─embed→ h0(seq×4096) │ ┌────────┴──────────── 32× decoder blocks ────────────┐ │ RMSNorm→ MHA(RoPE, causal mask) →+ RMSNorm→SwiGLU→+ │ └────────┬──────────────────────────────────────────────┘ final RMSNorm │ lm_head (4096→32000) logits ─softmax→ probs ─argmax/sample→ next id ─decode→ text ``

Everything above is decoder-only self-attention — no encoder, no cross-attention. Next: NB 02 turns the text-to-SQL datasets into exactly this prompt format so we can fine-tune the model to emit SQL instead of arbitrary text.

02 · Text-to-SQL data — load, format, tokenize, split

We fine-tune on two Hugging Face datasets (fetch with sql_grpo/download_datasets.sh):

Both share the same shape: (natural-language question, schema, gold SQL). That triple is exactly what we need to (a) build a prompt for SFT and (b) later build a runnable sqlite database for the execution reward in GRPO.

At least one dataset must be present locally; the loader raises if neither is found.

1 · Load and normalize to (question, schema, sql, seed)

Different datasets name their columns differently. We map each to the same dict so the rest of the pipeline is dataset-agnostic.

2 · The instruction prompt

build_prompt / build_example wrap each triple in the Llama-2 chat format with an [INST] ... [/INST] block, embed the schema, and put the gold SQL inside a fenced ``` `sql ... ` `` block terminated by </s>. Fencing makes it trivial to extract the query from a completion (the extract_sql` helper) for the reward function later.

3 · Tokenize and inspect lengths

The token-length histogram tells us a good max_seq_len — long enough to keep the answer, short enough to save memory.

4 · Completion-only masking (labels)

For SFT we do not penalize the model for failing to predict the prompt — only the SQL answer. So labels = input_ids but with prompt tokens set to -100 (PyTorch's cross-entropy ignore_index). Only tokens after [/INST] contribute to the loss. TRL's SFTTrainer can do this automatically with a DataCollatorForCompletionOnlyLM; here we show the mechanism explicitly.

5 · Train / validation split & save

We shuffle, hold out a validation slice for the eval-loss curve in NB 03, and persist the normalized rows so every downstream notebook reads the same data. A separate prompts-only file is what GRPO (NB 04) samples from.

6 · Sanity-check the execution sandbox

The reward in NB 04 runs the gold SQL against a sqlite DB built from the schema (+seed rows). Verify a couple of gold queries actually execute here so we know the reward signal is meaningful.

Note on synthetic schemas: gretelai schemas sometimes use Postgres-only types/functions sqlite rejects, and rarely ship seed rows. That is fine — in NB 04 the reward gives partial credit for a query that merely parses/executes and full credit when result sets match. Next: NB 03 fine-tunes Llama-2 on train.jsonl with QLoRA.

03 · Supervised fine-tuning (SFT) with QLoRA

Stage 1 of 2. Before any RL, we teach the base Llama-2 to imitate gold SQL with plain supervised learning on sql_data/train.jsonl (from NB 02).

The LLM fine-tuning loss

An autoregressive LM is trained with next-token cross-entropy (a.k.a. the causal-LM / negative-log-likelihood loss). For a target sequence $y_1..y_T$ and model $p_\theta$:

$$\mathcal{L}=-\frac{1}{|\mathcal{S}|}\sum_{t\in\mathcal{S}}\log p_\theta\!\left(y_t\mid y_{<t},\text{prompt}\right)$$

where $\mathcal{S}$ is the set of supervised positions — with completion-only masking (NB 02), that is only the SQL-answer tokens. Internally this is F.cross_entropy(logits[:-1], labels[1:], ignore_index=-100) (the shift lines up each position's logits with the next token).

Metrics

• train loss / eval loss — the cross-entropy above on train vs held-out val.
• perplexity $=e^{\text{loss}}$ — average branching factor; lower is better.
• token accuracy — fraction of answer tokens where argmax(logits) equals

the gold token (a quick readable proxy).

• Task-level (NB 04/05): execution accuracy — does the SQL run and return

the right rows.

QLoRA (why we can train 7B on one GPU)

• 4-bit NF4 quantization shrinks the frozen base weights from ~14 GB (fp16)

to ~4 GB.

• LoRA adapters: each targeted linear $W$ gets a low-rank update

$W + \frac{\alpha}{r}BA$ with $r\ll d$. Only $A,B$ (a few million params) get gradients; the 6.7B base stays frozen and quantized.

1 · Load the SFT data produced by NB 02

2 · Cross-entropy loss & perplexity — the mechanism, on real logits

Before a full run, let's see the loss: one forward pass on a formatted example, completion-only mask applied, then cross-entropy + perplexity + token accuracy by hand — exactly what the trainer optimizes each step.

3 · Attach LoRA adapters (the only trainable weights)

Target the attention projections (q,k,v,o_proj) and MLP projections (gate,up,down_proj). print_trainable_parameters shows the <1% that gets gradients.

4 · Inference snapshot BEFORE fine-tuning

Capture what the base model produces on a few held-out questions so we can compare against the fine-tuned model later. We also log the snapshot to TensorBoard as text.

5 · Train with SFTTrainer → TensorBoard (train + eval loss)

SFTConfig(report_to="tensorboard") streams train loss, eval loss, and learning rate to the TensorBoard event files under sft_sql_out/tb. A custom callback additionally logs eval perplexity and mirrors the scalars into our history dict for the inline plot. eval_strategy="steps" lets us watch for over-fitting (train ↓ while eval ↑).

6 · Inference snapshot AFTER fine-tuning + before/after comparison

Same held-out questions, now through the fine-tuned model — logged to TensorBoard and compared with the base-model execution accuracy.

📈 SFT levers — every hyperparameter and how it moves accuracy

The SFT quality caps what GRPO can achieve later (RL only sharpens an existing skill). Here is the full set of knobs in this notebook, what each does, and which way to turn it to raise text→SQL execution accuracy.

A. Data / sequence

B. LoRA capacity

C. Optimization

D. Evaluation / selection

In this notebook we changed: max_seq_length 512→1024, LoRA r 16→32 / alpha 32→64, and max_steps 120→400 — plus a schema-grounding aid in the prompt (build_prompt now lists the exact table(columns) and instructs the model to copy names verbatim). Together these target the specific failure you saw (invented/misspelled table & column names). Re-run this notebook to retrain the SFT adapter, then re-run NB 04 so GRPO starts from the stronger base.

---

Why SFT alone is not enough → RL

SFT maximizes token-likelihood of one gold query. But many SQL strings are equally correct (aliasing, join order, COUNT() vs COUNT(id)), and SFT gives no credit for a different-but-correct query and no penalty for a query that looks* like the gold text yet fails to execute. We need a reward tied to execution correctness — exactly what GRPO optimizes in NB 04.

04 · GRPO — reinforcement learning with an execution reward

Stage 2 of 2. We now optimize the SFT model with Group Relative Policy Optimization (GRPO) — the algorithm popularized by DeepSeek-R1/-Math and documented in Cameron Wolfe's write-up and Hugging Face TRL. The reward is does the generated SQL execute correctly?

GRPO in one page

For each prompt $q$, sample a group of $G$ completions $o_1,\dots,o_G$ from the current policy $\pi_\theta$. Score each with a reward $r_i$. GRPO's key idea: instead of a learned value/critic (as in PPO), use the group statistics as the baseline. The advantage of completion $i$ is the z-scored reward within its group:

$$A_i=\frac{r_i-\operatorname{mean}(r_1..r_G)}{\operatorname{std}(r_1..r_G)+\varepsilon}$$

A completion better than its group-mean gets a positive advantage (reinforced); worse than average gets negative (suppressed). No critic network needed.

The policy is updated with the PPO-style clipped surrogate, applied per-token (ratio $\rho_{i,t}=\pi_\theta/\pi_{\theta_{old}}$):

$$\mathcal{L}^{\text{GRPO}}=-\mathbb{E}\Big[\tfrac{1}{G}\sum_i \tfrac{1}{|o_i|}\sum_t \min\big(\rho_{i,t}A_i,\ \operatorname{clip}(\rho_{i,t},1-\epsilon,1+\epsilon)A_i\big) -\beta\,\mathrm{KL}\!\big(\pi_\theta\,\|\,\pi_{\text{ref}}\big)\Big]$$

• Clipping ($\epsilon\approx0.2$) stops any single update from moving the

policy too far — the PPO trust-region trick.

• KL penalty ($\beta$) to a frozen reference (the SFT model) keeps the

policy from drifting into gibberish that happens to game the reward.

Metrics to watch during GRPO

• mean reward (and execution accuracy) per step — should climb.
• reward std within groups — the learning signal; collapses toward 0 as the

policy masters a prompt.

• KL(π‖π_ref) — should stay small/bounded (not blow up).
• clip fraction — how often the ratio is clipped.
• completion length — watch for reward-hacking via degenerate outputs.

torch          : 2.3.0 | device: cuda | N_GPU: 1 | multi-GPU: False
Llama-2 weights: /mnt/mydata/huggingface/Llama-2-7b-hf
synthetic_sql  : /mnt/mydata/huggingface/datasets/synthetic_text_to_sql
sqale_sql      : /mnt/mydata/huggingface/datasets/sqale_text_to_sql

1 · The reward function = execution correctness

This is the heart of the method. Given a completion, we:

1. extract_sql the query,
2. build a sqlite DB from the schema (+seed rows),
3. run both the candidate and the gold query,
4. score:

Order-insensitive multiset comparison (rows_equal) avoids penalizing a different-but-equivalent ORDER BY.

reward=1.10  ok=True err=None  <= '```sql\nSELECT country, COUNT(*) FROM users GROUP BY country\n```'
reward=0.40  ok=True err=None  <= '```sql\nSELECT country FROM users\n```'
reward=0.00  ok=False err=OperationalError: incomplete input  <= 'SELECT nonsense FROM'

2 · Group-relative advantages — the core computation, visualized

Given a group of rewards, compute z-scored advantages. This cell has no model dependency, so it always runs and makes the math concrete.

mixed rewards=[1.0, 0.3, 0.0, 0.4]  advantages=[ 1.58 -0.34 -1.17 -0.07]
all good rewards=[1.0, 1.0, 1.0, 1.0]  advantages=[0. 0. 0. 0.]
one win rewards=[0.0, 0.0, 1.0, 0.0]  advantages=[-0.58 -0.58  1.73 -0.58]
saved nb04_advantages.png  (note: 'all good' -> ~0 advantage -> no gradient)

3 · The clipped surrogate + KL — a scalar demo

We plot the per-token GRPO objective as a function of the probability ratio $\rho$, for a positive and a negative advantage, to see the clipping.

saved nb04_clip.png  (flat regions = clipped: update is capped)

4 · Load the SFT policy + prompts, then run GRPO

GRPO is run with trl.GRPOTrainer — the single, production path. If TRL is not importable, the cell raises a clear ImportError with install guidance (there is no alternative loop).

RL prompts: 512
trl.GRPOTrainer available -> production path

loaded SFT adapter

4a · Run GRPO with trl.GRPOTrainer

Point the trainer at our sql_reward (wrapped to TRL's signature) and the prompts dataset. This is the code you'd scale up on the full dataset. A custom callback records the mean reward / KL that TRL logs so we can plot the curves below and mirror them to TensorBoard.

TensorBoard -> /workspace/notebooks/llms/sql_grpo/grpo_sql_out/tb
GRPOConfig: dropping kwargs not in this TRL: ['max_prompt_length']

The tokenizer has new PAD/BOS/EOS tokens that differ from the model config and generation config. The model config and generation config were aligned accordingly, being updated with the tokenizer's values. Updated tokens: {'pad_token_id': 2}.
`generation_config` default values have been modified to match model-specific defaults: {'max_length': 4096}. If this is not desired, please set these values explicitly.
Could not estimate the number of tokens of the input, floating-point operations will not be computed

<IPython.core.display.HTML object>

GRPO (TRL) done -> grpo_sql_out/adapter

saved nb04_grpo_curves.png
reward 0.463 -> 0.212

Inference snapshot AFTER GRPO (logged to TensorBoard)

Generate SQL for a few held-out prompts with the RL-tuned policy and log the text to TensorBoard so you can eyeball the qualitative gain over the SFT model.

The following generation flags are not valid and may be ignored: ['temperature', 'top_p']. Set `TRANSFORMERS_VERBOSITY=info` for more details.

[GRPO] Q: Reveal vessels with decreasing violations.
   pred: SELECT * FROM VesselInspectons WHERE ViolationCount = LAG(ViolationCount) OVER (PARTITION 
   exec_ok=False correct=False err=OperationalError: no such table: VesselInspectons
[GRPO] Q: How many defense diplomacy events were held in 2021 in the Asia-Pacifi
   pred: SELECT COUNT(*) FROM DefenseDiplomacies WHERE EventName LIKE '%Defense%'
   exec_ok=False correct=False err=OperationalError: no such table: DefenseDiplomacies
[GRPO] Q: How many hospitals are there in the rural region of "Southside"?
   pred: SELECT COUNT(*) FROM RuralHospital WHERE Location = 'SouthSide'
   exec_ok=False correct=False err=OperationalError: no such table: RuralHospital
[GRPO] Q: What is the average age of players who have played a game in the "Raci
   pred: SELECT p.platform, AVG(p.age) FROM Players p JOIN Games g ON p.favorite_game_Id = g.id WHE
   exec_ok=True correct=False err=None
logged 4 post-GRPO inference snapshots to TensorBoard

What we should observe

• mean reward and execution accuracy rise as the policy learns to emit

runnable, correct SQL.

• KL to reference grows slowly — the $\beta$ term (and clipping) keeps the

policy anchored to the fluent SFT model.

• advantage std provides the gradient; when a prompt is "solved" its group

rewards become uniform and its advantage → 0 (no further push).

Next, NB 05 closes the loop: when the SQL errors, feed the error text back into the prompt for a second attempt (test-time self-repair).

---

🐞 Lesson learned — degenerate [ [ [ ... output right after training

Symptom. Running the inference cell immediately after trainer.train() produced garbage greedy output like pred: [ [ [ [ [ [ ..., and every query failed to execute (unrecognized token). It looked like GRPO had collapsed the policy.

Investigation. We reloaded the saved adapters into a fresh 4-bit base model and generated on the same prompts:

• pre-GRPO (sft_sql_out/adapter) → clean SQL.
• post-GRPO (grpo_sql_out/adapter) → identical clean SQL.

The GRPO reward history was also healthy (rewards 0.13–1.01, KL ~1e-4), i.e. the learned weights were fine. The garbage only appeared when generating from the in-memory policy object used for training.

Root cause. After trainer.train(), TRL leaves the live model in its training-time state: .train() mode, gradient checkpointing enabled, and config.use_cache = False. Greedy .generate() on a QLoRA model in that state degenerates into a single repeated token. It is not a weights / reward-hacking problem — reloading the saved adapter (as NB 05 does) is fine.

Fix. Before generating from the just-trained model, restore inference state — policy.eval(), gradient_checkpointing_disable(), and config.use_cache = True (on both the PEFT wrapper and its base model) — and pass use_cache=True to .generate(). A light repetition_penalty / no_repeat_ngram_size is added as defense-in-depth. This is exactly what the updated inference cell above now does.

Takeaway. When a model looks "broken" only right after training, suspect the runtime state (train mode / cache / gradient checkpointing) before concluding the weights are bad — verify by reloading the saved checkpoint.

📓 Appendix — Debugging journal: making this notebook run in the container

This notebook originally failed to run inside the shared deeplearning-notebooks Docker image, and even once it ran, the final inference cell produced garbage. Below is the full sequence of bugs and the fixes, in the order they surfaced — a realistic RL-debugging trail.

The environment

The mismatch between a modern TRL and torch 2.3.0 is the source of the first three bugs.

---

Bug 1 — ImportError: cannot import name 'FSDPModule'

``text from trl import GRPOTrainer, GRPOConfig ImportError: cannot import name 'FSDPModule' from 'torch.distributed.fsdp' ``

Why. TRL >= 0.16 imports FSDPModule (the FSDP2 API) at module load. That symbol only exists in torch >= 2.4. On torch 2.3.0 the import blows up.

Fix (runtime shim, cell 1 bootstrap). Register a harmless stub so the import resolves; GRPO here is single-GPU and never uses FSDP2:

``python import torch.distributed.fsdp as _fsdp if not hasattr(_fsdp, "FSDPModule"): class _NoFSDPModule: pass _fsdp.FSDPModule = _NoFSDPModule ``

Fix (durable, image). requirements.txt now pins trl>=0.14,<0.16 — has GRPOTrainer (added in 0.14) but does not require FSDP2.

---

Bug 2 — GRPOConfig(max_prompt_length=...) -> TypeError

max_prompt_length was removed in newer TRL. Fix: keep only the kwargs the installed GRPOConfig accepts (via inspect.signature), and auto-pick the tokenizer arg name (tokenizer -> processing_class).

---

Bug 3 — generation_batch_size (2) must be divisible by num_generations (4)

Newer TRL requires per_device_train_batch_size % num_generations == 0. Fix: set the train batch equal to the group size and pin generation_batch_size when the field exists.

After Bugs 1–3 the notebook ran to completion.

---

Bug 4 — degenerate [ [ [ ... output from the post-GRPO model

The final inference cell printed pred: [ [ [ [ .... It looked like GRPO had collapsed the policy.

How we proved it wasn't. Reloading the saved adapters into a fresh 4-bit base model gave clean SQL for both pre- and post-GRPO checkpoints, and the reward history was healthy (rewards 0.13->1.01, KL ~1e-4). The garbage only came from the in-memory policy used for training.

Root cause. After trainer.train(), TRL leaves the model in training-time state: .train() mode, gradient checkpointing on, config.use_cache = False. Greedy .generate() in that state degenerates into a repeated token.

Fix. policy.eval(), gradient_checkpointing_disable(), config.use_cache = True (wrapper + base), and use_cache=True on .generate() (plus light repetition_penalty/no_repeat_ngram_size).

---

Cheat-sheet

Meta-lesson. Two failure classes: library-version drift (defend with capability detection + pins) and stateful objects (reset the model before inference; cross-check against a reloaded checkpoint).

🎯 RL for LLMs — pitfalls & how to get better convergence, lower loss, higher accuracy

GRPO/PPO-style RL on language models is finicky: the objective is non-stationary (the policy generates its own data), the reward is often sparse or gameable, and one bad hyperparameter can silently wreck a run.

---

1. Reward design pitfalls (the #1 source of failure)

Rule of thumb: if you can verify correctness programmatically (code, SQL, math), do that instead of a learned reward model.

---

2. Group / advantage pitfalls (GRPO-specific)

• Zero-variance groups -> zero gradient. If all G samples get the same

reward, std=0 => A_i=0 and that prompt gives no signal. Fix: raise num_generations (8–16), curriculum on medium-difficulty prompts, rollout temperature 0.7–1.0 for diversity.

• Too few generations -> high-variance advantage estimates.
• Advantage normalization — watch scale_rewards / loss_type in newer

TRL; mismatched scaling changes effective LR.

---

3. Stability pitfalls (KL, clipping, LR)

How the KL metric is computed

The kl value TRL logs (and what our GRPOHist callback records) is the mean per-token KL divergence between the current policy $\pi_\theta$ and the frozen reference $\pi_{\text{ref}}$ (the SFT model), averaged over the generated completion tokens and the batch.

Step by step, for one generated token $t$ of completion $o_i$:

1. During the rollout we have the reference and old-policy log-probs; on the

update we recompute $\log \pi_\theta(o_{i,t}\mid o_{i,<t}, q)$ for the same sampled token.

1. Form the log-ratio $\;s_{i,t} = \log\pi_{\text{ref}} - \log\pi_\theta$.
2. TRL uses the k3 unbiased, non-negative estimator (Schulman):

$$\mathrm{KL}_{i,t} = \exp(s_{i,t}) - s_{i,t} - 1 = \frac{\pi_{\text{ref}}}{\pi_\theta} - \log\frac{\pi_{\text{ref}}}{\pi_\theta} - 1 \ge 0.$$

1. Mask prompt tokens and padding (only completion tokens count).
2. Average over completion length and the group/batch to get the scalar

kl that is logged and plotted.

Equivalently, over sampled completions:

$$\widehat{\mathrm{KL}} = \frac{1}{\sum_i |o_i|}\sum_i\sum_{t=1}^{|o_i|} \Big(e^{s_{i,t}} - s_{i,t} - 1\Big),\qquad s_{i,t}=\log\frac{\pi_{\text{ref}}(o_{i,t}\mid\cdot)}{\pi_\theta(o_{i,t}\mid\cdot)}.$$

How it enters training. GRPO adds $\beta\cdot\mathrm{KL}_{i,t}$ as a per-token penalty to the loss (the $-\beta\,\mathrm{KL}$ term at the top of this notebook) — a larger beta pulls $\pi_\theta$ back toward the reference harder.

Practical notes.

• Masked + averaged, so kl is a small number (our demo ~1e-4); compare it

relatively across steps, not to an absolute threshold.

• The reference $\pi_{\text{ref}}$ is frozen (the SFT policy). With a LoRA

policy, TRL gets the reference distribution "for free" by disabling the adapter instead of holding a second full model in memory.

• Per-token estimator sketch:

``python s = ref_logp - pol_logp # log(pi_ref/pi_theta) kl = (s.exp() - s - 1) # k3, >= 0 kl = (kl * completion_mask).sum() / completion_mask.sum() # masked mean ``

• KL is your safety rail. KL climbing fast -> raise beta (0.04->0.1) or

lower LR. KL stuck at 0 with flat reward -> lower beta or raise LR.

• Clip fraction >0.3 -> steps too large; ~0 -> updates trivial.
• LR is much smaller than SFT: 1e-6–1e-5. Too high -> collapse in a few

steps; too low -> no movement.

• Reference refresh (sync_ref_model) keeps KL meaningful on long runs.
• Entropy. Collapsing entropy -> mode collapse (repeats); keep sampling

temperature up / add a small entropy bonus.

---

4. Optimization / systems pitfalls

• Start from a good SFT checkpoint — RL sharpens, rarely teaches from

scratch (this series does SFT in NB 03 first).

• Warmup + cosine/linear LR decay; flat high LR early causes divergence.
• Gradient clipping (max_grad_norm ~ 1.0).
• Effective batch size — larger num_generations x prompts (grad

accumulation) smooths reward curves.

• Prefer bf16 over fp16 (fp16 overflow -> silent NaNs in long rollouts).
• Train vs inference state (Bug 4): grad checkpointing + use_cache=False

are for training, not generation.

• Determinism for debugging — fix seeds; log prompt->completion->reward for

the first steps; most "mysterious" RL bugs are visible in the raw text.

---

5. Evaluation pitfalls

• Don't trust training reward — evaluate on a held-out set with

deployment decoding settings.

• Reload the checkpoint for eval (Bug 4).
• Watch length/format drift alongside accuracy.

---

6. A practical recipe for this task (text->SQL)

1. Strong SFT first (NB 03) until SFT execution accuracy plateaus.
2. num_generations = 8–16, rollout temperature 0.7–1.0.
3. LR 1e-6 -> 1e-5, warmup 3–5%, cosine decay, max_grad_norm=1.0.
4. beta 0.02–0.1, tuned by watching KL (slow bounded climb).
5. Reward = exec match (1.0) + valid-SQL partial (0.3) + format bonus −

length penalty; evaluate on held-out schemas.

1. bf16, gradient accumulation, hundreds–thousands of steps (not 20).
2. Then add NB 05's execution-feedback self-repair at inference.

📊 Evaluate after GRPO — train vs. validation execution accuracy

Why not "loss"? In policy-gradient RL the GRPO/PPO loss is not a quality metric (advantages are re-centered every step, so it can rise while the model improves). The meaningful post-training metrics are the reward and, for this verifiable task, execution accuracy — measured on a held-out validation split, not the prompts we trained on.

The cell below computes, with the current (post-GRPO) policy:

• train exec-accuracy on the RL prompts the policy saw, and
• validation exec-accuracy on sql_data/val.jsonl (unseen),

and logs both to TensorBoard (eval/train_exec_acc, eval/val_exec_acc). A big train≫val gap = overfitting/memorization; both low = underfit (train longer / stronger SFT base).

post-GRPO TRAIN exec-acc: 0.375  (runs-ok 0.688, n=16)
post-GRPO VAL   exec-acc: 0.312  (runs-ok 0.750, n=16)
train-val gap: +0.062  (ok)
saved nb04_eval_accuracy.png

Reading this chart. exec-correct = fraction whose rows exactly match gold; runs-without-error = fraction that at least execute. For a short demo run you'll typically see modest, similar train/val numbers (the model is underfit, not overfit). To move them up: stronger/longer SFT (NB 03) + more GRPO steps & num_generations — then this same cell should show both bars climbing, with the train-val gap staying small.

� Analysis + fix: no such table/column and an inline self-repair retry

After the runtime-state fix (Bug 4), the post-GRPO model emits well-formed SQL, but several queries still fail at execution:

``text pred: SELECT FROM VesselInspectons WHERE ... -> no such table: VesselInspectons pred: SELECT COUNT() FROM DefenseDiplomacies ... -> no such table: DefenseDiplomacies pred: SELECT COUNT(*) FROM RuralHospital ... -> no such table: RuralHospital pred: SELECT p.age AS AvgAge, ... -> no such column: p.age ``

What's happening (and what's not)

• This is not the degeneration bug — the syntax is valid SQL.
• The model is inventing near-miss identifiers: VesselInspectons (missing

an i), DefenseDiplomacies (pluralized), RuralHospital (singular), aliasing p.age without binding p. These are schema-grounding errors.

• Root cause: a 20-step GRPO demo barely moves the policy; it hasn't

learned to copy table/column names verbatim from the schema. With a real run (hundreds–thousands of steps, num_generations 8–16 — see the recipe above) these largely disappear.

The cheap, training-free fix: execution-feedback self-repair

We don't need more training to recover many of these. The sqlite error message (no such table: VesselInspectons) plus the schema is enough signal for the model to correct itself on a second attempt. This is exactly NB 05's idea, shown inline here so you can see the accuracy lift immediately:

1. generate SQL, execute it;
2. if it errors, append the error text (and re-show the schema) to the

prompt and ask for a corrected query;

1. repeat up to MAX_ROUNDS.

The cell below runs this loop on the same held-out prompts and reports first-try vs. after-repair execution accuracy.

Q: Reveal vessels with decreasing violations.
   first_try_correct=False after_repair_correct=False final_runs=False rounds=3
   final_pred: SELECT Vessel FROM VesselInspections2 WHERE ViolationCount = LAG(ViolationCount) OVER (PAR
Q: How many defense diplomacy events were held in 2021 in the A
   first_try_correct=False after_repair_correct=False final_runs=True rounds=1
   final_pred: SELECT COUNT(*) FROM DefenseDiplomacy WHERE EventName LIKE '%Defense%Diplomacy%' AND Year 
Q: How many hospitals are there in the rural region of "Southsi
   first_try_correct=True after_repair_correct=True final_runs=True rounds=1
   final_pred: SELECT COUNT(*) FROM RuralHospitals WHERE Location = 'Southside'
Q: What is the average age of players who have played a game in
   first_try_correct=False after_repair_correct=False final_runs=True rounds=1
   final_pred: SELECT p.platform, AVG(p.age) FROM Players p JOIN Games g ON p.favorite_game_id = g.id WHE
Q: What is the average time between maintenance for each type o
   first_try_correct=False after_repair_correct=False final_runs=False rounds=3
   final_pred: SELECT type, AVG(DATEDIFF('day', last_maintenance, next_maintenance)) FROM fleet GROUP BY 
Q: Identify suppliers who provide products with the highest and
   first_try_correct=False after_repair_correct=False final_runs=False rounds=3
   final_pred: SELECT SupplierName FROM Suppliers INNER JOIN NutritionScores ON Suppliers.SupplierID = Nu
Q: Which country had the highest average donation in 2021?
   first_try_correct=False after_repair_correct=False final_runs=False rounds=3
   final_pred: SELECT AVG(donation_amount) FROM donations WHERE YEAR(donation_date) = 2021 GROUP BY donor
Q: What's the total revenue for movies released between 2018 an
   first_try_correct=True after_repair_correct=True final_runs=True rounds=1
   final_pred: SELECT SUM(revenue) FROM movies WHERE release_year BETWEEN 2018 AND 2020
execution-correct: first-try 2/8 -> after self-repair 2/8
runs-without-error after repair: 4/8

Why this works. The failing query plus the concrete error (no such table: DefenseDiplomacies) turns an open-ended generation into a constrained correction — the model now has explicit evidence of which name was wrong and is nudged to copy the schema verbatim. Execution-feedback repair typically recovers a chunk of the first-try failures within 1–3 rounds, then plateaus.

Complementary, not a substitute. Self-repair is an inference-time patch; the durable fix is a longer GRPO run so the first-try output is already schema-grounded. In production you'd do both — train properly, then keep a 1–2 round repair loop as a safety net (that full loop, with round-by-round metrics, is NB 05).

05 · Execution-feedback loop — feed the error back into the prompt

Yes — the execution output can be fed back to improve the answer. Because SQL is executable, we get a free, high-quality critic: run the query and read sqlite's error message (or the returned rows). This notebook builds a self-repair / agentic loop:

`` ┌───────────────────────────────────────────────┐ ▼ │ question + schema ─► model ─► SQL ─► EXECUTE ─► error? ─yes─► append error ─┘ │ └─ no ─► return rows ✅ ``

On each round we append the previous (wrong) SQL and the exact sqlite error to the prompt and ask the model to fix it. This is the same signal GRPO used as a reward (NB 04); here we use it at inference time as an observation.

Two complementary uses of execution feedback

1. Training-time (NB 04): turn correctness into a scalar reward for GRPO.
2. Inference-time (this NB): turn the error text into extra prompt context

for a retry — no weight updates needed. Often combined: an RL model that was trained with the reward is also better at using the feedback text.

torch          : 2.3.0 | device: cuda | N_GPU: 1 | multi-GPU: False
Llama-2 weights: /mnt/mydata/huggingface/Llama-2-7b-hf
synthetic_sql  : /mnt/mydata/huggingface/datasets/synthetic_text_to_sql
sqale_sql      : /mnt/mydata/huggingface/datasets/sqale_text_to_sql

1 · The repair prompt

We extend the base prompt with a short "your previous attempt failed with this error, fix it" block. Keeping the schema in-context is essential — most errors are no such column/table, which the model can only fix by re-reading the schema.

<s>[INST] <<SYS>>
You are a helpful assistant that writes a single, correct SQLite query for the user's question given the schema. Use ONLY the table and column names exactly as they appear in the schema; do not invent, pluralize, or re-spell them.
<</SYS>>
### Schema (SQLite):
CREATE TABLE users(id INTEGER, country TEXT);
### Question:
Count users per country.
Your previous SQL was:
```sql
SELECT nation, COUNT(*) FROM users GROUP BY nation
```
It did not work. Error: OperationalError: no such column: nation
Return a CORRECTED SQL query only. [/INST]

2 · The feedback loop

solve_with_feedback tries up to max_rounds times, executing after each attempt and folding any error back into the next prompt. It returns the final SQL, whether it succeeded, and the per-round trace.

3 · Run the loop

With the real model this exercises Llama-2's ability to read an error and fix its query: it may start with a wrong column name and correct it once it "sees" the no such column error.

`torch_dtype` is deprecated! Use `dtype` instead!

The following generation flags are not valid and may be ignored: ['temperature', 'top_p']. Set `TRANSFORMERS_VERBOSITY=info` for more details.

loaded adapter: grpo_sql_out/adapter
round 0: correct=True ok=True err=None
   sql: SELECT country, COUNT(*) FROM users GROUP BY country
FINAL: SELECT country, COUNT(*) FROM users GROUP BY country | correct: True

4 · Does feedback help? pass@k vs repair@k

We compare single-shot execution accuracy to **accuracy after up to k repair rounds** over a batch of prompts. The gap is the value added by feeding execution output back into the prompt.

repair@k execution accuracy: {1: 0.5, 2: 0.5, 3: 0.5}
saved nb05_repair.png

Wrap-up

• SQL is verifiable: running it yields a correctness signal and a concrete

error message.

• Training-time, that signal is the GRPO reward (NB 04).
• Inference-time, the error text becomes extra prompt context that lets the

model self-repair (this NB) — typically lifting execution accuracy across the first 1–3 rounds before it plateaus.

• The two compound: a GRPO-tuned model both emits better first-shot SQL and

makes better use of the feedback.

That completes the pipeline: architecture (01) → data (02) → SFT (03) → GRPO (04) → execution feedback (05).

Production Training — Multi‑GPU SFT + GRPO via accelerate.notebook_launcher

This notebook trains a production Llama‑2‑7B text‑to‑SQL model on all your GPUs from inside the notebook, using accelerate.notebook_launcher.

Why notebook_launcher and not raw DDP? Native PyTorch DistributedDataParallel needs one OS process per GPU, spawned by torchrun/mp.spawn. That does not play well with a single long‑lived notebook kernel (zombie process groups, CUDA re‑init errors, hangs). notebook_launcher forks the worker processes correctly from the notebook, so multi‑GPU "just works" — which is exactly why it is the recommended path on Colab Pro/Pro+.

Prereqs: pip install "accelerate>=0.30" "trl>=0.9" peft bitsandbytes datasets optuna tensorboard, Llama‑2 weights at the model_id path, and the sql_data/ splits from 02_text_to_sql_data.ipynb.

1. Config + how many GPUs

2. Import the training entry points

We reuse the exact same functions the CLI scripts call, so the notebook and the accelerate launch path run identical code.

3. Launch SFT on all GPUs

notebook_launcher(fn, num_processes=NUM_GPUS) forks one worker per GPU. Each worker runs train_sft.main(), which builds the SFTTrainer; TRL + Accelerate shard the data and all‑reduce gradients across GPUs automatically. Only the main process writes checkpoints/TensorBoard.

4. Launch GRPO on all GPUs (starts from the SFT adapter)

5. TensorBoard — train & validation loss, reward, KL

Run in a terminal (or the magic below). Open http://localhost:6006 locally, or http://<server-ip>:6006 remotely (the --host 0.0.0.0 exposes it).

6. (Optional) Hyperparameter sweep

For a systematic search over lora_r, learning_rate, lora_dropout, warmup_ratio (minimizing eval loss), use the CLI sweeper — it launches each trial under accelerate so every trial is itself multi‑GPU:

``bash python production/sweep.py --trials 20 --gpus {NUM_GPUS} --max_rows 8000 --max_steps 300 ``

Then copy the best params into config.yaml (remember alpha = 2*r) and rerun sections 3–4 for the full training.

7. Checkpoints & resume

Checkpoints are written to production/runs/<run_name>/checkpoint-<step>/ (save_steps, keep last save_total_limit, best restored via load_best_model_at_end). To resume after a crash:

``bash accelerate launch --config_file production/accelerate_config.yaml \ production/train_sft.py --config production/config.yaml --resume_from_checkpoint auto ``

The final best adapter is at production/runs/<run_name>/adapter/.