train_fable.py

#!/usr/bin/env python3
"""
train_fable.py — QLoRA fine-tune HPC-injected Ornith on FABLE 5 traces

Loads the HPC-injected model (Ornith-1.0-9B-hpc), adds LoRA adapters,
and fine-tunes on FABLE 5 assistant conversations.

Usage:
  python3 train_fable.py --model ./Ornith-1.0-9B-hpc \\
                         --data /tmp/fable5_sft.jsonl \\
                         --output ./Ornith-1.0-9B-fable \\
                         --epochs 1 \\
                         --lr 2e-4
"""

import argparse, json, gc, math, os, sys, time
from functools import partial

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    BitsAndBytesConfig,
    get_linear_schedule_with_warmup,
)
from peft import LoraConfig, get_peft_model


# ── Dataset ──────────────────────────────────────────────────────────────────

class FableDataset(Dataset):
    """Tokenized FABLE 5 assistant conversations."""

    def __init__(self, data_path, tokenizer, max_length=2048):
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.samples = []

        # System prompt used by Claude-style models
        system_msg = "You are a helpful, harmless, and honest assistant."

        with open(data_path) as f:
            for line in f:
                d = json.loads(line)
                text = d.get("text", "")
                if not text:
                    continue

                # Reconstruct structured conversation
                turns = text.split("<|im_start|>")
                messages = [{"role": "system", "content": system_msg}]
                for turn in turns:
                    turn = turn.strip()
                    if not turn:
                        continue
                    if turn.startswith("user\n"):
                        messages.append({"role": "user", "content": turn[len("user\n"):].replace("<|im_end|>", "").strip()})
                    elif turn.startswith("assistant\n"):
                        messages.append({"role": "assistant", "content": turn[len("assistant\n"):].replace("<|im_end|>", "").strip()})

                if len(messages) <= 1:
                    continue

                # Format with chat template
                formatted = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)
                tokens = tokenizer.encode(formatted, add_special_tokens=False, truncation=True, max_length=max_length)
                self.samples.append(tokens)

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx):
        tokens = self.samples[idx]
        return torch.tensor(tokens, dtype=torch.long)


def collate_fn(batch, pad_token_id):
    """Pad batch to uniform length."""
    max_len = max(len(x) for x in batch)
    padded = torch.full((len(batch), max_len), pad_token_id, dtype=torch.long)
    for i, seq in enumerate(batch):
        padded[i, :len(seq)] = seq
    return padded


# ── Training ─────────────────────────────────────────────────────────────────

