DPO Training Test: Qwen3-4B¶

# Environment Setup
import os
from dotenv import load_dotenv
load_dotenv()

# Force text-based progress instead of HTML widgets
os.environ["TQDM_NOTEBOOK"] = "false"

# CRITICAL: Import unsloth FIRST for proper TRL patching
import unsloth
from unsloth import FastLanguageModel, is_bf16_supported

import torch
from trl import DPOConfig, DPOTrainer
from datasets import Dataset

# Environment summary
gpu = torch.cuda.get_device_name(0) if torch.cuda.is_available() else "CPU"
print(f"Environment: unsloth {unsloth.__version__}, PyTorch {torch.__version__}, {gpu}")
print(f"HF_TOKEN loaded: {'Yes' if os.environ.get('HF_TOKEN') else 'No'}")

# Load Qwen3-4B with 4-bit quantization
MODEL_NAME = "unsloth/Qwen3-4B-unsloth-bnb-4bit"
print(f"\nLoading {MODEL_NAME.split('/')[-1]}...")

model, tokenizer = FastLanguageModel.from_pretrained(
    MODEL_NAME,
    max_seq_length=512,
    load_in_4bit=True,
    dtype=None,  # Auto-detect
)

# Ensure pad token is set
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.pad_token_id = tokenizer.eos_token_id

print(f"Model loaded: {type(model).__name__}")

# Apply LoRA adapters for DPO training
model = FastLanguageModel.get_peft_model(
    model,
    r=16,
    lora_alpha=16,
    lora_dropout=0,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                    "gate_proj", "up_proj", "down_proj"],
    bias="none",
    use_gradient_checkpointing="unsloth",
    random_state=42,
)

trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
total = sum(p.numel() for p in model.parameters())
print(f"LoRA applied: {trainable:,} trainable / {total:,} total ({100*trainable/total:.2f}%)")

# Create minimal synthetic preference dataset (5 samples)
# DPO requires: prompt, chosen response, rejected response

preference_data = [
    {
        "prompt": "Explain recursion in programming.",
        "chosen": "Recursion is when a function calls itself with a simpler version of the problem, including a base case to stop infinite loops.",
        "rejected": "Recursion is just loops."
    },
    {
        "prompt": "What is an API?",
        "chosen": "An API (Application Programming Interface) is a set of protocols that allows different software applications to communicate with each other.",
        "rejected": "API is code."
    },
    {
        "prompt": "Describe version control.",
        "chosen": "Version control is a system that records changes to files over time, allowing you to recall specific versions and collaborate with others.",
        "rejected": "Version control saves files."
    },
    {
        "prompt": "What is a database?",
        "chosen": "A database is an organized collection of structured data stored electronically, typically managed by a database management system (DBMS).",
        "rejected": "A database stores stuff."
    },
    {
        "prompt": "Explain object-oriented programming.",
        "chosen": "Object-oriented programming (OOP) is a paradigm that organizes code into objects containing data (attributes) and behavior (methods).",
        "rejected": "OOP uses objects."
    },
]

# Format for DPO
def format_for_dpo(sample):
    prompt = tokenizer.apply_chat_template(
        [{"role": "user", "content": sample["prompt"]}],
        tokenize=False,
        add_generation_prompt=True
    )
    return {
        "prompt": prompt,
        "chosen": sample["chosen"],
        "rejected": sample["rejected"],
    }

dataset = Dataset.from_list(preference_data)
dataset = dataset.map(format_for_dpo)

print(f"Dataset created: {len(dataset)} preference pairs")
print(f"Sample prompt: {dataset[0]['prompt'][:80]}...")

# Post-training inference test
FastLanguageModel.for_inference(model)

test_prompt = "What is machine learning?"
messages = [{"role": "user", "content": test_prompt}]
prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

with torch.no_grad():
    outputs = model.generate(
        **inputs,
        max_new_tokens=64,
        temperature=0.7,
        top_p=0.9,
        do_sample=True,
        pad_token_id=tokenizer.pad_token_id,
    )

response = tokenizer.decode(outputs[0], skip_special_tokens=True)
print("=" * 60)
print("DPO Training Pipeline Test PASSED")
print("=" * 60)
print(f"Sample generation:\n{response[-200:]}")