train_sft.py

#!/usr/bin/env python3
"""
SFT Training Script — Qwen3.8-27B Instruction Tuning
──────────────────────────────────────────────────────────────────────────────
Trains Qwen3.8-27B with LoRA on a custom instruction dataset.
Uses ZeroGPU (A100 80GB) via HF Spaces GPU mount.

Usage:
  python3 train_sft.py

Expected output:
  - Adapter weights saved to ./adapters/qwen3.8-27b-sft-lora
  - Full fine-tuned model saved to ./models/qwen3.8-27b-sft-full
"""

import os
import json
import random
import hashlib
from dataclasses import dataclass, field
from typing import List, Optional
from pathlib import Path

import torch
import torch.nn as nn
from datasets import Dataset, DatasetDict
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    DataCollatorForLanguageModeling,
    TrainingArguments,
    Trainer,
)
from peft import LoraConfig, TaskType, get_peft_model

# ─── Configuration ──────────────────────────────────────────────────────────
@dataclass
class TrainingConfig:
    model_name: str = "unsloth/Qwen3.8-27B-GGUF"  # or "Qwen/Qwen3.8-27B" for full
    output_dir: str = "./models/qwen3.8-27b-sft-full"
    adapter_dir: str = "./adapters/qwen3.8-27b-sft-lora"
    batch_size: int = 1
    gradient_accumulation_steps: int = 4
    learning_rate: float = 1e-4
    lr_scheduler: str = "cosine"
    num_train_epochs: int = 3
    save_steps: int = 100
    eval_steps: int = 50
    log_steps: int = 10
    weight_decay: float = 0.05
    beta1: float = 0.9
    beta2: float = 0.999
    eps: float = 1e-8
    max_grad_norm: float = 1.0
    seed: int = 42
    lora_rank: int = 64
    lora_alpha: int = 128
    lora_dropout: float = 0.05
    target_modules: List[str] = field(default_factory=lambda: [
        "q_proj", "k_proj", "v_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj",
    ])
    use_8bit_adam: bool = True
    fp16: bool = True
    bf16: bool = True
    use_llama_flash_attn2: bool = True
    use_dora: bool = False
    use_rope_scaling: bool = False

# ─── Data Loading ───────────────────────────────────────────────────────────
def load_sft_dataset(data_path: str = "./data/sft_instructions.json") -> Dataset:
    """Load SFT training data from JSONL or JSON."""
    if not os.path.exists(data_path):
        # Generate synthetic data if none exists
        print(f"[INFO] Data not found at {data_path}, generating synthetic dataset...")
        return generate_synthetic_dataset()

    with open(data_path, "r") as f:
        data = json.load(f)

    # Ensure it's a list of dicts
    if isinstance(data, dict) and "messages" in data[0]:
        return Dataset.from_list(data)
    elif isinstance(data[0], dict):
        return Dataset.from_list(data)
    else:
        raise ValueError(f"Unexpected data format: {type(data[0])}")


