GRPO Training Test: Ministral (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 GRPOConfig, GRPOTrainer
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 Ministral-3B with FastVisionModel for vision capabilities
MODEL_NAME = "unsloth/Ministral-3-3B-Reasoning-2512"
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 GRPO with vision
# Note: GRPO typically uses text prompts - testing if vision context works
# Using text-only prompts as GRPO generates completions during training

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 GRPO (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(f"Note: Using text prompts with vision model (GRPO generates completions)")

# Define a simple reward function for testing
def simple_reward_fn(completions, prompts=None, **kwargs):
    """Simple reward function for testing GRPO with vision model."""
    rewards = []
    for completion in completions:
        length = len(completion.split())
        if length < 5:
            reward = -1.0
        elif length < 20:
            reward = 0.5
        elif length < 50:
            reward = 1.0
        else:
            reward = 0.8
        rewards.append(reward)
    return rewards

print("Reward function defined: simple_reward_fn")

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