DPO Training Test: Qwen3-4B-Thinking-2507¶

# 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-Thinking-2507 with 4-bit quantization
MODEL_NAME = "unsloth/Qwen3-4B-Thinking-2507-unsloth-bnb-4bit"
print(f"\nLoading {MODEL_NAME.split('/')[-1]}...")

model, tokenizer = FastLanguageModel.from_pretrained(
    MODEL_NAME,
    max_seq_length=1024,  # Increased for thinking content
    load_in_4bit=True,
    dtype=None,
)

# 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 with thinking content (5 samples)
# DPO requires: prompt, chosen response (with quality thinking), rejected response (poor/no thinking)

preference_data = [
    {
        "prompt": "Explain recursion in programming.",
        "chosen": "<think>\nWhat is recursion exactly? It's when a function calls itself. But why would you do that? To break down problems into smaller pieces. What's the key thing users need to understand? The base case - without it you get infinite loops.\n</think>\n\nRecursion 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": "<think>\nHow do I explain API to someone? What does it stand for? Application Programming Interface. But what does that mean practically? It's like a contract between software systems. What's a good analogy? Like a waiter taking orders between you and the kitchen.\n</think>\n\nAn 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": "<think>\nWhat's the core purpose of version control? Tracking changes over time. Why is that useful? You can go back to previous versions, see who changed what. What systems exist? Git is the most popular. How should I frame this?\n</think>\n\nVersion 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": "<think>\nWhat is the essential definition of a database? It stores data, but that's too simple. What makes it different from just files? It's organized and structured. What manages it? A DBMS. How do I make this clear?\n</think>\n\nA 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": "<think>\nWhat are the key concepts of OOP? Objects, classes, encapsulation, inheritance, polymorphism. But what's the core idea? Organizing code around objects that have both data and behavior. How do I explain this simply?\n</think>\n\nObject-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"Chosen responses include <think> blocks with self-questioning")
print(f"Rejected responses have minimal/no reasoning")

# 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=1024,  # Increased to allow full thinking + response
        temperature=0.6,
        top_p=0.95,
        do_sample=True,
        pad_token_id=tokenizer.pad_token_id,
    )

# Get generated token IDs only (exclude prompt)
input_length = inputs["input_ids"].shape[1]
generated_ids = outputs[0][input_length:].tolist()

# Token-based parsing using </think> token ID
THINK_END_TOKEN_ID = 151668

if THINK_END_TOKEN_ID in generated_ids:
    end_idx = generated_ids.index(THINK_END_TOKEN_ID)
    thinking = tokenizer.decode(generated_ids[:end_idx], skip_special_tokens=True).strip()
    final_resp = tokenizer.decode(generated_ids[end_idx + 1:], skip_special_tokens=True).strip()
    think_tag_found = True
else:
    thinking = tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
    final_resp = "(Model did not complete thinking - increase max_new_tokens)"
    think_tag_found = False

print("=" * 60)
print("DPO Training Pipeline Test (Thinking Mode)")
print("=" * 60)
print(f"</think> token found: {'✅ YES' if think_tag_found else '❌ NO'}")
print(f"Output tokens: {len(generated_ids)}")
print(f"\nTHINKING: {thinking[:300]}..." if len(thinking) > 300 else f"\nTHINKING: {thinking}")
print(f"\nRESPONSE: {final_resp[:200]}..." if len(final_resp) > 200 else f"\nRESPONSE: {final_resp}")

if think_tag_found and thinking and final_resp:
    print("\n✅ DPO Training Pipeline Test PASSED")
else:
    print("\n⚠️ Test completed but output may need review")

# Shutdown kernel to release all GPU memory
import IPython
print("Shutting down kernel to release GPU memory...")
app = IPython.Application.instance()
app.kernel.do_shutdown(restart=False)