Cognitive Alignment & PEFT Optimization

Overview

• Goal: take base Llama-3.1-8B-Instruct, fine-tune it in two stages (SFT → DPO) so its outputs match a preferred style/persona, then measure if it actually got better.
• Pipeline shape: Setup → Baseline generation → Data prep → SFT → DPO → Evaluation (metric + LLM judge).
• Core idea: SFT teaches the model what good answers look like (imitation). DPO then teaches it to prefer good answers over bad ones when both are shown (preference learning) — DPO alone from a random base model is unstable, so SFT first gives it a sane starting point.
• Everything trains via QLoRA: the base model stays frozen and quantized to 4-bit; only small adapter matrices (LoRA) are trained. This is why an 8B model can train on a single Colab GPU.
• Evaluation is two-pronged: a cheap automatic depth score (no API needed) + a GPT-4 pairwise judge (costs API calls, more reliable) with a bootstrap confidence interval so the win-rate number isn't just a lucky sample.

---

1. Setup

cap = torch.cuda.get_device_capability(0)
USE_BF16 = torch.cuda.get_device_capability(0)[0] >= 8
COMPUTE_DTYPE = torch.bfloat16 if USE_BF16 else torch.float16

• get_device_capability(0)[0] >= 8 checks if the GPU is Ampere or newer (A100, etc.) — these support bf16 natively. Older GPUs fall back to fp16. This decides COMPUTE_DTYPE used everywhere later.
• Why bf16 over fp16: bf16 has the same exponent range as fp32 (less overflow risk during training), fp16 has more precision but overflows/underflows more easily — bf16 is standard for LLM training on modern GPUs.

from google.colab import drive
drive.mount('/content/drive')
ROOT = "/content/drive/MyDrive/xlon_dpo"
for sub in ["data", "adapters", "generations", "logs"]:
    os.makedirs(f"{ROOT}/{sub}", exist_ok=True)

• Drive is mounted so adapters/logs survive Colab disconnects — Colab's local disk is wiped on session reset.

2. Quantization (the "Q" in QLoRA)

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True,
    bnb_4bit_compute_dtype=COMPUTE_DTYPE,
)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID, quantization_config=bnb_config,
    device_map="auto", dtype=COMPUTE_DTYPE,
    attn_implementation="sdpa",
)

• 4-bit: each weight stored in 4 bits instead of 16/32 → ~4-8x smaller memory footprint, letting an 8B model fit on one Colab GPU.
• nf4 (NormalFloat4): a 4-bit format tuned for weights that follow a roughly normal distribution — more accurate than plain int4 at the same bit-width.
• double quant: quantizes the quantization constants themselves too, squeezing out more memory for free.
• compute_dtype: weights are stored in 4-bit, but matrix multiplies are computed in bf16/fp16 — 4-bit is a storage format, not a compute format.

tokenizer.pad_token = "<|finetune_right_pad_id|>"
tokenizer.padding_side = "right"

• Llama-3.1 has a dedicated reserved pad token; using it (instead of reusing EOS as pad, a common hack) avoids the model confusing "end of sequence" with "padding."
• Padding on the right keeps real tokens left-aligned so loss masking lines up correctly during training.

3. Baseline Generation (Cell 7)

GEN_KWARGS = dict(max_new_tokens=512, do_sample=True,
                  temperature=0.7, top_p=0.9,
                  pad_token_id=tokenizer.pad_token_id)

def generate(prompt):
    msgs = [{"role": "user", "content": prompt}]
    inputs = tokenizer.apply_chat_template(
        msgs, add_generation_prompt=True,
        return_tensors="pt", return_dict=True).to(model.device)
    with torch.no_grad():
        out = model.generate(**inputs, **GEN_KWARGS)
    return tokenizer.decode(out[0][inputs["input_ids"].shape[1]:],
                            skip_special_tokens=True)

torch.manual_seed(0)
with open(f"{ROOT}/generations/baseline.jsonl", "w") as f:
    for p in eval_prompts:
        f.write(json.dumps({"prompt": p, "model": "base",
                            "response": generate(p)}) + "\n")

