train_qwen_codeforces.py

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

from datasets import load_dataset
from peft import LoraConfig
from trl import SFTTrainer, SFTConfig
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"📊 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
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")