def train():
    parser = argparse.ArgumentParser(description="Fine-tune HPC-injected Ornith on FABLE 5")
    parser.add_argument("--model", default="./Ornith-1.0-9B-hpc", help="Injected model path")
    parser.add_argument("--data", default="/tmp/fable5_sft.jsonl", help="FABLE 5 JSONL path")
    parser.add_argument("--output", default="./Ornith-1.0-9B-fable", help="Output path")
    parser.add_argument("--epochs", type=int, default=1, help="Training epochs")
    parser.add_argument("--lr", type=float, default=2e-4, help="Peak learning rate")
    parser.add_argument("--batch_size", type=int, default=1, help="Per-device batch size")
    parser.add_argument("--grad_accum", type=int, default=8, help="Gradient accumulation steps")
    parser.add_argument("--max_length", type=int, default=2048, help="Max sequence length")
    parser.add_argument("--lora_r", type=int, default=16, help="LoRA rank")
    parser.add_argument("--lora_alpha", type=int, default=32, help="LoRA alpha")
    parser.add_argument("--lora_dropout", type=float, default=0.05, help="LoRA dropout")
    parser.add_argument("--save_steps", type=int, default=200, help="Checkpoint interval (steps)")
    args = parser.parse_args()

    t0 = time.time()

    # ── 1. Tokenizer & 4-bit model ──
    print("[1/6] Loading tokenizer & 4-bit model...")
    tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True, use_fast=False)
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "right"
    if tokenizer.chat_template is None:
        tokenizer.chat_template = "{% for message in messages %}{% if message['role'] == 'system' %}<|im_start|>system\n{{ message['content'] }}<|im_end|>\n{% elif message['role'] == 'user' %}<|im_start|>user\n{{ message['content'] }}<|im_end|>\n{% elif message['role'] == 'assistant' %}<|im_start|>assistant\n{{ message['content'] }}<|im_end|>\n{% endif %}{% endfor %}{% if add_generation_prompt %}<|im_start|>assistant\n{% endif %}"

    bnb = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_use_double_quant=True,
        bnb_4bit_compute_dtype=torch.bfloat16,
    )
    # Leave ~3 GiB headroom on GPU for activations/gradients
    max_memory = {0: f"{torch.cuda.get_device_properties(0).total_memory // (1024**3) - 3}GiB", "cpu": "64GiB"}
    model = AutoModelForCausalLM.from_pretrained(
        args.model,
        trust_remote_code=True,
        quantization_config=bnb,
        device_map="auto",
        max_memory=max_memory,
        torch_dtype=torch.bfloat16,
        low_cpu_mem_usage=True,
    )
    model.config.use_cache = False  # required for gradient checkpointing
    model.gradient_checkpointing_enable()

    # ── 2. LoRA config — target gate_proj (the injected weights) ──
    print(f"[2/6] Adding LoRA (r={args.lora_r}, alpha={args.lora_alpha})...")
    lora_config = LoraConfig(
        r=args.lora_r,
        lora_alpha=args.lora_alpha,
        lora_dropout=args.lora_dropout,
        bias="none",
        task_type="CAUSAL_LM",
        target_modules=["gate_proj", "up_proj", "down_proj"],
    )
    model = get_peft_model(model, lora_config)
    model.print_trainable_parameters()

    # ── 4. Data ──
    print("[3/6] Loading FABLE 5 dataset...")
    dataset = FableDataset(args.data, tokenizer, max_length=args.max_length)
    loader = DataLoader(
        dataset,
        batch_size=args.batch_size,
        shuffle=True,
        collate_fn=partial(collate_fn, pad_token_id=tokenizer.pad_token_id),
        num_workers=2,
        pin_memory=True,
    )
    print(f"  {len(dataset)} samples, {len(loader)} batches/epoch")

    # ── 4. Optimizer & scheduler (only trainable LoRA params) ──
    print("[4/6] Setting up optimizer...")
    opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=args.lr)
    total_steps = len(loader) * args.epochs // args.grad_accum
    scheduler = get_linear_schedule_with_warmup(opt, num_warmup_steps=int(0.05 * total_steps), num_training_steps=total_steps)

    # ── 6. Training loop ──
    print(f"[5/6] Training ({args.epochs} epoch(s))...")
    os.makedirs(args.output, exist_ok=True)
    global_step = 0
    best_loss = float("inf")

    for epoch in range(args.epochs):
        model.train()
        total_loss = 0.0
        n_batches = 0
        epoch_t0 = time.time()

        for batch_idx, batch in enumerate(loader):
            batch = batch.to(model.device)
            labels = batch.clone()

            loss = model(input_ids=batch, labels=labels).loss
            loss = loss / args.grad_accum
            loss.backward()

            total_loss += loss.item() * args.grad_accum
            n_batches += 1

            if (batch_idx + 1) % args.grad_accum == 0 or (batch_idx + 1) == len(loader):
                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
                opt.step()
                scheduler.step()
                opt.zero_grad()
                global_step += 1

                if global_step % args.save_steps == 0:
                    avg_loss = total_loss / n_batches
                    ppl = math.exp(avg_loss)
                    save_path = os.path.join(args.output, f"checkpoint-{global_step}")
                    model.save_pretrained(save_path)
                    tokenizer.save_pretrained(save_path)
                    print(f"  Step {global_step}: loss={avg_loss:.4f}, ppl={ppl:.2f}, lr={scheduler.get_last_lr()[0]:.2e}")

            if (batch_idx + 1) % 20 == 0:
                current_loss = total_loss / n_batches
                print(f"  Epoch {epoch+1}, batch {batch_idx+1}/{len(loader)}: loss={current_loss:.4f}")

        avg_loss = total_loss / n_batches
        ppl = math.exp(avg_loss)
        epoch_time = time.time() - epoch_t0
        print(f"  Epoch {epoch+1} done: loss={avg_loss:.4f}, ppl={ppl:.2f}, time={epoch_time:.0f}s")

        if avg_loss < best_loss:
            best_loss = avg_loss
            model.save_pretrained(os.path.join(args.output, "best"))
            tokenizer.save_pretrained(os.path.join(args.output, "best"))

    # Save final
    model.save_pretrained(os.path.join(args.output, "final"))
    tokenizer.save_pretrained(os.path.join(args.output, "final"))
    print(f"\nDone in {time.time()-t0:.0f}s. Final model: {args.output}/final")
    print(f"Best loss: {best_loss:.4f} (PPL={math.exp(best_loss):.2f})")


if __name__ == "__main__":
    train()