• Before touching the model, generate answers to the held-out eval prompts with the untouched base model and save them (baseline.jsonl).
• Why: without a baseline, "the DPO model sounds good" is meaningless — you need a fixed reference point to prove the fine-tune actually changed something.
• torch.manual_seed(0): makes sampling reproducible across baseline vs. later DPO generation runs, so any difference is due to the model changing, not seed luck.
• apply_chat_template(..., add_generation_prompt=True): formats the raw prompt into Llama's chat format and appends the "assistant, start talking now" marker — required because Llama-3.1-Instruct was trained on this specific chat format, not raw text.

4. Data Format

dpo_ds = load_dataset("json", data_files=f"{ROOT}/data/train_pairs.jsonl", split="train")
# prompt / chosen / rejected  -> for DPO

sft_ds = dpo_ds.map(lambda r: {
    "prompt":     [{"role": "user",      "content": r["prompt"]}],
    "completion": [{"role": "assistant", "content": r["chosen"]}],
}, remove_columns=dpo_ds.column_names)

• Source data: train_pairs.jsonl — each row has prompt, chosen, rejected. This triplet format is exactly what DPO needs directly.
• SFT doesn't use preference pairs — it just needs "good examples to imitate," so the same data is reshaped by dropping rejected, keeping only prompt + chosen in a prompt/completion split (TRL's preferred SFT format over a flat messages list).
• Net effect: one dataset, two derived views — full triplets for DPO, chosen-only pairs for SFT. This is why SFT must run first: it's built from the same preference data, just the "positive examples" half of it.

5. Stage 2 — SFT (Supervised Fine-Tuning)

model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=True)
model.config.use_cache = False

peft_config = LoraConfig(
    r=16, lora_alpha=32, lora_dropout=0.05, bias="none",
    task_type="CAUSAL_LM",
    target_modules=["q_proj","k_proj","v_proj","o_proj",
                    "gate_proj","up_proj","down_proj"],
)

• prepare_model_for_kbit_training: bookkeeping needed to train a quantized model — casts norm layers to fp32 for stability, enables gradient checkpointing, makes input embeddings require grad.
• use_cache = False: KV-cache speeds up generation (reusing past key/values) but must be off during training — training needs gradients through the full sequence.
• LoRA config — freezes base weights, injects small trainable low-rank matrices next to them:
  • r=16: rank (bottleneck size) of the adapter matrices — higher r = more trainable capacity, more memory.
  • lora_alpha=32: scaling factor on the LoRA update (effective scale = alpha/r); alpha = 2×r is a common default ratio.
  • target_modules: adapters applied to all attention projections (q,k,v,o_proj) and all MLP projections (gate,up,down_proj) [ These are weight matrices inside each transformer block that LoRA attaches its small trainable adapters to.]— broad coverage across the transformer block, not just attention.

    • MLP block (3) — Llama uses a "gated" MLP (SwiGLU), so it's 3 matrices instead of the usual 2:

      gate_proj — computes a gate that decides how much of each signal gets through

      up_proj — expands the hidden size up to a bigger intermediate size

      down_proj — projects back down from that intermediate size to the normal hidden size, after the gate is applied

sft_args = SFTConfig(
    output_dir=f"{ROOT}/adapters/sft",
    num_train_epochs=3, per_device_train_batch_size=4,
    gradient_accumulation_steps=2, learning_rate=2e-4,
    lr_scheduler_type="cosine", warmup_ratio=0.1,
    max_length=1024, bf16=True, logging_steps=1,
    save_strategy="no", report_to="none",
)
trainer = SFTTrainer(model=model, args=sft_args,
                      train_dataset=sft_ds, peft_config=peft_config)
trainer.train()
trainer.save_model(f"{ROOT}/adapters/sft")

• SFTTrainer(model, args, train_dataset, peft_config): passing peft_config here means the trainer wraps the base model into a PeftModel internally — no separate get_peft_model call needed.
• LR 2e-4 (typical for LoRA — LoRA tolerates higher LR than full fine-tuning since only small matrices move), cosine schedule with warmup, 3 epochs, max_length=1024 (longer sequences get truncated).

model.config.use_cache = True
model.eval()
print(generate(eval_prompts[0])[:600])

• Immediately eyeballs one post-SFT generation as a sanity check before moving to DPO — cheap way to catch a broken run early.

6. Stage 3 — DPO (Direct Preference Optimization)

