Fine-tuning LLMs
from base to aligned

# ── Step 1: Quantize base model to 4-bit NF4 ────────────────
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
import torch

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",              # NormalFloat4 — optimal for NN weights
    bnb_4bit_compute_dtype=torch.float16    # adapter forward pass in fp16
)
model = AutoModelForCausalLM.from_pretrained(
    "TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T",
    quantization_config=bnb_config, device_map="auto"
)
model = prepare_model_for_kbit_training(model)  # cast LayerNorms to fp32

# ── Step 2: Attach LoRA₁ adapter ────────────────────────────
# Only trains B and A matrices (r×d and d×r) — everything else frozen
lora_config = LoraConfig(
    r=8,               # rank — adapter capacity (higher = more params)
    lora_alpha=16,    # scaling = alpha/r = 2.0 (lr multiplier effect)
    target_modules=["q_proj", "v_proj"],  # query + value attention layers
    lora_dropout=0.05,
    task_type="CAUSAL_LM"
)
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()
# trainable: 2,097,152 || all: 1,100,048,384 || trainable%: 0.1907

# ── Step 3: Tokenize raw text and train ─────────────────────
from transformers import Trainer, TrainingArguments, DataCollatorForLanguageModeling

tokenizer = AutoTokenizer.from_pretrained("TinyLlama/TinyLlama-1.1B-...")
tokenizer.pad_token = tokenizer.eos_token

trainer = Trainer(
    model=model,
    train_dataset=tokenized_dataset,          # raw text, 512-token chunks
    data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
    args=TrainingArguments(
        num_train_epochs=2,
        per_device_train_batch_size=4,
        learning_rate=2e-4, fp16=True,
        output_dir="./stage1_output"
    )
)
trainer.train()

# ── Step 4: Merge LoRA₁ into base → save as M1 ──────────────
# W_new = W_base + B·A  (the rank-decomposed update is baked in)
model = model.merge_and_unload()
model.save_pretrained("pharma_tinyllama_noninstruct_merged")   # ← M1
tokenizer.save_pretrained("pharma_tinyllama_noninstruct_merged")

from trl import SFTTrainer, SFTConfig
from datasets import load_dataset

# ── Load M1 (Stage 1 merged model) as the base for Stage 2 ──
model = AutoModelForCausalLM.from_pretrained(
    "pharma_tinyllama_noninstruct_merged",   # ← M1
    quantization_config=bnb_config, device_map="auto"
)
model = prepare_model_for_kbit_training(model)
model = get_peft_model(model, lora_config)    # ← LoRA₂, fresh adapter on M1

# ── Alpaca prompt formatter ───────────────────────────────────
def format_alpaca(example):
    return (
        f"### Instruction:\n{example['instruction']}\n\n"
        f"### Input:\n{example['input']}\n\n"
        f"### Response:\n{example['output']}"
    )

# ── Train — SFTTrainer masks instruction tokens automatically ─
instruction_ds = load_dataset("json",
    data_files="pharma_instruction_dataset.jsonl")

trainer = SFTTrainer(
    model=model,
    train_dataset=instruction_ds["train"],
    formatting_func=format_alpaca,
    args=SFTConfig(
        num_train_epochs=3,
        per_device_train_batch_size=4,
        learning_rate=2e-4,
        max_seq_length=512,
        fp16=True
    )
)
trainer.train()

# ── Merge LoRA₂ into M1 → save as M2 ────────────────────────
model = model.merge_and_unload()
model.save_pretrained("pharma_tinyllama_instruct_merged")   # ← M2
tokenizer.save_pretrained("pharma_tinyllama_instruct_merged")

from trl import DPOTrainer, DPOConfig

# ── Load M2 (Stage 2 merged model) as the base for Stage 3 ──
model = AutoModelForCausalLM.from_pretrained(
    "pharma_tinyllama_instruct_merged",   # ← M2
    quantization_config=bnb_config, device_map="auto"
)
model = prepare_model_for_kbit_training(model)
model = get_peft_model(model, lora_config)   # ← LoRA₃, fresh adapter on M2

# ── Load preference dataset (prompt / chosen / rejected) ─────
pref_ds = load_dataset("json",
    data_files="pharma_preference_dataset.jsonl").train_test_split(0.1)
# 48 rows in demo — DPOTrainer expects exactly these three column names

# ── DPO training — no reward model needed ───────────────────
trainer = DPOTrainer(
    model=model,
    ref_model=None,             # None = implicit reference from frozen base
    train_dataset=pref_ds["train"],
    eval_dataset=pref_ds["test"],
    args=DPOConfig(
        beta=0.1,               # KL penalty — lower = more alignment aggression
        num_train_epochs=3,
        per_device_train_batch_size=4,
        learning_rate=1e-4,     # lower LR than SFT for fine-grained tuning
        max_prompt_length=256,
        max_length=512,
        fp16=True
    )
)
trainer.train()

# ── Merge LoRA₃ into M2 → M3 (final aligned model) ──────────
model = model.merge_and_unload()
model.save_pretrained("pharma_tinyllama_preference_dpo_merged")  # ← M3
model.push_to_hub("your-hf-username/pharma-tinyllama-dpo")    # optional