def generate_synthetic_dataset(num_samples: int = 500) -> Dataset:
    """Generate synthetic instruction-following data for Qwen3.8-27B."""
    prompts = [
        "Explain quantum entanglement in simple terms.",
        "Write a Python function to reverse a string without using built-in reverse.",
        "Translate 'La vida es bella' to English.",
        "Summarize the following paragraph in one sentence: {text}",
        "Solve this math problem: {math_problem}",
        "Write a haiku about {topic}",
        "What is the capital of {country}?",
        "Explain how photosynthesis works.",
        "Debug this code: {code}",
        "Write a SQL query to find the top 5 customers by total purchase amount.",
    ]

    topics = ["mountains", "ocean", "space", "forest", "city", "rain", "sunrise", "robot", "ai"]
    countries = ["France", "Japan", "Brazil", "Australia", "Egypt", "Canada", "India", "Norway"]
    math_problems = [
        "What is 123456789 * 987654321?",
        "Calculate the derivative of x^3 + 2x^2 + x + 1 with respect to x.",
        "Integrate sin(x) from 0 to pi/2.",
    ]
    code_snippets = [
        "def fib(n):\n    if n <= 1:\n        return n\n    return fib(n-1) + fib(n-2)",
        "def bubble_sort(arr):\n    for i in range(len(arr)):\n        for j in range(len(arr)-1-i):\n            if arr[j] > arr[j+1]:\n                arr[j], arr[j+1] = arr[j+1], arr[j]",
    ]

    examples = []
    for _ in range(num_samples):
        prompt = random.choice(prompts)
        if "{text}" in prompt:
            text = "Artificial intelligence is transforming how we live and work. " \
                    "It powers everything from smartphone assistants to autonomous vehicles. " \
                    "Machine learning models can now generate text, create images, and even play games. " \
                    "However, AI also raises concerns about job displacement, bias, and misinformation."
        elif "{math_problem}" in prompt:
            math_problem = random.choice(math_problems)
            prompt = prompt.replace("{math_problem}", math_problem)
        elif "{code}" in prompt:
            code = random.choice(code_snippets)
            prompt = prompt.replace("{code}", code)
        elif "{topic}" in prompt:
            topic = random.choice(topics)
            prompt = prompt.replace("{topic}", topic)
        elif "{country}" in prompt:
            country = random.choice(countries)
            prompt = prompt.replace("{country}", country)
        elif "{text}" in prompt:
            text = "This is a placeholder text for summarization tasks."
            prompt = prompt.replace("{text}", text)
        elif "{math_problem}" in prompt:
            prompt = prompt.replace("{math_problem}", "Compute the factorial of 10.")
        elif "{topic}" in prompt:
            prompt = prompt.replace("{topic}", "mountains")

        examples.append({"instruction": prompt, "output": "This is a synthetic response."})

    return Dataset.from_list(examples)


# ─── Training Loop ──────────────────────────────────────────────────────────
@dataclass
class SFTResult:
    model: Optional[nn.Module] = None
    tokenizer: Optional[AutoTokenizer] = None
    best_loss: float = float("inf")
    best_model_path: Optional[str] = None
    history: List[dict] = field(default_factory=list)


