A fictional application turns a short request into a SQLite SELECT query on a local inventory. The same model is trained with the same examples and adapters, first on a 16-bit base, then on a 4-bit quantized base. No produced SQL is executed with write permissions.
A request produces a proposed query, validated syntactically then evaluated read-only on several small inventories. Tables separate the effect of quantization without training from the effect of each adaptation.
Download the complete project, revision r002
The problem and its data
The supplied schema is inventory(sku TEXT, zone TEXT, qty INTEGER). A request specifies a zone, a quantity threshold and an output: SKU, SKU/quantity with explicit ordering, row count, sum, minimum or maximum. A request for price or supplier must return REFUSE because those fields do not exist. The model never receives the target query in its prompt.
StockSQL-v1 contains 300 families of eight phrasings, or 2,400 fictional requests. Families 0–209 provide 1,680 training examples, 210–254 provide 360 validation examples and 255–299 provide 360 test examples. SELECT type, zone and threshold define 300 distinct family SQL queries. Training zones are A to G; test uses I and J. Query types and language patterns remain shared. An identifier change alone is not a new family.
Among 360 test cases, 330 expect SQL and 30 a refusal. Five independent inventories are generated from fixed seeds: each has 300 SKUs in ten zones, with quantities from 0 to 20. Their contents are written to five real SQLite databases and reread to check correspondence with the original JSON rows. SHA-256 before and after evaluation checks their integrity.
The supervised signal
The tokenizer and chat template belong to Qwen2.5-1.5B-Instruct. The prompt contains the system instruction, schema, request, then the boundary of an assistant response. Prompt IDs come from apply_chat_template; target SQL or REFUSE is encoded separately, then EOS=151645 is added. This explicit contract is not presented as equivalent to retokenizing a concatenated string.
The prompt is limited to 256 tokens, the response to 128 including EOS, and the total to 384. An overflow raises an identified error, never silent truncation. Prompt and padding labels are −100. Other labels supervise every response token and EOS. PAD=151643 and EOS=151645 are distinct; attention and supervision nevertheless depend on real positions.
Cross-entropy aligns logits[:, :-1] and labels[:, 1:], then selects active positions and computes in FP32. The model pass uses BF16 autocast. The eight examples in an update can have different response lengths: their cross-entropy sums are divided by the same total number of response tokens. Averaging eight averages would change that weight. The test compares gradients with a reference calculation.
LoRA: learn a low-rank correction
For a projection W of output × input shape, the correction is (alpha/r) B A, with A of r × input shape and B of output × r shape. The layer computes xWᵀ + (alpha/r)xAᵀBᵀ. W remains frozen; only A and B are optimized. The small LoRALinear module shows this calculation, then compares output and input/A/B gradients with a real PEFT layer in FP32. B is non-zero in that comparison so its gradients are observable.
For actual training, A follows Kaiming initialization and B starts at zero. At that point the correction is zero and A's first gradient can be zero; this is not unintended freezing. Later steps must genuinely change A and B. Checks cover the real tensors and their hashes.
Adapters cover q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj and down_proj in 28 blocks: 196 projections and 392 A/B matrices. With r8 and alpha 16, they represent 9,232,384 FP32 parameters. With r16 and alpha 32, they represent 18,464,768 FP32 parameters. Safetensors files add metadata to the logical size of these values; their actual size is measured.
BF16 and NF4: what actually changes
The frozen base is Qwen/Qwen2.5-1.5B-Instruct, revision 989aa7980e4cf806f80c7fef2b1adb7bc71aa306. Width 1536, 28 blocks, 12 Q heads, two K/V heads, MLP 8960 and vocabulary 151936 are checked after loading. Sharing between input and output embeddings must be preserved.
In the BF16 branch, the 196 projection weights remain stored in BF16. In the NF4 branch, bitsandbytes quantizes them in NF4 with double quantization; Linear4bit classes, nested quantized state, storage and BF16 compute dtype are audited. Loading explicitly designates CUDA 0 inside the container, which is itself limited to the selected GPU. It does not use automatic placement.
prepare_model_for_kbit_training is called before PEFT installation. This preparation converts non-quantized parameters to FP32. To avoid an additional difference, the BF16 branch also keeps embeddings/tied output, norms and biases in FP32 while leaving its 196 projection weights in BF16. It would therefore be false to say the whole model uses four bits. Quantization metadata is retained.
PEFT can convert adapters according to loading dtype. After installation, we explicitly reinitialize A/B in FP32 with the same CPU generators by module name in both branches. Their initial hashes must be exactly equal at the same rank; a BF16-rounded initialization in just one branch would blur the comparison. Gradient checkpointing is disabled.
A comparable budget and verifiable resume
Both branches use seed 45, the same example order, two epochs, microbatch 1 and accumulation of eight examples: 420 updates and 3,360 presentations. AdamW uses lr=0.0001 and weight_decay=0.01; clipping at 1; 42 warm-up steps then linear decay to 10% of the maximum rate. Only FP32 adapters enter the optimizer. Response-only validation is measured after each epoch.
A forward/backward pass on one training example warms computation without an optimizer step. Gradients are cleared, RNG restored and weights checked unchanged. Its duration and extra tokens are recorded separately. Allocated/reserved peaks are reset afterwards; reserved can retain allocator cache created during this warm-up. Optimization duration is distinguished from this preparation.
Interruption occurs at the end of the first epoch. Resume reloads the same verified base, then adapters, AdamW, sampler, scheduler, histories and Python/NumPy/PyTorch CPU/CUDA RNG. All final states must be exactly equal to those of the continuous two-epoch path. Frozen parameters and quantization metadata must retain the same hashes. This contract does not promise mid-update recovery.
The .pt states are produced locally and written through a temporary file then atomic rename. They must not be replaced by untrusted objects. Final adapters are also saved as Safetensors, with explicit base/configuration/hashes and strict reload. No frozen weight is duplicated in these adapters and no merge into another base occurs.
Generate and check, without giving the model the answers
The two bases without adaptation are evaluated separately, then both r8 and both r16 variants. Each variant freely generates the same 360 test requests with greedy decoding, one beam, at most 128 new tokens and EOS=151645. One fixed training request is a warm-up, recorded but excluded from test denominators and throughput. IDs, messages, raw text, stop reason and durations are retained.
An output stopped at the limit is TRUNCATED even if it already resembles valid SELECT or REFUSE. The parser removes no Markdown tags and repairs nothing. Apart from outer spaces, REFUSE must be exact. SQL output must belong to a complete bounded grammar: SELECT on inventory, known columns or COUNT(*)/SUM(qty)/MIN(qty)/MAX(qty), at most four predicates.
Accepted queries are passed to a separate Python process without a model. SQLite opens local databases with URI mode=ro and immutable=1; query_only, disabled extensions, a strict authorizer and VM deadline apply. Configuration PRAGMAs are trusted code run before the authorizer; model inputs cannot provide them. The authorizer accepts only SELECT and reads.
A timeout or process error retains an explicit failure. If the process stops before publishing its JSON, any intermediate operations are not counted as observed results; no result is invented. Inventories remain read-only and their integrity is checked again.
Measure the right thing
The report retains 1,800 snapshot slots per variant: five for each of 360 cases. The 30 out-of-schema requests give 150 explicitly ignored slots whatever their output; they evaluate refusal, not execution of an irrelevant query. Invalid outputs or unexpected refusals in the other 330 cases are not executed but remain failures in their denominators.
Normalized SQL exactness compares canonical textual structure without claiming to resolve every equivalence. Permuting two AND conditions can fail this exact match while producing the same results. Execution accuracy requires expected rows in expected order on all five inventories. Expected results are also computed by an independent Python interpreter.
Reported separately are valid SQL, truncated outputs, refusals, correct refusals among 30, exact SQL among 330, correct results on five inventories among 330, task accuracy among 360, actually observed executions and timeouts. A constant query that happens to work on one database is insufficient. Passing five fixtures is nevertheless neither general SQL-equivalence proof nor a guarantee on another database.
Exercise and limits
The exercise trains new r16 adapters with alpha 32, thus the same alpha/r=2 as r8/alpha 16. Request order, update count, supervised responses and optimizer remain identical. Doubling rank doubles adapter parameters, not base-model parameters. It does not automatically erase a quantization error; quality, size, memory and duration are compared.
These requests, schemas and inventories are synthetic and small. Phrasings and query types are shared; one training seed does not measure all variability. Grammar rejection reduces the admissible SQL space: a legal SQLite query outside the grammar deliberately fails. No customer SQL, database network, external tool, server or business action is used.
Create requests and reserved inventories
SQL validity and semantics
Common error: paraphrases of the same SQL are split at random; group families before splitting.
Invariant: disjoint families, oracle queries valid on their snapshots.
File src/lora_qlora_sql/data.py
import hashlib, random, sqlite3
from pathlib import Path
from .config import save_json, sha
SCHEMA = "inventory(sku TEXT, zone TEXT, qty INTEGER)"
def split_queries(family):
return "train" if family < 210 else "validation" if family < 255 else "test"
def query_spec(family):
kind = family % 10
zone = chr(65 + family // 30)
threshold = 10 + (family // 10) % 3
op = "<" if kind in [0, 2, 4, 6, 9] else ">="
projection = [
"sku",
"sku",
"sku, qty",
"sku, qty",
"COUNT(*)",
"COUNT(*)",
"SUM(qty)",
"SUM(qty)",
"MIN(qty)",
"MAX(qty)",
][kind]
order = (
" ORDER BY sku"
if kind in [0, 1]
else (
" ORDER BY qty, sku"
if kind == 2
else " ORDER BY qty DESC, sku" if kind == 3 else ""
)
)
sql = f"SELECT {projection} FROM inventory WHERE zone = '{zone}' AND qty {op} {threshold}{order};"
return {
"kind": kind,
"zone": zone,
"threshold": threshold,
"operator": op,
"sql": sql,
}
def request_text(spec, j):
kind = spec["kind"]
zone = spec["zone"]
threshold = spec["threshold"]
condition = (
f"fewer than {threshold} items"
if spec["operator"] == "<"
else f"at least {threshold} items"
)
outputs = [
"the SKUs sorted by SKU",
"the SKUs sorted by SKU",
"SKU and quantity sorted by quantity then SKU",
"SKU and quantity sorted by quantity descending then SKU",
"the number of matching rows",
"the number of matching rows",
"the sum of their quantities",
"the sum of their quantities",
"the smallest matching quantity",
"the largest matching quantity",
]
result = outputs[kind]
templates = [
f"For zone {zone}, return {result} for rows with {condition}.",
f"Find rows in zone {zone} with {condition}; report {result}.",
f"Using zone {zone} and {condition} as filters, give {result}.",
f"I need {result} among inventory rows in zone {zone} holding {condition}.",
f"Select entries with {condition} in zone {zone}, then provide {result}.",
f"In the inventory for zone {zone}, consider rows containing {condition}. Return {result}.",
f"Report {result}; restrict the stock to zone {zone} and rows with {condition}.",
f"Filter zone {zone} to stock rows with {condition}. What are {result}?",
]
return templates[j]
def generate_query_families():
pages = []
gold = []
family_sql = []
for family in range(300):
spec = query_spec(family)
family_sql.append(spec["sql"])
for j in range(8):
outside = family % 3 == 0 and j in [6, 7]
text = (
(
f"List the selling prices for SKUs in zone {spec['zone']}."
if j == 6
else f"Give the supplier names for stock in zone {spec['zone']}."
)
if outside
else request_text(spec, j)
)
# The family tag is an original ticket reference, not a SQL label.
# It keeps repeated unsupported requests traceable without claiming
# that unsupported semantic patterns are novel across partitions.
text = f"Ticket T{family:03d}-{j}. " + text
identifier = f"S{family:03d}-{j}"
pages.append(
{
"document_id": identifier,
"family": family,
"formulation": j,
"split": split_queries(family),
"request": text,
"content_sha256": hashlib.sha256(text.encode()).hexdigest(),
}
)
gold.append(
{
"document_id": identifier,
"outside_schema": outside,
"target_response": "REFUSE" if outside else spec["sql"],
"query_spec": None if outside else spec,
}
)
if len(set(family_sql)) != 300 or len({p["request"] for p in pages}) != 2400:
raise ValueError("Duplicate family SQL or ticket")
return pages, gold
def build_inventory_snapshots():
snapshots = []
for index in range(5):
generator = random.Random(45000 + index)
rows = []
for zone in range(10):
for offset in range(30):
rows.append(
{
"sku": f"K{zone*30+offset:04d}",
"zone": chr(65 + zone),
"qty": generator.randrange(21),
}
)
snapshots.append({"snapshot_id": index, "rows": rows})
return snapshots
def prepare(config):
root = Path("data/generated")
root.mkdir(parents=True, exist_ok=True)
if (root / "manifest.json").exists():
raise ValueError("Preserve generated corpus; use a fresh extraction")
pages, gold = generate_query_families()
snapshots = build_inventory_snapshots()
save_json(root / "input/pages.json", pages)
save_json(root / "labels/gold.json", gold)
save_json(root / "input/inventories.json", snapshots)
lookup = {g["document_id"]: g for g in gold}
audit = {}
for split in ["train", "validation", "test"]:
selected = [p for p in pages if p["split"] == split]
labels = [lookup[p["document_id"]] for p in selected]
save_json(root / f"{split}/pages.json", selected)
if split != "test":
save_json(root / f"labels/{split}.json", labels)
audit[split] = {
"documents": len(selected),
"families": len({p["family"] for p in selected}),
"outside_schema": sum(g["outside_schema"] for g in labels),
}
databases = []
directory = root / "databases"
directory.mkdir(parents=True, exist_ok=True)
for snapshot in snapshots:
path = directory / f"snapshot-{snapshot['snapshot_id']}.sqlite"
if path.exists():
raise ValueError("Never overwrite an inventory database")
connection = sqlite3.connect(path)
try:
connection.execute(
"CREATE TABLE inventory(sku TEXT PRIMARY KEY, zone TEXT NOT NULL, qty INTEGER NOT NULL)"
)
connection.executemany(
"INSERT INTO inventory(sku,zone,qty) VALUES(?,?,?)",
[(r["sku"], r["zone"], r["qty"]) for r in snapshot["rows"]],
)
connection.commit()
actual = connection.execute(
"SELECT sku,zone,qty FROM inventory ORDER BY sku"
).fetchall()
finally:
connection.close()
if actual != [(r["sku"], r["zone"], r["qty"]) for r in snapshot["rows"]]:
raise ValueError("SQLite fixture differs from generated rows")
databases.append(
{
"snapshot_id": snapshot["snapshot_id"],
"file": str(path),
"sha256": sha(path),
"bytes": path.stat().st_size,
"rows": len(actual),
}
)
save_json(
"outputs/database-audit.json",
{"schema": SCHEMA, "databases": databases, "generated_only": True},
)
manifest = {
"dataset": "StockSQL-v1",
"rights": "CC0-1.0",
"seed": config["seed"],
"split_audit": audit,
"snapshots": 5,
"inventory_rows_per_snapshot": 300,
"unique_family_sql": 300,
"files": {
str(f.relative_to(root)): sha(f) for f in root.rglob("*") if f.is_file()
},
}
save_json(root / "manifest.json", manifest)
save_json("outputs/data-audit.json", manifest)
return audit
Show and verify the SFT signal
Response only and boundary control
Common error: the prompt accidentally provides the target query; audit the fields actually formatted.
Invariant: branches receive the same IDs and labels.
File src/lora_qlora_sql/collator.py
import torch
from .data import SCHEMA
SYSTEM = "Translate the request into one read-only SELECT over the supplied SQLite schema. Return SQL only, with no explanation or code fence. If the request needs a column absent from the schema, return exactly REFUSE."
def format_sql_prompt(tokenizer, request):
messages = [
{"role": "system", "content": SYSTEM},
{"role": "user", "content": "Schema: " + SCHEMA + "\nRequest: " + request},
]
ids = tokenizer.apply_chat_template(
messages, tokenize=True, add_generation_prompt=True
)
return messages, ids
def make_sql_labels(tokenizer, request, response, config, identifier):
messages, prompt = format_sql_prompt(tokenizer, request)
target = tokenizer.encode(response, add_special_tokens=False) + [151645]
if (
not prompt
or len(prompt) > config["max_prompt_tokens"]
or not 1 < len(target) <= config["max_response_tokens_including_eos"]
or len(prompt) + len(target) > config["context"]
):
raise ValueError(identifier + ": token budget exceeded; no truncation")
return {
"document_id": identifier,
"messages": messages,
"input_ids": prompt + target,
"labels": [-100] * len(prompt) + target,
"prompt_tokens": len(prompt),
"response_tokens": len(target),
}
def collate(rows, device="cuda", pad_to=None):
if not rows:
raise ValueError("Empty supervised batch")
width = max(len(r["input_ids"]) for r in rows) if pad_to is None else pad_to
ids = torch.full((len(rows), width), 151643, dtype=torch.long, device=device)
mask = torch.zeros_like(ids, dtype=torch.bool)
labels = torch.full_like(ids, -100)
for i, row in enumerate(rows):
n = len(row["input_ids"])
if n > width:
raise ValueError("Do not truncate a response while padding")
ids[i, :n] = torch.tensor(row["input_ids"], device=device)
mask[i, :n] = True
labels[i, :n] = torch.tensor(row["labels"], device=device)
return ids, mask, labels
def response_loss(logits, labels):
active = labels[:, 1:] != -100
if not active.any():
raise ValueError("No active response token")
targets = labels[:, 1:][active]
selected = logits[:, :-1][active].float()
values = torch.nn.functional.cross_entropy(selected, targets, reduction="none")
return values.sum(), int(active.sum()), values, selected, targets
def inspect_sql_batch(config):
from .adapters import load_tokenizer, compare_peft_gradients
from .config import save_json
tokenizer = load_tokenizer()
request = "List the SKUs in zone B with fewer than 5 items, sorted by SKU."
response = "SELECT sku FROM inventory WHERE zone = 'B' AND qty < 5 ORDER BY sku;"
row = make_sql_labels(tokenizer, request, response, config, "illustrative-query")
ids, mask, labels = collate([row], "cpu")
save_json(
"outputs/mechanism.json",
{
"illustrative_query": True,
"request": request,
"target_response": response,
"row": row,
"tokens": tokenizer.convert_ids_to_tokens(row["input_ids"]),
"labels": labels[0].tolist(),
"attention_mask": mask[0].tolist(),
"eos": 151645,
"pad": 151643,
"prompt_and_response_encoded_separately": True,
"lora": compare_peft_gradients(),
},
)
Code LoRA, then load both bases
Low rank and quantization
Common error: k-bit loading also freezes adapters added too early; check operation order and requires_grad.
Invariant: layer control passes and measurements of both untrained bases are recorded.
File src/lora_qlora_sql/adapters.py
import base64, hashlib, json, math, os
from pathlib import Path
import numpy as np
import torch
from torch import nn
import bitsandbytes as bnb
from peft import (
LoraConfig,
get_peft_model,
prepare_model_for_kbit_training,
get_peft_model_state_dict,
set_peft_model_state_dict,
)
from peft.tuners.lora.layer import Linear as PeftLinear
from transformers import Qwen2ForCausalLM, AutoTokenizer, BitsAndBytesConfig
from safetensors.torch import save_file, load_file
from .config import save_json, read_json, sha, seed_all
TARGETS = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
class LoRALinear(nn.Module):
def __init__(self, weight, rank, alpha):
super().__init__()
if rank < 1 or alpha <= 0:
raise ValueError("Rank and scale must be positive")
self.weight = nn.Parameter(weight.clone(), requires_grad=False)
self.A = nn.Parameter(torch.empty(rank, weight.shape[1]))
self.B = nn.Parameter(torch.zeros(weight.shape[0], rank))
self.scale = alpha / rank
nn.init.kaiming_uniform_(self.A, a=math.sqrt(5))
def forward(self, x):
return torch.nn.functional.linear(
x, self.weight
) + self.scale * torch.nn.functional.linear(
torch.nn.functional.linear(x, self.A), self.B
)
def compare_peft_gradients():
torch.manual_seed(45)
weight = torch.randn(3, 5)
coded = LoRALinear(weight, 2, 4)
coded.B.data.normal_(0, 0.2)
base = nn.Linear(5, 3, bias=False)
base.weight.data.copy_(weight)
base.requires_grad_(False)
reference = PeftLinear(base, "default", r=2, lora_alpha=4, lora_dropout=0.0)
reference.lora_A["default"].weight.data.copy_(coded.A)
reference.lora_B["default"].weight.data.copy_(coded.B)
x = torch.randn(4, 5, requires_grad=True)
y = x.detach().clone().requires_grad_(True)
probe = torch.randn(4, 3)
a = coded(x)
b = reference(y)
(a * probe).sum().backward()
(b * probe).sum().backward()
pairs = {
"forward": (a, b),
"input_gradient": (x.grad, y.grad),
"A_gradient": (coded.A.grad, reference.lora_A["default"].weight.grad),
"B_gradient": (coded.B.grad, reference.lora_B["default"].weight.grad),
}
errors = {}
for name, (first, second) in pairs.items():
errors[name] = float((first - second).abs().max())
if not torch.allclose(first, second, atol=1e-6, rtol=1e-5):
raise ValueError("Coded LoRA differs from PEFT: " + name)
if coded.weight.grad is not None or reference.base_layer.weight.grad is not None:
raise ValueError("Frozen teaching base got gradients")
return {
"illustrative_fp32_layer": True,
"weight": weight.tolist(),
"A": coded.A.detach().tolist(),
"B": coded.B.detach().tolist(),
"scale": 2.0,
"input": x.detach().tolist(),
"upstream_gradient": probe.tolist(),
"output": a.detach().tolist(),
"input_gradient": x.grad.tolist(),
"A_gradient": coded.A.grad.tolist(),
"B_gradient": coded.B.grad.tolist(),
"maximum_errors": errors,
"atol": 1e-6,
"rtol": 1e-5,
"base_frozen": True,
}
def tensor_bytes(tensor):
return (
tensor.detach()
.cpu()
.contiguous()
.reshape(-1)
.view(torch.uint8)
.numpy()
.tobytes()
)
def fingerprint(state):
rows = []
for name, tensor in sorted(state.items()):
raw = tensor_bytes(tensor)
rows.append(
{
"name": name,
"shape": list(tensor.shape),
"dtype": str(tensor.dtype),
"bytes": len(raw),
"sha256": hashlib.sha256(raw).hexdigest(),
}
)
digest = hashlib.sha256(
json.dumps(rows, sort_keys=True, separators=(",", ":")).encode()
).hexdigest()
return {"sha256": digest, "tensors": rows}
def frozen_fingerprint(model):
return fingerprint(
{
k: v
for k, v in model.state_dict().items()
if "lora_A" not in k and "lora_B" not in k
}
)
def adapter_state(model):
return get_peft_model_state_dict(
model, adapter_name="default", save_embedding_layers=False
)
def adapter_fingerprint(model):
return fingerprint(adapter_state(model))
def load_tokenizer():
tokenizer = AutoTokenizer.from_pretrained(
"assets/model", local_files_only=True, trust_remote_code=False, use_fast=True
)
if not tokenizer.is_fast or tokenizer.eos_token_id != 151645:
raise ValueError("Unexpected instruct tokenizer")
tokenizer.pad_token_id = 151643
tokenizer.padding_side = "left"
return tokenizer
def target_modules(model):
result = {
name: module
for name, module in model.named_modules()
if name.split(".")[-1] in TARGETS
}
expected = {}
for i in range(28):
for name in TARGETS:
scope = (
"self_attn"
if name in ["q_proj", "k_proj", "v_proj", "o_proj"]
else "mlp"
)
shape = (
(1536, 1536)
if name in ["q_proj", "o_proj"]
else (
(256, 1536)
if name in ["k_proj", "v_proj"]
else (
(8960, 1536)
if name in ["gate_proj", "up_proj"]
else (1536, 8960)
)
)
)
expected[f"model.layers.{i}.{scope}.{name}"] = shape
if set(result) != set(expected):
raise ValueError("Unexpected target module names")
for name, module in result.items():
if (module.out_features, module.in_features) != expected[name]:
raise ValueError("Target dimensions changed")
return result
def load_precision_branch(config, precision):
if precision not in ["bf16", "nf4"] or not torch.cuda.is_bf16_supported():
raise ValueError("Expected BF16-capable CUDA device and explicit precision")
seed_all(config["seed"])
kwargs = {
"local_files_only": True,
"trust_remote_code": False,
"torch_dtype": torch.bfloat16,
"attn_implementation": "eager",
"output_loading_info": True,
}
if precision == "nf4":
kwargs.update(
{
"device_map": {"": 0},
"quantization_config": BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_use_double_quant=True,
bnb_4bit_compute_dtype=torch.bfloat16,
bnb_4bit_quant_storage=torch.uint8,
),
}
)
model, info = Qwen2ForCausalLM.from_pretrained("assets/model", **kwargs)
if (
info.get("unexpected_keys")
or info.get("mismatched_keys")
or info.get("error_msgs")
or set(info.get("missing_keys", [])) - {"lm_head.weight"}
):
raise ValueError("Incomplete original checkpoint loading: " + str(info))
c = model.config
if (
c.hidden_size,
c.num_hidden_layers,
c.num_attention_heads,
c.num_key_value_heads,
c.intermediate_size,
c.vocab_size,
) != (1536, 28, 12, 2, 8960, 151936):
raise ValueError("Unexpected model architecture")
modules = target_modules(model)
if precision == "nf4":
model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=False)
for name, module in modules.items():
if (
not isinstance(module, bnb.nn.Linear4bit)
or module.weight.quant_state is None
or module.weight.quant_type != "nf4"
or not module.weight.quant_state.nested
or module.compute_dtype != torch.bfloat16
):
raise ValueError("NF4/double quantization not active: " + name)
else:
# Match the non-quantized FP32 policy of prepare_model_for_kbit_training
# without promoting the BF16 transformer linear weights themselves.
linear_weights = {id(module.weight) for module in modules.values()}
for parameter in model.parameters():
parameter.requires_grad_(False)
if id(parameter) not in linear_weights:
parameter.data = parameter.data.float()
model = model.cuda()
if (
any(p.requires_grad for p in model.parameters())
or model.is_gradient_checkpointing
):
raise ValueError("Frozen base/checkpointing policy changed")
if model.lm_head.weight is not model.model.embed_tokens.weight:
raise ValueError("Tied embedding/output lost")
model.config.use_cache = False
model.eval()
audit = {
"precision": precision,
"target_modules": {
n: {
"in_features": m.in_features,
"out_features": m.out_features,
"class": type(m).__name__,
"weight_dtype": str(m.weight.dtype),
"weight_shape": list(m.weight.shape),
"bias_dtype": str(m.bias.dtype) if m.bias is not None else None,
"nested_quantization": (
bool(m.weight.quant_state.nested) if precision == "nf4" else False
),
}
for n, m in modules.items()
},
"embedding_dtype": str(model.model.embed_tokens.weight.dtype),
"head_dtype": str(model.lm_head.weight.dtype),
"compute_dtype": "torch.bfloat16",
"all_base_frozen": True,
"gradient_checkpointing": False,
"model_file_sha256": sha("assets/model/model.safetensors"),
"tokenizer_sha256": sha("assets/model/tokenizer.json"),
"unique_parameter_storage_bytes_excluding_external_quant_state": sum(
p.numel() * p.element_size() for p in model.parameters()
),
"library_memory_footprint_bytes": model.get_memory_footprint(),
"loading_info": info,
}
return model, audit
def install_adapters(model, config, rank):
alpha = 2 * rank
peft_config = LoraConfig(
r=rank,
lora_alpha=alpha,
lora_dropout=0.0,
target_modules=TARGETS,
bias="none",
task_type="CAUSAL_LM",
init_lora_weights=True,
use_rslora=False,
use_dora=False,
)
model = get_peft_model(model, peft_config, autocast_adapter_dtype=True)
modules = [
(name, module)
for name, module in model.named_modules()
if hasattr(module, "lora_A") and "default" in module.lora_A
]
if len(modules) != 196:
raise ValueError("Expected196adapted projections")
for index, (name, module) in enumerate(sorted(modules)):
A = module.lora_A["default"].weight
B = module.lora_B["default"].weight
if A.dtype != torch.float32 or B.dtype != torch.float32:
raise ValueError("Adapter masters must be FP32")
generator = torch.Generator(device="cpu").manual_seed(
config["seed"] * 1000 + index
)
initial = torch.empty(A.shape, dtype=torch.float32, device="cpu")
nn.init.kaiming_uniform_(initial, a=math.sqrt(5), generator=generator)
with torch.no_grad():
A.copy_(initial)
B.zero_()
trainable = {
name: parameter
for name, parameter in model.named_parameters()
if parameter.requires_grad
}
if (
any("lora_A" not in name and "lora_B" not in name for name in trainable)
or len(trainable) != 392
or sum(p.numel() for p in trainable.values()) != 1154048 * rank
):
raise ValueError("Unexpected trainable adapter state")
return model
def restore_adapter(model, state):
expected = adapter_state(model)
if expected.keys() != state.keys() or any(
expected[k].shape != state[k].shape or state[k].dtype != torch.float32
for k in state
):
raise ValueError("Adapter keys/shapes/dtypes changed")
result = set_peft_model_state_dict(model, state, adapter_name="default")
if result.unexpected_keys or any(
"lora_A" in k or "lora_B" in k for k in result.missing_keys
):
raise ValueError("Incomplete adapter load")
if fingerprint(adapter_state(model)) != fingerprint(state):
raise ValueError("Adapter reload mismatch")
def save_adapter(model, path, identity):
path = Path(path)
if path.exists():
raise ValueError("Preserve existing adapter")
state = adapter_state(model)
temporary = path.with_suffix(".pending")
save_file(
{k: v.detach().cpu().contiguous() for k, v in state.items()}, str(temporary)
)
os.replace(temporary, path)
save_json(
path.with_suffix(".json"),
{
"identity": identity,
"sha256": sha(path),
"bytes": path.stat().st_size,
"adapter": fingerprint(state),
},
)
def export_adapter_json(model, path):
state = adapter_state(model)
entries = {}
for name, tensor in state.items():
raw = (
tensor.detach()
.cpu()
.contiguous()
.numpy()
.astype("<f4", copy=False)
.tobytes()
)
entries[name] = {
"shape": list(tensor.shape),
"dtype": "float32-little-endian",
"bytes": len(raw),
"sha256": hashlib.sha256(raw).hexdigest(),
"base64": base64.b64encode(raw).decode(),
}
save_json(
path,
{
"weights": entries,
"adapter": fingerprint(state),
"parameters": sum(t.numel() for t in state.values()),
},
)
Train both comparable adaptations
Same data, budgets and optimizer
Common error: an old cached base loads on resume; reject it if SHA does not match.
Invariant: both branches consumed the planned budget and frozen weights are unchanged.
File src/lora_qlora_sql/train.py
import math, time
from pathlib import Path
import numpy as np
import os
import torch
from .config import read_json, save_json, sha, seed_all, rng_state, restore_rng
from .adapters import (
load_tokenizer,
load_precision_branch,
install_adapters,
adapter_state,
adapter_fingerprint,
frozen_fingerprint,
restore_adapter,
save_adapter,
export_adapter_json,
)
from .collator import make_sql_labels, collate
def equal(a, b):
if isinstance(a, torch.Tensor):
return isinstance(b, torch.Tensor) and torch.equal(a, b)
if isinstance(a, np.ndarray):
return isinstance(b, np.ndarray) and np.array_equal(a, b)
if isinstance(a, dict):
return (
isinstance(b, dict)
and a.keys() == b.keys()
and all(equal(a[k], b[k]) for k in a)
)
if isinstance(a, (list, tuple)):
return (
type(a) == type(b)
and len(a) == len(b)
and all(equal(x, y) for x, y in zip(a, b))
)
return a == b
def save_resume_state(path, state):
path = Path(path)
path.parent.mkdir(parents=True, exist_ok=True)
if path.exists():
raise ValueError("Existing training state is preserved")
temporary = path.with_suffix(".pending")
torch.save(state, temporary)
os.replace(temporary, path)
def encode_split(config, tokenizer, split):
pages = read_json(f"data/generated/{split}/pages.json")
labels = {
r["document_id"]: r for r in read_json(f"data/generated/labels/{split}.json")
}
if any(p["split"] != split for p in pages):
raise ValueError("Incorrect learning partition")
return [
make_sql_labels(
tokenizer,
p["request"],
labels[p["document_id"]]["target_response"],
config,
p["document_id"],
)
for p in pages
]
def active_response_logits(model, ids, mask, labels):
active = labels[:, 1:] != -100
if not active.any():
raise ValueError("No supervised response token")
with torch.autocast("cuda", dtype=torch.bfloat16):
full = model(
input_ids=ids, attention_mask=mask, use_cache=False, return_dict=True
).logits
logits = full[:, :-1][active]
targets = labels[:, 1:][active]
values = torch.nn.functional.cross_entropy(
logits.float(), targets, reduction="none"
)
return values.sum(), len(targets), values, logits, targets
def learning_rate(config, step, total):
warm = config["warmup_updates"]
peak = config["learning_rate"]
if step < warm:
return peak * (step + 1) / warm
return peak * (
1
- min(1.0, (step - warm) / max(1, total - warm - 1))
* (1 - config["minimum_lr_ratio"])
)
@torch.no_grad()
def validate_response(model, rows):
model.eval()
total = 0.0
count = 0
for row in rows:
_, n, values, _, _ = active_response_logits(model, *collate([row]))
total += float(values.double().sum())
count += n
model.train()
return {
"documents": len(rows),
"response_tokens": count,
"response_nll_nats_per_token": total / count,
}
def warm_up_training(model, row):
before = adapter_fingerprint(model)
rng = rng_state()
model.train()
torch.cuda.synchronize()
begun = time.perf_counter()
loss, n, _, _, _ = active_response_logits(model, *collate([row]))
(loss / n).backward()
torch.cuda.synchronize()
seconds = time.perf_counter() - begun
model.zero_grad(set_to_none=True)
restore_rng(rng)
if adapter_fingerprint(model) != before:
raise ValueError("Warm-up modified adapter weights")
return {
"document_id": row["document_id"],
"response_tokens": n,
"seconds": seconds,
"optimizer_steps": 0,
"rng_restored": True,
"excluded_from_optimization_budget": True,
}
def train_adapters(config, precision="bf16", rank=8, mode="continuous"):
if (
Path("data/generated/test").exists()
or Path("data/generated/labels/gold.json").exists()
):
raise ValueError("Learning service must not expose test inputs or gold")
if mode not in ["interrupt", "continuous", "resume"] or rank not in [
config["rank"],
config["exercise_rank"],
]:
raise ValueError("Unknown mode/rank")
seed_all(config["seed"])
tokenizer = load_tokenizer()
rows = encode_split(config, tokenizer, "train")
validation = encode_split(config, tokenizer, "validation")
if len(rows) != 1680 or len(validation) != 360:
raise ValueError("Expected1680/360 frozen population")
model, base_audit = load_precision_branch(config, precision)
model = install_adapters(model, config, rank)
initial = adapter_fingerprint(model)
frozen = frozen_fingerprint(model)
parameters = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.AdamW(
parameters, lr=config["learning_rate"], weight_decay=config["weight_decay"]
)
sampler = torch.Generator().manual_seed(config["seed"] + 100)
total = math.ceil(len(rows) / config["accumulation_examples"]) * config["epochs"]
branch = f"{precision}-r{rank}"
root = Path("assets/checkpoints")
destination = root / f"{branch}-{mode}.pt"
if destination.exists():
raise ValueError("Existing checkpoint is preserved")
identity = {
"branch": branch,
"precision": precision,
"rank": rank,
"alpha": rank * 2,
"config": config,
"initial_adapter": initial,
"frozen_base": frozen,
"model_sha256": sha("assets/model/model.safetensors"),
"tokenizer_sha256": sha("assets/model/tokenizer.json"),
"lock_sha256": sha("locks/models.lock.json"),
"implementation_sha256": {
name: sha("src/lora_qlora_sql/" + name)
for name in ["train.py", "adapters.py", "collator.py"]
},
"data_sha256": {
name: sha("data/generated/" + name)
for name in [
"train/pages.json",
"validation/pages.json",
"labels/train.json",
"labels/validation.json",
]
},
}
start = updates = 0
history = []
validation_history = []
if mode == "resume":
state = torch.load(
root / f"{branch}-interrupt.pt", map_location="cpu", weights_only=False
)
if state["identity"] != identity:
raise ValueError("Resume identity changed")
restore_adapter(model, state["adapter"])
optimizer.load_state_dict(state["optimizer"])
sampler.set_state(state["sampler"])
restore_rng(state["rng"])
start = state["epoch"]
updates = state["updates"]
history = state["history"]
validation_history = state["validation_history"]
if state["scheduler"] != {
"next_update": updates,
"next_lr": learning_rate(config, updates, total),
}:
raise ValueError("Resume scheduler changed")
del state
warmup = warm_up_training(model, rows[0])
torch.cuda.reset_peak_memory_stats()
torch.cuda.synchronize()
begun = time.perf_counter()
optimization_seconds = 0.0
stop = config["interrupt_epoch"] if mode == "interrupt" else config["epochs"]
model.train()
for epoch in range(start, stop):
permutation = torch.randperm(len(rows), generator=sampler).tolist()
torch.cuda.synchronize()
epoch_start = time.perf_counter()
for offset in range(0, len(rows), config["accumulation_examples"]):
selected = permutation[offset : offset + config["accumulation_examples"]]
denominator = sum(rows[i]["response_tokens"] for i in selected)
lr = learning_rate(config, updates, total)
for group in optimizer.param_groups:
group["lr"] = lr
optimizer.zero_grad(set_to_none=True)
summed = 0.0
for i in selected:
loss, n, values, _, _ = active_response_logits(
model, *collate([rows[i]])
)
if n != rows[i]["response_tokens"] or not torch.isfinite(loss):
raise ValueError("Incorrect/nonfinite response loss")
(loss / denominator).backward()
summed += float(values.detach().double().sum())
if any(
p.dtype != torch.float32
or p.grad is None
or p.grad.dtype != torch.float32
for p in parameters
):
raise ValueError("Expected FP32 adapter masters and gradients")
norm = torch.nn.utils.clip_grad_norm_(
parameters, config["gradient_clip"], error_if_nonfinite=True
)
optimizer.step()
updates += 1
history.append(
{
"update": updates,
"epoch": epoch + 1,
"document_ids": [rows[i]["document_id"] for i in selected],
"response_tokens": denominator,
"nll_nats_per_response_token": summed / denominator,
"learning_rate": lr,
"gradient_norm_before_clip": float(norm),
}
)
if updates % 42 == 0:
print(history[-1], flush=True)
torch.cuda.synchronize()
optimization_seconds += time.perf_counter() - epoch_start
observed = validate_response(model, validation)
validation_history.append({"epoch": epoch + 1, **observed})
print(validation_history[-1], flush=True)
torch.cuda.synchronize()
elapsed = time.perf_counter() - begun
peak_allocated = torch.cuda.max_memory_allocated()
peak_reserved = torch.cuda.max_memory_reserved()
final = adapter_fingerprint(model)
frozen_after = frozen_fingerprint(model)
initial_map = {r["name"]: r for r in initial["tensors"]}
changed = [
r["name"]
for r in final["tensors"]
if r["sha256"] != initial_map[r["name"]]["sha256"]
]
if (
frozen_after != frozen
or not any("lora_A" in n for n in changed)
or not any("lora_B" in n for n in changed)
):
raise ValueError("Frozen base changed or A/B learning absent")
state = {
"identity": identity,
"epoch": stop,
"updates": updates,
"adapter": adapter_state(model),
"optimizer": optimizer.state_dict(),
"sampler": sampler.get_state(),
"scheduler": {
"next_update": updates,
"next_lr": learning_rate(config, updates, total),
},
"rng": rng_state(),
"history": history,
"validation_history": validation_history,
}
save_resume_state(destination, state)
if mode == "resume":
save_adapter(model, root / f"{branch}-resume.safetensors", identity)
export_adapter_json(model, f"outputs/adapter-{branch}.json")
save_json(
f"outputs/train-{branch}-{mode}.json",
{
"branch": branch,
"mode": mode,
"identity": identity,
"base_audit": base_audit,
"epochs": stop,
"updates": updates,
"history": history,
"validation_history": validation_history,
"tokenization": rows,
"initial_adapter": initial,
"final_adapter": final,
"changed_parameters": changed,
"frozen_base_unchanged": True,
"trainable_parameters": sum(p.numel() for p in parameters),
"checkpoint_sha256": sha(destination),
"checkpoint_bytes": destination.stat().st_size,
"warmup": warmup,
"optimization_seconds_this_process": optimization_seconds,
"training_and_validation_seconds_this_process": elapsed,
"max_torch_allocated_bytes": peak_allocated,
"max_torch_reserved_bytes": peak_reserved,
"reserved_includes_warmup_allocator_cache": True,
"test_visible": False,
},
)
def verify_resume(precision, rank):
branch = f"{precision}-r{rank}"
root = Path("assets/checkpoints")
a = torch.load(
root / f"{branch}-continuous.pt", map_location="cpu", weights_only=False
)
b = torch.load(root / f"{branch}-resume.pt", map_location="cpu", weights_only=False)
checks = {k: equal(a[k], b[k]) for k in a}
if a.keys() != b.keys() or not all(checks.values()):
raise ValueError("Exact adapter resume differs: " + str(checks))
save_json(
f"outputs/resume-{branch}.json",
{
"status": "PASS",
"all_state_exact": True,
"checks": checks,
"boundary": "complete epoch; adapter/optimizer/sampler/scheduler/all RNG; frozen base reloaded by verified identity",
},
)
Run SELECT queries under an independent control
Safe interpretation and measured memory
Common error: a string containing two statements passes a simple startswith SELECT; use a real parser/authorizer.
Invariant: every case has five results or an explicit refusal/failure; no write is possible.
File src/lora_qlora_sql/evaluate.py
import json, math, subprocess, sys, tempfile, time
from pathlib import Path
import torch
from safetensors.torch import load_file
from transformers import GenerationConfig
from .config import read_json, save_json, sha, seed_all
from .adapters import (
load_tokenizer,
load_precision_branch,
install_adapters,
restore_adapter,
adapter_fingerprint,
frozen_fingerprint,
)
from .collator import format_sql_prompt, make_sql_labels, collate
from .train import active_response_logits
from .sql_sandbox import validate_readonly_sql, normalize_sql
def parse_completion(text, finish):
if finish != "eos":
return {"status": "TRUNCATED", "sql": None, "tree": None}
text = text.strip()
if text == "REFUSE":
return {"status": "REFUSE", "sql": None, "tree": None}
try:
tree = validate_readonly_sql(text)
except ValueError as error:
return {"status": "INVALID_SQL", "sql": None, "tree": None, "error": str(error)}
return {"status": "VALID_SQL", "sql": normalize_sql(text), "tree": tree}
def load_variant(config, precision, rank):
model, audit = load_precision_branch(config, precision)
identity = {
"precision": precision,
"rank": rank,
"model_sha256": sha("assets/model/model.safetensors"),
}
if rank:
model = install_adapters(model, config, rank)
path = Path(f"assets/checkpoints/{precision}-r{rank}-resume.safetensors")
metadata = read_json(path.with_suffix(".json"))
if (
sha(path) != metadata["sha256"]
or metadata["identity"]["config"] != config
or metadata["identity"]["precision"] != precision
or metadata["identity"]["rank"] != rank
or metadata["identity"]["frozen_base"] != frozen_fingerprint(model)
):
raise ValueError("Learned adapter/base identity changed")
restore_adapter(model, load_file(str(path), device="cpu"))
if adapter_fingerprint(model) != metadata["adapter"]:
raise ValueError("Adapter manifest differs")
identity.update(
{
"adapter_sha256": sha(path),
"adapter_bytes": path.stat().st_size,
"adapter": metadata["adapter"],
}
)
return model.eval(), audit, identity
@torch.no_grad()
def generate_one(config, model, tokenizer, page, generation):
messages, ids = format_sql_prompt(tokenizer, page["request"])
if len(ids) > config["max_prompt_tokens"]:
raise ValueError("Generation prompt over token cap; do not truncate")
tensor = torch.tensor([ids], device="cuda")
torch.cuda.synchronize()
begun = time.perf_counter()
with torch.autocast("cuda", dtype=torch.bfloat16):
output = (
model.generate(
input_ids=tensor,
attention_mask=torch.ones_like(tensor),
generation_config=generation,
)[0, len(ids) :]
.cpu()
.tolist()
)
torch.cuda.synchronize()
seconds = time.perf_counter() - begun
finish = "eos" if output and output[-1] == 151645 else "length"
content = output[:-1] if finish == "eos" else output
text = tokenizer.decode(
content, skip_special_tokens=False, clean_up_tokenization_spaces=False
)
return {
"document_id": page["document_id"],
"messages": messages,
"input_ids": ids,
"output_ids": output,
"text": text,
"finish_reason": finish,
"parsed": parse_completion(text, finish),
"input_tokens": len(ids),
"output_tokens": len(output),
"seconds": seconds,
}
@torch.no_grad()
def predict(config, precision, rank=0):
if Path("data/generated/labels").exists():
raise ValueError("Blind generation must not mount labels")
seed_all(config["seed"])
tokenizer = load_tokenizer()
model, audit, identity = load_variant(config, precision, rank)
pages = read_json("data/generated/input/pages.json")
test = [p for p in pages if p["split"] == "test"]
warmup_page = next(p for p in pages if p["split"] == "train")
variant = f"{precision}-r{rank}"
if len(test) != 360:
raise ValueError("Keep all360heldout cases")
generation = GenerationConfig(
do_sample=False,
num_beams=1,
max_new_tokens=config["max_new_tokens"],
use_cache=True,
eos_token_id=151645,
pad_token_id=151643,
bos_token_id=151643,
)
warmup = generate_one(config, model, tokenizer, warmup_page, generation)
torch.cuda.reset_peak_memory_stats()
torch.cuda.synchronize()
begun = time.perf_counter()
rows = []
for index, page in enumerate(test):
rows.append(generate_one(config, model, tokenizer, page, generation))
if (index + 1) % 60 == 0:
print(
{"variant": variant, "generated": index + 1, "total": 360}, flush=True
)
torch.cuda.synchronize()
seconds = time.perf_counter() - begun
tokens = sum(r["output_tokens"] for r in rows)
save_json(
f"outputs/predictions-{variant}.json",
{
"variant": variant,
"identity": identity,
"base_audit": audit,
"tokenizer_sha256": sha("assets/model/tokenizer.json"),
"input_sha256": sha("data/generated/input/pages.json"),
"generation": generation.to_dict(),
"warmup": warmup,
"warmup_excluded": True,
"rows": rows,
"seconds": seconds,
"generated_tokens_including_eos": tokens,
"generated_tokens_per_second": tokens / seconds,
"max_torch_allocated_bytes": torch.cuda.max_memory_allocated(),
"max_torch_reserved_bytes": torch.cuda.max_memory_reserved(),
"reserved_includes_warmup_allocator_cache": True,
"gold_visible": False,
},
)
@torch.no_grad()
def response_nll(config, precision, rank=0):
seed_all(config["seed"])
tokenizer = load_tokenizer()
model, _, identity = load_variant(config, precision, rank)
variant = f"{precision}-r{rank}"
pages = [
p for p in read_json("data/generated/input/pages.json") if p["split"] == "test"
]
gold = {g["document_id"]: g for g in read_json("data/generated/labels/gold.json")}
rows = []
probes = []
for index, page in enumerate(pages):
row = make_sql_labels(
tokenizer,
page["request"],
gold[page["document_id"]]["target_response"],
config,
page["document_id"],
)
_, n, losses, logits, targets = active_response_logits(model, *collate([row]))
rows.append(
{
"document_id": page["document_id"],
"input_ids": row["input_ids"],
"prompt_tokens": row["prompt_tokens"],
"response_tokens": n,
"response_targets": targets.cpu().tolist(),
"response_nll": losses.cpu().tolist(),
}
)
if index == 0:
for position in [0, n - 1]:
probes.append(
{
"document_id": page["document_id"],
"response_position": position,
"target": int(targets[position]),
"logits": logits[position].float().cpu().tolist(),
"nll": float(losses[position]),
}
)
count = sum(r["response_tokens"] for r in rows)
total = math.fsum(v for row in rows for v in row["response_nll"])
save_json(
f"outputs/response-nll-{variant}.json",
{
"variant": variant,
"identity": identity,
"documents": len(rows),
"response_tokens": count,
"response_nll_nats_per_token": total / count,
"rows": rows,
"probes": probes,
},
)
def reference_result(tree, inventory):
def matches(row):
for p in tree["predicates"]:
a = row[p["column"]]
b = p["value"]
op = p["operator"]
ok = {
"=": a == b,
"<>": a != b,
"<": a < b,
">": a > b,
"<=": a <= b,
">=": a >= b,
}[op]
if not ok:
return False
return True
selected = [r for r in inventory if matches(r)]
aggregate = tree["aggregate"]
if aggregate:
values = [r["qty"] for r in selected]
fn = aggregate["function"]
value = (
len(values)
if fn == "COUNT"
else (
(
sum(values)
if fn == "SUM"
else min(values) if fn == "MIN" else max(values)
)
if values
else None
)
)
return [[value]]
for order in reversed(tree["order"]):
selected.sort(
key=lambda r: r[order["column"]], reverse=order["direction"] == "DESC"
)
return [[r[c] for c in tree["columns"]] for r in selected]
def run_sql_worker(config, sql):
validate_readonly_sql(sql)
with tempfile.TemporaryDirectory(prefix="sql-request-") as directory:
path = Path(directory) / "request.json"
path.write_text(
json.dumps(
{
"sql": sql,
"vm_seconds": config["sqlite_vm_seconds"],
"max_rows": config["sqlite_max_rows"],
}
)
)
begun = time.perf_counter()
try:
result = subprocess.run(
[sys.executable, "scripts/sql_worker.py", "--request", str(path)],
capture_output=True,
text=True,
timeout=config["sqlite_worker_timeout_seconds"],
check=False,
)
except subprocess.TimeoutExpired:
return {
"status": "PROCESS_TIMEOUT",
"results": [],
"seconds": time.perf_counter() - begun,
}
if result.returncode:
return {
"status": "WORKER_ERROR",
"results": [],
"returncode": result.returncode,
"stderr": result.stderr[:2000],
"seconds": time.perf_counter() - begun,
}
try:
payload = json.loads(result.stdout)
except ValueError:
return {
"status": "WORKER_ERROR",
"results": [],
"error": "Worker output was not JSON",
"seconds": time.perf_counter() - begun,
}
if (
not isinstance(payload, dict)
or set(payload) != {"results"}
or not isinstance(payload["results"], list)
or len(payload["results"]) != 5
or [r.get("snapshot_id") for r in payload["results"]] != list(range(5))
):
raise ValueError("Malformed SQL worker result")
return {
"status": "COMPLETED",
"results": payload["results"],
"seconds": time.perf_counter() - begun,
}
def evaluate_variants(config, variants, exercise=False):
pages = {
p["document_id"]: p
for p in read_json("data/generated/input/pages.json")
if p["split"] == "test"
}
gold = {g["document_id"]: g for g in read_json("data/generated/labels/gold.json")}
snapshots = read_json("data/generated/input/inventories.json")
databases = read_json("outputs/database-audit.json")["databases"]
before = {d["file"]: sha(d["file"]) for d in databases}
result = {}
if before != {d["file"]: d["sha256"] for d in databases}:
raise ValueError("Original inventory changed before evaluation")
for variant in variants:
predictions = read_json(f"outputs/predictions-{variant}.json")
records = predictions["rows"]
if len(records) != 360 or {r["document_id"] for r in records} != set(pages):
raise ValueError("Every test case must be retained exactly once")
counts = {
"documents": 360,
"eligible_sql": 330,
"outside_schema": 30,
"valid_sql": 0,
"truncated": 0,
"refusal": 0,
"correct_refusal": 0,
"normalized_sql_exact": 0,
"execution_correct_all_five": 0,
"task_correct": 0,
"snapshot_slots": 1800,
"outside_schema_skipped_slots": 150,
"actual_snapshot_executions": 0,
"successful_snapshot_executions": 0,
"worker_timeouts": 0,
}
details = []
for row in records:
label = gold[row["document_id"]]
parsed = parse_completion(row["text"], row["finish_reason"])
if parsed != row["parsed"]:
raise ValueError("Recorded SQL parse changed")
status = parsed["status"]
counts["valid_sql"] += status == "VALID_SQL"
counts["truncated"] += status == "TRUNCATED"
counts["refusal"] += status == "REFUSE"
outside = label["outside_schema"]
correct_refusal = outside and status == "REFUSE"
counts["correct_refusal"] += correct_refusal
exact = False
execution = False
worker = {
"status": (
"OUTSIDE_SCHEMA_SKIP" if outside else "INVALID_OR_REFUSED_SKIP"
),
"results": [],
}
comparisons = []
if not outside and status == "VALID_SQL":
exact = parsed["sql"] == normalize_sql(label["target_response"])
worker = run_sql_worker(config, parsed["sql"])
counts["worker_timeouts"] += worker["status"] == "PROCESS_TIMEOUT"
counts["actual_snapshot_executions"] += len(worker["results"])
expected_tree = validate_readonly_sql(label["target_response"])
for observed in worker["results"]:
expected = reference_result(
expected_tree, snapshots[observed["snapshot_id"]]["rows"]
)
correct = (
observed["status"] == "OK" and observed["rows"] == expected
)
comparisons.append(
{
"snapshot_id": observed["snapshot_id"],
"expected_rows": expected,
"correct": correct,
}
)
counts["successful_snapshot_executions"] += (
observed["status"] == "OK"
)
execution = len(comparisons) == 5 and all(
c["correct"] for c in comparisons
)
counts["normalized_sql_exact"] += exact
counts["execution_correct_all_five"] += execution
counts["task_correct"] += bool(correct_refusal or execution)
details.append(
{
"document_id": row["document_id"],
"outside_schema": outside,
"parsed_status": status,
"correct_refusal": correct_refusal,
"normalized_sql_exact": exact,
"execution_correct_all_five": execution,
"worker": worker,
"snapshot_comparisons": comparisons,
}
)
counts.update(
{
"normalized_sql_exact_rate_over_330": counts["normalized_sql_exact"]
/ 330,
"execution_accuracy_over_330": counts["execution_correct_all_five"]
/ 330,
"refusal_accuracy_over_30": counts["correct_refusal"] / 30,
"task_accuracy_over_360": counts["task_correct"] / 360,
}
)
result[variant] = {"metrics": counts, "rows": details}
print({"variant": variant, **counts}, flush=True)
after = {path: sha(path) for path in before}
if before != after:
raise ValueError("Read-only evaluation modified an inventory")
save_json(
"outputs/" + ("exercise-evaluation.json" if exercise else "evaluation.json"),
{
"results": result,
"database_sha256_before": before,
"database_sha256_after": after,
"database_unchanged": True,
"quality_gain_required": False,
"execution_equivalence_on_five_fixtures_is_not_a_general_SQL_proof": True,
},
)
Test rank capacity without changing scale
One-factor ablation
Common error: comparing an r16 trained longer distorts the rank effect; retain supervised tokens.
Invariant: identical alpha/r and newly calculated adapter size.
File src/lora_qlora_sql/exercise.py
from pathlib import Path
from .config import read_json, save_json, sha
def main_artifacts():
paths = ["outputs/evaluation.json"]
for precision in ["bf16", "nf4"]:
paths += [
f"assets/checkpoints/{precision}-r8-resume.safetensors",
f"outputs/resume-{precision}-r8.json",
f"outputs/adapter-{precision}-r8.json",
]
for rank in [0, 8]:
paths += [
f"outputs/predictions-{precision}-r{rank}.json",
f"outputs/response-nll-{precision}-r{rank}.json",
]
return paths
def preserve_main():
path = Path("outputs/exercise/protected.json")
if path.exists():
raise ValueError("Preserve existing exercise snapshot")
save_json(path, {name: sha(name) for name in main_artifacts()})
def compare_ranks():
before = read_json("outputs/exercise/protected.json")
if before != {name: sha(name) for name in main_artifacts()}:
raise ValueError("Rank exercise altered a main artifact")
comparisons = []
for rank in [8, 16]:
a = read_json(f"outputs/train-bf16-r{rank}-continuous.json")
b = read_json(f"outputs/train-nf4-r{rank}-continuous.json")
if a["initial_adapter"] != b["initial_adapter"]:
raise ValueError("Precision branches did not start from the same adapters")
if [
(r["document_ids"], r["response_tokens"], r["learning_rate"])
for r in a["history"]
] != [
(r["document_ids"], r["response_tokens"], r["learning_rate"])
for r in b["history"]
]:
raise ValueError("Precision branches used different learning schedules")
comparisons.append(
{
"rank": rank,
"alpha": 2 * rank,
"initial_adapters_exact": True,
"example_order_and_token_budget_exact": True,
"trainable_parameters": a["trainable_parameters"],
"updates": a["updates"],
"presentations": sum(len(r["document_ids"]) for r in a["history"]),
"response_tokens": sum(r["response_tokens"] for r in a["history"]),
}
)
save_json(
"outputs/exercise/rank-comparison.json",
{
"protected_sha256": before,
"main_artifacts_unchanged": True,
"fresh_rank16_both_precisions": True,
"same_scale_alpha_over_rank": 2.0,
"comparisons": comparisons,
"quality_gain_required": False,
"rank16_has_more_parameters": True,
},
)
Sources
LoRA, QLoRA, Qwen2.5-1.5B-Instruct, PEFT 0.17.1 k-bit preparation, PEFT 0.17.1 LoRA layers and bitsandbytes 0.47.0 layers. Local choices, including no paged optimizer and no gradient checkpointing, are explained in this tutorial.
What was observed during development
The complete path passed its 20 public tests. The independent audit recalculated SQL decisions and results on five inventories, denominators, response losses, adapter identities and resume checks. This technical validation does not require a quality gain.
The six conditions receive the same 360 test requests: 330 within the schema and 30 asking for an absent field. A query is counted correct only if its results match the reference on all five inventories. The five inventories are not proof of equivalence on every possible database.
| Condition | Correct SQL on the five bases | Correct refusal | Correct task | Truncated responses |
|---|---|---|---|---|
| bf16-r0 | 0/330 | 0/30 | 0/360 | 0 |
| nf4-r0 | 10/330 | 0/30 | 10/360 | 0 |
| bf16-r8 | 325/330 | 30/30 | 355/360 | 0 |
| nf4-r8 | 324/330 | 30/30 | 354/360 | 0 |
| bf16-r16 | 318/330 | 30/30 | 348/360 | 0 |
| nf4-r16 | 329/330 | 30/30 | 359/360 | 0 |
The score belongs to the complete “SQL only or REFUSE” contract. A Markdown block, explanation or filter outside the grammar is rejected; the score therefore does not isolate abstract SQL knowledge. For the first request S255-0, observed exact answers are retained below.
{
"bf16-r0": {
"document_id": "S255-0",
"text": "```sql\nSELECT COUNT(*) FROM inventory WHERE zone = 'I' AND qty >= 1\n```",
"status": "INVALID_SQL"
},
"nf4-r0": {
"document_id": "S255-0",
"text": "```sql\nSELECT COUNT(*) FROM inventory WHERE zone = 'I' AND qty >= 11;\n```",
"status": "INVALID_SQL"
},
"bf16-r8": {
"document_id": "S255-0",
"text": "SELECT COUNT(*) FROM inventory WHERE zone = 'I' AND qty >= 11;",
"status": "VALID_SQL"
},
"nf4-r8": {
"document_id": "S255-0",
"text": "SELECT COUNT(*) FROM inventory WHERE zone = 'I' AND qty >= 11;",
"status": "VALID_SQL"
},
"bf16-r16": {
"document_id": "S255-0",
"text": "SELECT COUNT(*) FROM inventory WHERE zone = 'I' AND qty >= 11;",
"status": "VALID_SQL"
},
"nf4-r16": {
"document_id": "S255-0",
"text": "SELECT COUNT(*) FROM inventory WHERE zone = 'I' AND qty >= 11;",
"status": "VALID_SQL"
}
}
The following loss covers the same reference responses with their EOS, not only generated answers that succeeded. Its unit is nat per response token. A lower loss does not guarantee a correct query in free generation.
| Condition | NLL per response token |
|---|---|
| bf16-r0 | 0.2835442012 |
| nf4-r0 | 0.2712743014 |
| bf16-r8 | 0.003191479554 |
| nf4-r8 | 0.002342658512 |
| bf16-r16 | 0.005727622692 |
| nf4-r16 | 0.0006154038 |
The times below belong to continuous training paths in this execution. Every condition has 420 updates. Interrupted/resumed paths, acquisitions, validations, inferences and exports also have a recorded cost; this table is not the total path cost.
| Condition | Trained parameters | Optimization (s) | Optimization and validation (s) | Peak PyTorch allocated (bytes) | Peak reserved (bytes) |
|---|---|---|---|---|---|
| bf16-r8 | 9232384 | 759.159 | 821.304 | 4780346880 | 5114953728 |
| nf4-r8 | 9232384 | 992.322 | 1080.726 | 3044170752 | 3351248896 |
| bf16-r16 | 18464768 | 766.232 | 829.986 | 4946857472 | 5284823040 |
| nf4-r16 | 18464768 | 980.464 | 1068.256 | 3210681344 | 3521118208 |
Allocated and reserved are PyTorch counters; they do not denote all process VRAM observed by the driver. The lesson's memory figure measures test generations, so it is another phase from this table. Its allocation cache can include the declared warm-up. One seed, a procedural corpus and one machine cannot promise these times or this quality.
Complete outputs allow errors to be examined without retroactively correcting the parser, test set or hyperparameters. The rank-16 exercise retains rank-8 results and repeats both precisions. No systematic advantage of rank, NF4 or LoRA is inferred from a simple PASS.
Development observations
Values recalculated from this run's raw outputs. They are not performance thresholds; the path from the ZIP recalculates them on the same populations.
{
"status": "PASS",
"documents": 2400,
"test_documents": 360,
"snapshots": 5,
"rank8_parameters": 9232384,
"rank16_parameters": 18464768,
"metrics": {
"bf16-r0": {
"documents": 360,
"eligible_sql": 330,
"outside_schema": 30,
"valid_sql": 0,
"truncated": 0,
"refusal": 0,
"correct_refusal": 0,
"normalized_sql_exact": 0,
"execution_correct_all_five": 0,
"task_correct": 0,
"snapshot_slots": 1800,
"outside_schema_skipped_slots": 150,
"actual_snapshot_executions": 0,
"successful_snapshot_executions": 0,
"worker_timeouts": 0,
"normalized_sql_exact_rate_over_330": 0.0,
"execution_accuracy_over_330": 0.0,
"refusal_accuracy_over_30": 0.0,
"task_accuracy_over_360": 0.0
},
"nf4-r0": {
"documents": 360,
"eligible_sql": 330,
"outside_schema": 30,
"valid_sql": 10,
"truncated": 0,
"refusal": 0,
"correct_refusal": 0,
"normalized_sql_exact": 8,
"execution_correct_all_five": 10,
"task_correct": 10,
"snapshot_slots": 1800,
"outside_schema_skipped_slots": 150,
"actual_snapshot_executions": 50,
"successful_snapshot_executions": 50,
"worker_timeouts": 0,
"normalized_sql_exact_rate_over_330": 0.024242424242424242,
"execution_accuracy_over_330": 0.030303030303030304,
"refusal_accuracy_over_30": 0.0,
"task_accuracy_over_360": 0.027777777777777776
},
"bf16-r8": {
"documents": 360,
"eligible_sql": 330,
"outside_schema": 30,
"valid_sql": 325,
"truncated": 0,
"refusal": 30,
"correct_refusal": 30,
"normalized_sql_exact": 325,
"execution_correct_all_five": 325,
"task_correct": 355,
"snapshot_slots": 1800,
"outside_schema_skipped_slots": 150,
"actual_snapshot_executions": 1625,
"successful_snapshot_executions": 1625,
"worker_timeouts": 0,
"normalized_sql_exact_rate_over_330": 0.9848484848484849,
"execution_accuracy_over_330": 0.9848484848484849,
"refusal_accuracy_over_30": 1.0,
"task_accuracy_over_360": 0.9861111111111112
},
"nf4-r8": {
"documents": 360,
"eligible_sql": 330,
"outside_schema": 30,
"valid_sql": 324,
"truncated": 0,
"refusal": 30,
"correct_refusal": 30,
"normalized_sql_exact": 324,
"execution_correct_all_five": 324,
"task_correct": 354,
"snapshot_slots": 1800,
"outside_schema_skipped_slots": 150,
"actual_snapshot_executions": 1620,
"successful_snapshot_executions": 1620,
"worker_timeouts": 0,
"normalized_sql_exact_rate_over_330": 0.9818181818181818,
"execution_accuracy_over_330": 0.9818181818181818,
"refusal_accuracy_over_30": 1.0,
"task_accuracy_over_360": 0.9833333333333333
},
"bf16-r16": {
"documents": 360,
"eligible_sql": 330,
"outside_schema": 30,
"valid_sql": 318,
"truncated": 0,
"refusal": 30,
"correct_refusal": 30,
"normalized_sql_exact": 318,
"execution_correct_all_five": 318,
"task_correct": 348,
"snapshot_slots": 1800,
"outside_schema_skipped_slots": 150,
"actual_snapshot_executions": 1590,
"successful_snapshot_executions": 1590,
"worker_timeouts": 0,
"normalized_sql_exact_rate_over_330": 0.9636363636363636,
"execution_accuracy_over_330": 0.9636363636363636,
"refusal_accuracy_over_30": 1.0,
"task_accuracy_over_360": 0.9666666666666667
},
"nf4-r16": {
"documents": 360,
"eligible_sql": 330,
"outside_schema": 30,
"valid_sql": 329,
"truncated": 0,
"refusal": 30,
"correct_refusal": 30,
"normalized_sql_exact": 329,
"execution_correct_all_five": 329,
"task_correct": 359,
"snapshot_slots": 1800,
"outside_schema_skipped_slots": 150,
"actual_snapshot_executions": 1645,
"successful_snapshot_executions": 1645,
"worker_timeouts": 0,
"normalized_sql_exact_rate_over_330": 0.996969696969697,
"execution_accuracy_over_330": 0.996969696969697,
"refusal_accuracy_over_30": 1.0,
"task_accuracy_over_360": 0.9972222222222222
}
},
"response_nll": {
"bf16-r0": 0.2835442011762571,
"nf4-r0": 0.2712743014305074,
"bf16-r8": 0.003191479553683838,
"nf4-r8": 0.002342658512238993,
"bf16-r16": 0.005727622691669397,
"nf4-r16": 0.0006154038000183541
},
"quality_gain_required": false,
"scope": "Data-only independent SQL grammar/results on all five fixtures, tokenizer byte decoding, adapter bytes/gradients, frozen base metadata, exact budget/resume evidence and all-case metrics; not BPE re-encoding or full quantized Transformer reimplementation"
}
