olympus/train_code_specialist.py

"""
Train code specialist: SmolLM3-3B + QLoRA on code data.

Run on GPU pod:
    python olympus/train_code_specialist.py

Expects SmolLM3-3B and datasets already cached from download_all.py.
"""

import torch
import time
import os
import json
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    TrainingArguments,
    Trainer,
    DataCollatorForLanguageModeling,
    BitsAndBytesConfig,
)
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from datasets import load_dataset, concatenate_datasets


def format_code_example(example):
    """Format CodeAlpaca/CodeFeedback examples as instruction-response pairs."""
    # Different datasets have different field names
    instruction = example.get("instruction", example.get("query", example.get("prompt", "")))
    output = example.get("output", example.get("answer", example.get("response", "")))
    inp = example.get("input", "")

    if inp and inp.strip():
        text = f"### Instruction:\n{instruction}\n\n### Input:\n{inp}\n\n### Response:\n{output}"
    else:
        text = f"### Instruction:\n{instruction}\n\n### Response:\n{output}"

    return {"text": text}


def main():
    t_start = time.time()
    model_id = "HuggingFaceTB/SmolLM3-3B"

    print("=" * 60)
    print("  OLYMPUS CODE SPECIALIST — QLoRA Training")
    print("=" * 60)

    # Load tokenizer
    print("\nLoading tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    # Load model in fp16 (4-bit has conversion issues with SmolLM3)
    # fp16 uses 6.2GB VRAM, leaves ~10GB for LoRA training on 16GB GPU
    print("Loading SmolLM3-3B in fp16...")
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        device_map="auto",
        torch_dtype=torch.float16,
    )

    # Freeze base model, only train LoRA adapters
    for param in model.parameters():
        param.requires_grad = False

    n_params = sum(p.numel() for p in model.parameters())
    print(f"  Base model: {n_params/1e9:.1f}B params")

    # Apply LoRA
    print("Applying LoRA (r=16)...")
    lora_config = LoraConfig(
        r=16,
        lora_alpha=32,
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                         "gate_proj", "up_proj", "down_proj"],
        lora_dropout=0.05,
        bias="none",
        task_type="CAUSAL_LM",
    )

    model = get_peft_model(model, lora_config)
    model.print_trainable_parameters()

    # Load and prepare data
    print("\nLoading code datasets...")
    ds1 = load_dataset("sahil2801/CodeAlpaca-20k", split="train")
    ds2 = load_dataset("m-a-p/CodeFeedback-Filtered-Instruction", split="train")

    print(f"  CodeAlpaca: {len(ds1)} examples")
    print(f"  CodeFeedback: {len(ds2)} examples")

    # Format and combine
    ds1 = ds1.map(format_code_example, remove_columns=ds1.column_names)
    ds2 = ds2.map(format_code_example, remove_columns=ds2.column_names)
    dataset = concatenate_datasets([ds1, ds2]).shuffle(seed=42)

    # Take a subset for faster training (can increase later)
    max_examples = 50000
    if len(dataset) > max_examples:
        dataset = dataset.select(range(max_examples))
    print(f"  Combined: {len(dataset)} examples (capped at {max_examples})")

    # Tokenize
    print("Tokenizing...")
    def tokenize(examples):
        return tokenizer(
            examples["text"],
            truncation=True,
            max_length=1024,
            padding=False,
        )

    dataset = dataset.map(tokenize, batched=True, remove_columns=["text"])

    # Split
    split = dataset.train_test_split(test_size=0.02, seed=42)
    train_dataset = split["train"]
    eval_dataset = split["test"]
    print(f"  Train: {len(train_dataset)}, Eval: {len(eval_dataset)}")

    # Training arguments
    output_dir = "/workspace/olympus_code_specialist"
    os.makedirs(output_dir, exist_ok=True)

    training_args = TrainingArguments(
        output_dir=output_dir,
        num_train_epochs=2,
        per_device_train_batch_size=2,
        per_device_eval_batch_size=2,
        gradient_accumulation_steps=8,  # effective batch = 16
        learning_rate=2e-4,
        lr_scheduler_type="cosine",
        warmup_ratio=0.05,
        weight_decay=0.01,
        logging_steps=50,
        eval_strategy="steps",
        eval_steps=500,
        save_strategy="steps",
        save_steps=500,
        save_total_limit=3,
        bf16=True,
        gradient_checkpointing=True,
        report_to="none",
        dataloader_num_workers=2,
        max_grad_norm=1.0,
    )

    # Data collator
    data_collator = DataCollatorForLanguageModeling(
        tokenizer=tokenizer,
        mlm=False,
    )

    # Trainer
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        data_collator=data_collator,
    )

    # Train
    print(f"\nStarting training...")
    print(f"  Epochs: {training_args.num_train_epochs}")
    print(f"  Batch: {training_args.per_device_train_batch_size} x {training_args.gradient_accumulation_steps} = {training_args.per_device_train_batch_size * training_args.gradient_accumulation_steps}")
    print(f"  LR: {training_args.learning_rate}")
    print(f"  Output: {output_dir}")

    trainer.train()

    # Save final model
    print("\nSaving final model...")
    trainer.save_model(os.path.join(output_dir, "final"))
    tokenizer.save_pretrained(os.path.join(output_dir, "final"))

    # Evaluate with generation
    print("\n" + "=" * 60)
    print("  CODE GENERATION TEST")
    print("=" * 60)

    model.eval()
    test_prompts = [
        "### Instruction:\nWrite a Python function that sorts a list using bubble sort.\n\n### Response:\n",
        "### Instruction:\nWrite a Python function to check if a number is prime.\n\n### Response:\n",
        "### Instruction:\nWrite a Python function that reverses a string.\n\n### Response:\n",
        "### Instruction:\nWrite a Python class for a binary search tree with insert and search methods.\n\n### Response:\n",
        "### Instruction:\nWrite a JavaScript function that debounces another function.\n\n### Response:\n",
    ]

    for prompt in test_prompts:
        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=256,
                temperature=0.3,
                top_p=0.9,
                do_sample=True,
            )
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        # Extract just the response part
        if "### Response:" in response:
            response = response.split("### Response:")[-1].strip()
        print(f"\nQ: {prompt.split('Instruction:')[1].split('Response:')[0].strip()}")
        print(f"A: {response[:300]}")
        print("-" * 40)

    elapsed = time.time() - t_start
    print(f"\n  Total time: {elapsed/3600:.1f} hours")
    print(f"  Output: {output_dir}/final")

    # Save metrics
    metrics = {
        "specialist": "code",
        "base_model": model_id,
        "training_hours": elapsed / 3600,
        "train_examples": len(train_dataset),
        "eval_examples": len(eval_dataset),
        "lora_rank": 16,
        "epochs": training_args.num_train_epochs,
    }
    with open(os.path.join(output_dir, "metrics.json"), "w") as f:
        json.dump(metrics, f, indent=2)


if __name__ == "__main__":
    main()