LoRA Fine-Tuning skill for Codex

Done well means: the base model’s own chat template formats every example, loss is computed only on assistant tokens, the run fits memory with headroom, and the adapter beats the base model on a held-out eval before anyone merges or ships it.

Before training

1. Read AGENTS.md, pyproject.toml or requirements.txt; pin versions of torch, transformers, peft, trl, bitsandbytes, accelerate, datasets.
2. Check hardware: nvidia-smi (VRAM, driver), python -c "import torch; print(torch.cuda.is_available(), torch.cuda.get_device_name(0))". bitsandbytes 4-bit needs CUDA (recent releases add limited ROCm/Intel/Apple support; verify).
3. Confirm the base model’s license allows fine-tuning and the intended use. Gated models need hf auth login (older: huggingface-cli login) or HF_TOKEN, and downloads need network.
4. Build the held-out eval set first and score the base model with a good prompt. No baseline, no claim of improvement.

Installs and model downloads need network; in a Codex sandbox ask for approval, or point HF_HOME at an existing cache and set HF_HUB_OFFLINE=1.

Method choice

Situation	Choice
7-8B model, one 24 GB GPU	QLoRA (4-bit NF4 base + LoRA adapters)
7-8B model, 48-80 GB GPU	LoRA on a bf16 base: faster, slightly better quality
70B model	QLoRA on 48-80 GB, or LoRA with FSDP across GPUs
Teach format, tone, labels, tool-call style	SFT
Prefer one style over another, reduce a failure mode	DPO after SFT, on preference pairs
Add knowledge the model must cite	Not fine-tuning: retrieval ($rag-pipelines)

Rough VRAM for QLoRA SFT at 2k tokens with gradient checkpointing: 7-8B about 10-16 GB, 13-14B about 18-24 GB, 70B about 45-60 GB. Measure; do not trust the table over nvidia-smi.

Dataset

Use the conversational format TRL understands and let the tokenizer’s chat template render it:

{"messages": [{"role": "system", "content": "You write SQL for Postgres 16."}, {"role": "user", "content": "Top 5 customers by revenue"}, {"role": "assistant", "content": "SELECT customer_id, SUM(amount) AS revenue FROM orders GROUP BY 1 ORDER BY 2 DESC LIMIT 5;"}]}

• Prompt-completion pairs: {"prompt": [...messages], "completion": [...messages]}.
• DPO: {"prompt": ..., "chosen": ..., "rejected": ...} (conversational or plain text).
• Split train/eval by source document or user, not by row, so near-duplicates do not leak.
• Inspect three rendered examples: tok.apply_chat_template(ex["messages"], tokenize=False). Check special tokens, the system prompt and the end-of-turn token appear exactly once.

SFT with QLoRA

import torch
from datasets import load_dataset
from peft import LoraConfig
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from trl import SFTConfig, SFTTrainer

base = "Qwen/Qwen2.5-7B-Instruct"
tok = AutoTokenizer.from_pretrained(base)

bnb = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True,
)
model = AutoModelForCausalLM.from_pretrained(base, quantization_config=bnb, device_map="auto")

peft_config = LoraConfig(
    r=16, lora_alpha=32, lora_dropout=0.05,
    target_modules="all-linear", task_type="CAUSAL_LM",
)

ds = load_dataset("json", data_files={"train": "train.jsonl", "eval": "eval.jsonl"})

args = SFTConfig(
    output_dir="out/qwen7b-sql-lora",
    num_train_epochs=2,
    per_device_train_batch_size=2,
    gradient_accumulation_steps=8,          # effective batch 16
    learning_rate=2e-4,
    lr_scheduler_type="cosine",
    warmup_ratio=0.03,
    max_length=2048,
    assistant_only_loss=True,               # needs a chat template with generation markers
    gradient_checkpointing=True,
    bf16=True,
    logging_steps=10,
    eval_strategy="steps", eval_steps=100,
    save_strategy="steps", save_steps=100, save_total_limit=3,
    load_best_model_at_end=True,
    report_to="none",                        # or "wandb"/"mlflow" (see ml-experiment-tracking)
    seed=42,
)

