train_qwen_codeforces.py

# /// script
# dependencies = ["trl>=0.12.0", "peft>=0.7.0", "trackio>=0.1.0", "datasets>=2.0.0", "transformers>=4.36.0"]
# ///

from datasets import load_dataset
from peft import LoraConfig
from trl import SFTTrainer, SFTConfig
from transformers import AutoTokenizer
import trackio

# Load dataset - 1000 examples for ~20 min training
print("📦 Loading dataset...")
dataset = load_dataset(
    "open-r1/codeforces-cots",
    "solutions_w_editorials_py_decontaminated",
    split="train[:1000]"
)
print(f"📊 Loaded {len(dataset)} examples")

# Load tokenizer to get chat template
print("🔤 Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B")

# Pre-process dataset - convert messages to text format
print("🔄 Converting messages to text format...")
def convert_messages_to_text(example):
    """Convert messages format to text using chat template."""
    if "messages" in example and example["messages"]:
        text = tokenizer.apply_chat_template(
            example["messages"],
            tokenize=False,
            add_generation_prompt=False
        )
        return {"text": text}
    return {"text": ""}

# Apply the conversion
dataset = dataset.map(convert_messages_to_text, remove_columns=dataset.column_names)
print(f"✅ Dataset preprocessed - training on {len(dataset)} examples for 3 epochs")

# LoRA configuration for efficient training
peft_config = LoraConfig(
    r=8,
    lora_alpha=16,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"]
)

# Training configuration - optimized for T4 small
config = SFTConfig(
    # Hub settings - CRITICAL for saving results
    output_dir="qwen-codeforces-finetuned",
    push_to_hub=True,
    hub_model_id="papebaba/qwen-codeforces-finetuned",
    hub_strategy="end",
    hub_private_repo=False,

    # Training parameters
    num_train_epochs=3,
    per_device_train_batch_size=1,
    gradient_accumulation_steps=8,  # Effective batch size = 8
    learning_rate=2e-4,
    max_length=512,  # Shorter sequences for T4 small

    # Checkpointing
    logging_steps=10,
    save_strategy="epoch",
    save_total_limit=1,

    # Optimization for T4 small
    gradient_checkpointing=True,
    bf16=True,
    max_grad_norm=1.0,
    warmup_ratio=0.1,
    lr_scheduler_type="cosine",
    optim="adamw_torch",

    # Trackio monitoring
    report_to="trackio",
    run_name="qwen-codeforces-sft-1k",
)

# Initialize trainer with preprocessed dataset
print("🎯 Initializing trainer...")
trainer = SFTTrainer(
    model="Qwen/Qwen2.5-0.5B",
    train_dataset=dataset,
    args=config,
    peft_config=peft_config,
)

# Train
print("🚀 Starting training on T4 small...")
trainer.train()

# Push to Hub
print("📤 Pushing final model to Hub...")
trainer.push_to_hub()

print("✅ Training complete!")
print("📊 View metrics at: https://huggingface.co/spaces/papebaba/trackio")
print("🤗 Model at: https://huggingface.co/papebaba/qwen-codeforces-finetuned")