def train_sft(config: Optional[TrainingConfig] = None) -> SFTResult:
    """Train Qwen3.8-27B with LoRA on instruction-following data."""
    cfg = config or TrainingConfig()

    # ─── Setup ───────────────────────────────────────────────────────────────
    print(f"[TRAIN] Model: {cfg.model_name}")
    print(f"[TRAIN] Output dir: {cfg.output_dir}")
    print(f"[TRAIN] Adapter dir: {cfg.adapter_dir}")
    print(f"[TRAIN] Epochs: {cfg.num_train_epochs}")
    print(f"[TRAIN] Batch size: {cfg.batch_size}")
    print(f"[TRAIN] LR: {cfg.learning_rate}")
    print(f"[TRAIN] LoRA rank: {cfg.lora_rank}")

    # Load tokenizer
    tokenizer = AutoTokenizer.from_pretrained(cfg.model_name, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    # Quantization config (optional — for 8-bit inference on limited VRAM)
    quant_config = BitsAndBytesConfig(
        load_in_8bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16,
        bnb_4bit_use_double_quant=True,
        llm_int8_enable_fp32_cpu_offload=False,
        llm_int8_threshold=6.0,
    )

    # Load model (full precision for training, 8-bit for inference)
    print("[TRAIN] Loading model...")
    model = AutoModelForCausalLM.from_pretrained(
        cfg.model_name,
        quantization_config=quant_config,
        trust_remote_code=True,
        torch_dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16,
        device_map="auto",
    )

    # ─── LoRA Config ─────────────────────────────────────────────────────────
    peft_config = LoraConfig(
        task_type=TaskType.CAUSAL_LM,
        inference_mode=False,
        r=cfg.lora_rank,
        lora_alpha=cfg.lora_alpha,
        lora_dropout=cfg.lora_dropout,
        target_modules=cfg.target_modules,
    )

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

    # ─── Dataset ─────────────────────────────────────────────────────────────
    dataset = load_sft_dataset()
    print(f"[TRAIN] Dataset size: {len(dataset)}")

    # Format: "instruction\noutput"
    def formatting_func(example):
        text = f"### Instruction:\n{example['instruction']}\n\n### Response:\n{example['output']}"
        return {"text": text}

    dataset = dataset.map(formatting_func)
    dataset = dataset.train_test_split(test_size=0.1, seed=cfg.seed)

    # ─── Collator ───────────────────────────────────────────────────────────
    collator = DataCollatorForLanguageModeling(
        tokenizer=tokenizer,
        mlm=False,
    )

    # ─── Training Arguments ──────────────────────────────────────────────────
    training_args = TrainingArguments(
        output_dir=cfg.output_dir,
        per_device_train_batch_size=cfg.batch_size,
        gradient_accumulation_steps=cfg.gradient_accumulation_steps,
        learning_rate=cfg.learning_rate,
        fp16=cfg.fp16,
        bf16=cfg.bf16,
        logging_steps=cfg.log_steps,
        save_steps=cfg.save_steps,
        save_total_limit=1,
        evaluation_strategy="steps",
        eval_steps=cfg.eval_steps,
        per_device_eval_batch_size=cfg.batch_size,
        num_train_epochs=cfg.num_train_epochs,
        weight_decay=cfg.weight_decay,
        lr_scheduler_type=cfg.lr_scheduler,
        load_best_model_at_end=True,
        metric_for_best_model="loss",
        greater_is_better=False,
        optim="paged_adamw_8bit" if cfg.use_8bit_adam else "adamw_torch",
        logging_strategy="steps",
        logging_first_step=True,
        remove_unused_columns=False,
        report_to="tensorboard",
        run_name=f"sft_qwen3.8-27b_lr{cfg.learning_rate} epochs{cfg.num_train_epochs}",
    )

    # ─── Trainer ─────────────────────────────────────────────────────────────
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=dataset["train"],
        eval_dataset=dataset["test"],
        tokenizer=tokenizer,
        data_collator=collator,
    )

    # ─── Train ───────────────────────────────────────────────────────────────
    print("[TRAIN] Starting training...")
    results = trainer.train()

    print("[TRAIN] Training complete!")
    print(f"[TRAIN] Best loss: {trainer.state.best_loss}")
    print(f"[TRAIN] Final loss: {trainer.state.log_history[-1]['loss']}")

    # Save model
    model.save_pretrained(cfg.output_dir)
    tokenizer.save_pretrained(cfg.output_dir)
    peft_config.save_pretrained(cfg.output_dir)

    # Save adapter config
    peft_config.save_pretrained(cfg.adapter_dir)

    return SFTResult(
        model=model,
        tokenizer=tokenizer,
        best_loss=results.best,
        best_model_path=cfg.output_dir,
        history=[
            {"epoch": h["epoch"], "loss": h["loss"], "learning_rate": h["learning_rate"]}
            for h in trainer.state.log_history
        ],
    )


# ─── Main ───────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Train Qwen3.8-27B with SFT")
    parser.add_argument("--model_name", type=str, default="unsloth/Qwen3.8-27B-GGUF")
    parser.add_argument("--output_dir", type=str, default="./models/qwen3.8-27b-sft-full")
    parser.add_argument("--adapter_dir", type=str, default="./adapters/qwen3.8-27b-sft-lora")
    parser.add_argument("--batch_size", type=int, default=1)
    parser.add_argument("--epochs", type=int, default=3)
    parser.add_argument("--lr", type=float, default=1e-4)
    parser.add_argument("--lora_rank", type=int, default=64)
    parser.add_argument("--lora_alpha", type=int, default=128)
    parser.add_argument("--data_path", type=str, default="./data/sft_instructions.json")
    parser.add_argument("--num_samples", type=int, default=500)

    args = parser.parse_args()

    config = TrainingConfig(
        model_name=args.model_name,
        output_dir=args.output_dir,
        adapter_dir=args.adapter_dir,
        batch_size=args.batch_size,
        num_train_epochs=args.epochs,
        learning_rate=args.lr,
        lora_rank=args.lora_rank,
        lora_alpha=args.lora_alpha,
    )

    result = train_sft(config)

    print(f"\n[COMPLETE] Model saved to: {result.best_model_path}")
    print(f"[COMPLETE] Adapter config saved to: {config.adapter_dir}")