trainer = SFTTrainer(
    model=model, args=args, processing_class=tok,
    train_dataset=ds["train"], eval_dataset=ds["eval"], peft_config=peft_config,
)
trainer.train()
trainer.save_model(args.output_dir)          # saves the adapter only

TRL renames arguments between releases (max_seq_length became max_length; dataset_text_field and completion-only collators changed). Check the installed version’s SFTConfig signature before copying older snippets. If the template lacks {% generation %} markers, assistant_only_loss fails; use prompt-completion data (completion-only loss is the default there) instead.

Hyperparameter defaults and fixes

Knob	Default	Adjust when
r / lora_alpha	16 / 32	Underfits on a complex task: 32-64; tiny data: 8
target_modules	"all-linear"	VRAM tight: attention projections only (q_proj, v_proj, …)
learning rate	1e-4 to 2e-4 (LoRA)	DPO: 5e-6 to 5e-5
epochs	1-3	Eval loss rises after epoch 1: stop at 1
effective batch	16-64	Noisy loss: raise via gradient accumulation
max_length	p95 of example lengths	Truncation silently drops the answer at the end

Symptom	Likely cause
Loss starts near 0	Loss computed on prompt + padding, or labels leaked; check masking
Loss NaN in fp16	Use bf16 on Ampere or newer; lower LR
Model never stops generating	EOS / end-of-turn token missing from training text
CUDA OOM	Halve batch, double accumulation; enable checkpointing; shorten max_length
Great eval loss, worse answers	Overfit on phrasing; fewer epochs; judge with task evals, not loss

DPO after SFT

from trl import DPOConfig, DPOTrainer

dpo_args = DPOConfig(output_dir="out/dpo", beta=0.1, learning_rate=5e-6,
                     per_device_train_batch_size=1, gradient_accumulation_steps=16,
                     num_train_epochs=1, bf16=True, max_length=2048)
trainer = DPOTrainer(model=sft_model, args=dpo_args, processing_class=tok,
                     train_dataset=prefs["train"], eval_dataset=prefs["eval"],
                     peft_config=peft_config)

With a PEFT model, TRL uses the frozen base as the reference model, so no second copy is loaded. Watch rewards/margins rise and rewards/accuracies move above 0.5.

Evaluate, merge, export

from peft import AutoPeftModelForCausalLM

model = AutoPeftModelForCausalLM.from_pretrained("out/qwen7b-sql-lora", dtype=torch.bfloat16, device_map="auto")
merged = model.merge_and_unload()
merged.save_pretrained("out/qwen7b-sql-merged", safe_serialization=True)
tok.save_pretrained("out/qwen7b-sql-merged")

• Merge into a bf16 copy of the base, not into the 4-bit model; quantize afterwards ($llama-cpp-gguf, $model-quantization).
• Or keep adapters separate and serve several on one base with vLLM --enable-lora ($vllm-serving).
• Evaluate the merged artifact you will ship, with greedy decoding and the production prompt, against the base-model baseline ($llm-evals).

Gotchas

• Training on the base (non-instruct) model with an instruct template, or the reverse, produces a model that ignores turns. Match the template to the checkpoint.
• Padding side: decoder-only generation wants left padding at inference; training collators handle it, custom loops often do not.
• device_map="auto" with multiple GPUs in one process is model parallel, not data parallel; use accelerate launch or torchrun for DDP.
• Pushing adapters to the Hub publishes them: ask before push_to_hub=True, and check private-repo settings.
• Never upload proprietary or personal data to a hosted notebook without user approval.

Report template

**Base:** Qwen/Qwen2.5-7B-Instruct (Apache-2.0) | **Method:** QLoRA SFT, r=16, all-linear, 2 epochs
**Data:** 3,200 train / 400 eval / 300 held-out test (split by schema)
**Hardware:** 1x L4 24 GB, peak 17.8 GB, 2h10m
**Result:** exact-match SQL 71% vs 52% base; eval loss best at step 600 (kept)
**Artifacts:** out/qwen7b-sql-lora (adapter), out/qwen7b-sql-merged (bf16)
**Versions:** transformers x.y, peft x.y, trl x.y, torch x.y