sft_model = trainer.model          # the PEFT-wrapped model from SFT
sft_model.config.use_cache = False
sft_model.train()

• Continuing training on the same PEFT-wrapped model that came out of SFT (not reloading from scratch) — DPO adapters build directly on top of the SFT adapters.

dpo_args = DPOConfig(
    output_dir=f"{ROOT}/adapters/dpo",
    beta=0.1, num_train_epochs=3,
    per_device_train_batch_size=2, gradient_accumulation_steps=4,
    learning_rate=1e-5, lr_scheduler_type="cosine",
    warmup_ratio=0.1, max_length=1024, bf16=True,
    logging_steps=1, save_strategy="no", report_to="none",
)
trainer = DPOTrainer(model=sft_model, ref_model=None,
                      args=dpo_args, train_dataset=dpo_ds)
trainer.train()
trainer.save_model(f"{ROOT}/adapters/dpo")

• What DPO optimizes (conceptually): given a (prompt, chosen, rejected) triplet, DPO increases the model's log-probability of chosen relative to rejected, compared against a frozen reference model's log-probabilities for the same pair — this ratio-of-ratios becomes the loss, no separate reward model needed (that's the "direct" in DPO).
• beta=0.1: controls how hard the model is pushed toward the preference vs. how much it's kept close to the reference model — lower beta = more aggressive preference-fitting, higher beta = more conservative. 0.1 is a standard starting value.
• ref_model=None: normally DPO needs a frozen copy of the pre-DPO model for reference log-probs. With PEFT/LoRA, None tells TRL to reuse the same model with adapters temporarily disabled as the reference — avoids loading a second full model copy, saving GPU memory.
• LR 1e-5 (10x smaller than SFT's 2e-4): DPO is more sensitive/unstable than SFT, needs gentler updates — a standard SFT-vs-DPO LR ratio.
• Batch size drops to 2 (vs. 4 for SFT) with more gradient accumulation (4 vs. 2) — DPO needs both chosen and rejected sequences per example in memory at once, roughly doubling per-sample memory cost.

torch.manual_seed(0)
with open(f"{ROOT}/generations/dpo.jsonl", "w") as f:
    for p in eval_prompts:
        f.write(json.dumps({"prompt": p, "model": "dpo",
                            "response": generate(p)}) + "\n")

• Same held-out prompts, same seed as the baseline generation — apples-to-apples comparison.

7. Stage 4 — Evaluation

7a. Depth score (cheap, no API, deterministic)

emb = SentenceTransformer("all-mpnet-base-v2")

def depth_score(text):
    sents = [s for s in re.split(r"(?<=[.!?])\s+", text.strip())
             if len(s.split()) > 3]
    if len(sents) < 2: return None
    v = emb.encode(sents, normalize_embeddings=True)
    return float(np.mean((v[:-1] * v[1:]).sum(axis=1)))

• Splits each response into sentences (ignoring very short ones), embeds each with all-mpnet-base-v2 (a general-purpose sentence embedding model), averages the cosine similarity between each sentence and the next one.
• Interpretation: high score → consecutive sentences are semantically close → the answer stays on one throughline. Low score → sentences jump between unrelated ideas → "parallel facts" style rather than sustained reasoning.
• Proxy metric standing in for "sustained attentional depth" — cheap and repeatable, but a heuristic, not ground truth quality.

7b. GPT-4 pairwise judge (Cells 18–19)

RUBRIC = """You are comparing two answers to the same question.
Judge ONLY this: does the answer sustain one line of reasoning...
Ignore length. Ignore writing polish. Ignore how many facts are included.
Reply with exactly one character: A or B."""

def judge(q, a, b):
    r = client.chat.completions.create(
        model="gpt-4o", temperature=0, max_tokens=1,
        messages=[{"role": "system", "content": RUBRIC},
                  {"role": "user", "content":
                   f"Question: {q}\n\nAnswer A:\n{a}\n\nAnswer B:\n{b}"}])
    return r.choices[0].message.content.strip().upper()

for b_row, d_row in zip(base, dpo):
    v1 = judge(q, b_row["response"], d_row["response"])   # base first
    v2 = judge(q, d_row["response"], b_row["response"])   # dpo first
    results.append({"prompt": q,
        "dpo_wins": (v1 == "B") + (v2 == "A"),
        "consistent": (v1 == "B") == (v2 == "A")})

• Rubric isolates one axis only: sustained reasoning vs. drift/listing — explicitly told to ignore length/polish so the judge can't just reward "sounds more fluent" or "longer."
• Each pair judged twice, order swapped: LLM judges have a known position bias (favoring whichever answer is shown first/second) — judging both orders and checking agreement (consistent) controls for this rather than trusting one judgment.
• dpo_wins per prompt is 0, 1, or 2 (how many of the two judgments favored DPO).
• temperature=0, max_tokens=1: judge forced to be as deterministic as possible, outputs a single character — minimizes noise and cost per call.

n = len(results)
wins  = sum(r["dpo_wins"] == 2 for r in results)
loss  = sum(r["dpo_wins"] == 0 for r in results)
ties  = n - wins - loss
print(f"DPO wins {wins}/{n}, loses {loss}, judge inconsistent {ties}")

• A "true win" only counts prompts where DPO won both orderings (dpo_wins == 2) — a stricter, bias-resistant win condition than a simple majority vote.

7c. Bootstrap confidence interval

scores = np.array([r["dpo_wins"] for r in results]) / 2
boot = [np.mean(np.random.choice(scores, n, replace=True)) for _ in range(5000)]
lo, hi = np.percentile(boot, [2.5, 97.5])
print(f"win rate {scores.mean():.0%}  (95% CI {lo:.0%} – {hi:.0%})")

• A single win-rate number (e.g. "DPO wins 70%") hides how much it could shift with a different sample of prompts.

• Bootstrap resampling: repeatedly resample the existing results with replacement (5000 times), recompute the win rate each time, take the 2.5th/97.5th percentile of that distribution as a 95% CI : "We are 95% confident that DPO's true win rate lies somewhere between 50% and 94%."

• Getting to a reasonable expected DPO-win number :

  1. The Core Problem: Small Samples Lie

  Imagine you flip a coin 8 times and it lands on Heads 6 times (75%). Does that prove it's a biased coin? Not necessarily—it could easily just be a lucky streak.

  The same thing happens when you evaluate an AI model on only 50 prompts:

  • Raw Win Rate (75%): "Out of 8 matches scored, DPO won 5 times, lost 1, and 2 were inconsistent."
  • The Reality: 8 matches (or even 50 prompts) is a tiny sample size. If you tested it on 50 different prompts tomorrow, DPO might win 50% of the time, or 90% of the time.

  2. What is "Bootstrap Resampling"?

  Since you don't have time to create 5,000 brand-new test prompts from scratch, bootstrapping is a statistical trick that simulates running the experiment 5,000 times using the data you already have.

  How the code does it:

  1. Put all your scores in a hat (e.g., your actual win/loss results).
  2. Pull scores out randomly one by one, write them down, and put them back in the hat (this is called "sampling with replacement").
  3. Do this until you have a fake batch of results, and compute the win rate.
  4. Repeat that whole process 5,000 times!

  You end up with a list of 5,000 slightly different win rates (e.g., 62%, 75%, 88%, 50%...).

  3. What is the 95% Confidence Interval (CI)?

  Once you have those 5,000 simulated win rates, you line them up from lowest to highest and chop off the extreme 2.5% at the bottom and 2.5% at the top.

  What remains in the middle is your 95% Confidence Interval:

  95% CI 50% – 94%

  "We are 95% confident that DPO's true win rate lies somewhere between 50% and 94%."

  Summary: Why This Matters for Your AI Blog/Paper

  • A raw number lies: Saying "DPO achieved a 75% win rate!" sounds incredible.
  • The CI reveals the truth: Showing (95% CI 50% – 94%) proves you are honest enough to admit: "Look, our sample size was small. DPO is definitely better than baseline (at least 50%), but it could be anywhere from slightly better to dominant."

  It turns a simple metric into statistically meaningful evidence rather than a lucky fluke!

DPO wins 5/8, loses 1, judge inconsistent 2 judge self-consistency: 75% win rate 75% 

(95% CI : 50% – 94%)

• Why it matters: with only ~50 eval prompts, a raw win rate can look impressive but have a wide CI — the CI tells you if the result is statistically meaningful or noise.

---