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

# Environment Setup
import os

# FIX: Set ACCELERATE_MIXED_PRECISION BEFORE importing unsloth
os.environ['ACCELERATE_MIXED_PRECISION'] = 'bf16'

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 RLOOConfig, RLOOTrainer
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"ACCELERATE_MIXED_PRECISION: {os.environ.get('ACCELERATE_MIXED_PRECISION', 'not set')}")
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 RLOO 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 prompt dataset for RLOO (5 prompts)
# RLOO requires prompts only - completions are generated during training

prompts = [
    "Explain the concept of recursion in programming.",
    "What are the benefits of using version control?",
    "Describe how a hash table works.",
    "What is the difference between a stack and a queue?",
    "Explain what an API is to a beginner.",
]

# Format prompts for RLOO (requires "prompt" field)
dataset = Dataset.from_dict({
    "prompt": [
        tokenizer.apply_chat_template(
            [{"role": "user", "content": p}],
            tokenize=False,
            add_generation_prompt=True
        ) for p in prompts
    ]
})

print(f"Dataset created: {len(dataset)} prompts")
print(f"Sample prompt:\n{dataset[0]['prompt'][:150]}...")

# Define thinking-aware reward function using token IDs
# TRL passes completion_ids directly - no re-tokenization needed!

THINK_END_TOKEN_ID = 151668  # </think> token for Qwen3-Thinking models

def thinking_reward_fn(completions, prompts=None, completion_ids=None, **kwargs):
    """
    Token-based reward function using completion_ids provided by TRL.
    
    - Uses token ID 151668 for </think> boundary detection
    - Rewards longer, more detailed reasoning (measured in tokens)
    - Bonus for self-questioning (question marks in thinking content)
    
    Benefits over string matching:
    - No re-tokenization overhead (faster training)
    - Exact token boundaries (no regex edge cases)
    - Consistent with inference code pattern
    """
    rewards = []
    
    for completion, comp_ids in zip(completions, completion_ids):
        # Token-based detection: check for </think> token (ID 151668)
        if THINK_END_TOKEN_ID in comp_ids:
            end_idx = comp_ids.index(THINK_END_TOKEN_ID)
            thinking_length = end_idx  # Token count before </think>
            
            # String-based content analysis for question detection
            # (using string here is fine since we already know boundary from tokens)
            thinking_content = completion.split('</think>')[0]
            question_marks = thinking_content.count('?')
            has_self_questions = question_marks >= 1
            
            # Reward based on thinking token count
            if thinking_length < 10:
                reward = 0.3  # Minimal thinking
            elif thinking_length < 30:
                reward = 0.7 + (0.1 if has_self_questions else 0)
            else:
                reward = 1.0 + (0.1 if has_self_questions else 0)
        else:
            reward = -1.0  # No </think> token found
        
        rewards.append(reward)
    
    return rewards

print("Token-based thinking reward function defined")
print(f"Using THINK_END_TOKEN_ID = {THINK_END_TOKEN_ID}")
print("Rewards: thinking quality (token count) + self-questioning bonus")

# 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("RLOO 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✅ RLOO 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)