RLOO Training Test: Pixtral (Vision)¶

# 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 FastVisionModel, is_bf16_supported

import torch
from trl import RLOOConfig, RLOOTrainer
from datasets import Dataset, load_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 Pixtral-12B with FastVisionModel for vision capabilities
MODEL_NAME = "unsloth/pixtral-12b-2409-bnb-4bit"
print(f"\nLoading {MODEL_NAME.split('/')[-1]} with FastVisionModel...")

model, tokenizer = FastVisionModel.from_pretrained(
    MODEL_NAME,
    load_in_4bit=True,
    use_gradient_checkpointing="unsloth",
)

# 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 vision training
model = FastVisionModel.get_peft_model(
    model,
    finetune_vision_layers=True,
    finetune_language_layers=True,
    finetune_attention_modules=True,
    finetune_mlp_modules=True,
    r=16,
    lora_alpha=16,
    lora_dropout=0,
    bias="none",
    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 prompt dataset for RLOO with vision
# Note: RLOO typically uses text prompts - testing if vision model works

prompts = [
    "Describe what you see in the image.",
    "What mathematical expression is shown?",
    "Explain the visual content briefly.",
    "What is the main element in this image?",
    "Describe the pattern or structure shown.",
]

# Format prompts for RLOO (text-only prompts, vision model)
dataset = Dataset.from_dict({
    "prompt": [
        tokenizer.apply_chat_template(
            [{"role": "user", "content": [{"type": "text", "text": p}]}],
            tokenize=False,
            add_generation_prompt=True
        ) for p in prompts
    ]
})

print(f"Dataset created: {len(dataset)} prompts")
print("Note: Using text prompts with vision model (RLOO generates completions)")

# Define a simple reward function for testing
def simple_reward_fn(completions, prompts=None, **kwargs):
    """Simple reward function for testing RLOO with vision model."""
    rewards = []
    for completion in completions:
        length = len(completion.split())
        score = 0.0
        if 10 <= length <= 50:
            score += 1.0
        elif length < 10:
            score -= 0.5
        if completion.strip().endswith("."):
            score += 0.5
        rewards.append(score)
    return rewards

print("Reward function defined: simple_reward_fn")

# Post-training inference test with vision
FastVisionModel.for_inference(model)

# Load a test image for vision inference
vision_dataset = load_dataset("unsloth/LaTeX_OCR", split="train[:1]")
test_image = vision_dataset[0]["image"]

messages = [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": "Describe what you see."}]}]
input_text = tokenizer.apply_chat_template(messages, add_generation_prompt=True)
inputs = tokenizer(test_image, input_text, add_special_tokens=False, return_tensors="pt").to("cuda")

with torch.no_grad():
    output = model.generate(**inputs, max_new_tokens=64, temperature=1.5, min_p=0.1)

response = tokenizer.decode(output[0], skip_special_tokens=True)
# Clean up BPE artifacts from Mistral tokenizer family (Ġ=space, Ċ=newline)
response = response.replace('Ġ', ' ').replace('Ċ', '\n').strip()

# Clear success/failure banner
print("=" * 60)
if RLOO_VISION_SUPPORTED:
    print("RLOO Training: SUPPORTED for Pixtral (Vision)")
    print("Model: FastVisionModel + pixtral-12b-2409-bnb-4bit")
else:
    print("RLOO Training: NOT SUPPORTED for Pixtral (Vision)")
    print("Reason: RLOOTrainer may not have native vision support")
print("=" * 60)
print(f"Sample generation:\n{response[-200:]}")

# 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)