train.py

#!/usr/bin/env python3
"""
STEP 4 — LoRA/QLoRA fine-tune with Unsloth
"""

from __future__ import annotations

import json
from pathlib import Path

from utils import ensure_dir, load_config, read_jsonl


def main() -> None:
    cfg = load_config()
    output_dir = ensure_dir(cfg["output_dir"])
    adapter_dir = ensure_dir(cfg["adapter_output_dir"])

    train_path = output_dir / "train.jsonl"
    val_path = output_dir / "val.jsonl"
    if not train_path.exists():
        raise FileNotFoundError(f"Missing {train_path}. Run build_dataset.py first.")

    try:
        from unsloth import FastLanguageModel
        from datasets import Dataset
        from trl import SFTTrainer, SFTConfig
    except ImportError as e:
        raise ImportError(
            "Training dependencies missing. Install with:\n"
            "  pip install torch --index-url https://download.pytorch.org/whl/cu130\n"
            "  pip install unsloth transformers datasets peft trl bitsandbytes"
        ) from e

    train_records = read_jsonl(train_path)
    val_records = read_jsonl(val_path) if val_path.exists() else []

    max_samples = cfg.get("max_train_samples", None)
    if max_samples:
        train_records = train_records[:int(max_samples)]
        val_records = val_records[:max(1, int(max_samples) // 10)]
        print(f"[smoke test] Capped to {len(train_records)} train / {len(val_records)} val samples.")

    print(f"Loading base model: {cfg['base_model']}")
    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name=cfg["base_model"],
        max_seq_length=cfg.get("max_seq_length", 2048),
        dtype=None,
        load_in_4bit=True,
    )

    model = FastLanguageModel.get_peft_model(
        model,
        r=cfg.get("lora_rank", 16),
        target_modules=[
            "q_proj", "k_proj", "v_proj", "o_proj",
            "gate_proj", "up_proj", "down_proj",
        ],
        lora_alpha=cfg.get("lora_alpha", 32),
        lora_dropout=cfg.get("lora_dropout", 0.05),
        bias="none",
        use_gradient_checkpointing="unsloth",
        random_state=42,
    )

    train_ds = Dataset.from_list(train_records)
    val_ds = Dataset.from_list(val_records) if val_records else None

    has_val = val_ds is not None

    # Use SFTConfig (trl's own config class) to avoid pickle mismatch
    # when saving checkpoints with newer trl+transformers combos.
    sft_cfg = SFTConfig(
        output_dir=str(adapter_dir),
        # --- dataset ---
        dataset_text_field="text",
        max_seq_length=cfg.get("max_seq_length", 2048),
        # --- training ---
        per_device_train_batch_size=cfg.get("batch_size", 2),
        per_device_eval_batch_size=cfg.get("batch_size", 2),
        gradient_accumulation_steps=cfg.get("gradient_accumulation_steps", 4),
        num_train_epochs=cfg.get("epochs", 3),
        learning_rate=cfg.get("learning_rate", 2e-4),
        warmup_ratio=cfg.get("warmup_ratio", 0.03),
        logging_steps=1,
        save_steps=cfg.get("save_steps", 500),
        eval_steps=cfg.get("eval_steps", 500) if has_val else None,
        eval_strategy="steps" if has_val else "no",
        save_total_limit=2,
        fp16=False,
        bf16=True,
        optim="adamw_8bit",
        weight_decay=0.01,
        lr_scheduler_type="cosine",
        seed=42,
        report_to="none",
        # Packing can sometimes trigger CPU gradient offloading on 12GB cards, killing speed.
        # Disabling it relies on natural short padding, which is way faster here.
        packing=False,
    )

    import sys
    import time
    import torch
    from transformers import TrainerCallback

    class CliGuiCallback(TrainerCallback):
        def __init__(self, total_steps):
            self.total_steps = total_steps
            self.start_time = time.time()
            self.has_printed = False

        def on_step_end(self, args, state, control, **kwargs):
            if state.global_step == 0:
                return

            # Calculate metrics
            step = state.global_step
            total = self.total_steps
            pct = (step / total) * 100
            
            # Progress bar
            bar_len = 30
            filled_len = int(bar_len * step // total)
            bar = '=' * filled_len + '>' + ' ' * (bar_len - filled_len - 1)
            if filled_len == bar_len:
                bar = '=' * bar_len

            # Timing & ETA
            elapsed = time.time() - self.start_time
            steps_per_sec = step / elapsed if elapsed > 0 else 0
            eta_sec = (total - step) / steps_per_sec if steps_per_sec > 0 else 0
            
            def format_time(seconds):
                if seconds < 60: return f"{seconds:.0f}s"
                if seconds < 3600: return f"{seconds//60:.0f}m {seconds%60:.0f}s"
                return f"{seconds//3600:.0f}h {(seconds%3600)//60:.0f}m"

            eta_str = format_time(eta_sec)
            elapsed_str = format_time(elapsed)

            # Metrics
            loss = "N/A"
            for log in reversed(state.log_history):
                if "loss" in log:
                    loss = f"{log['loss']:.4f}"
                    break

            lr = "N/A"
            for log in reversed(state.log_history):
                if "learning_rate" in log:
                    lr = f"{log['learning_rate']:.2e}"
                    break

            # VRAM
            vram_alloc = torch.cuda.memory_allocated() / 1024**3
            vram_max = torch.cuda.max_memory_allocated() / 1024**3
            vram_str = f"{vram_alloc:.2f}GB / {vram_max:.2f}GB"

            # ANSI coloring & positioning
            # Move cursor up by 8 lines if we already printed once, so we overwrite the panel
            if self.has_printed:
                sys.stdout.write("\033[8A")
            else:
                self.has_printed = True

            panel = (
                f"\033[1;36m┌────────────────────────────────────────────────────────┐\033[0m\n"
                f"\033[1;36m│\033[0m \033[1;33mLLM-Imitate Fine-Tuning Dashboard\033[0m                      \033[1;36m│\033[0m\n"
                f"\033[1;36m├────────────────────────────────────────────────────────┤\033[0m\n"
                f"\033[1;36m│\033[0m Progress: [{bar}] {pct:.1f}% ({step}/{total}) \033[1;36m│\033[0m\n"
                f"\033[1;36m│\033[0m Epoch: {state.epoch:.2f} | Loss: \033[1;31m{loss}\033[0m | LR: \033[1;32m{lr}\033[0m               \033[1;36m│\033[0m\n"
                f"\033[1;36m│\033[0m Elapsed: {elapsed_str} | ETA: {eta_str} | Speed: {steps_per_sec:.2f} step/s \033[1;36m│\033[0m\n"
                f"\033[1;36m│\033[0m GPU VRAM: {vram_str}                               \033[1;36m│\033[0m\n"
                f"\033[1;36m└────────────────────────────────────────────────────────┘\033[0m\n"
            )
            sys.stdout.write(panel)
            sys.stdout.flush()

    # Calculate steps
    # Packed sequence dataset length might change, but we can approximate it or use state.max_steps
    # Actually state.max_steps is set once training starts. So we can grab it from trainer.
    # To pass it, we can initialize callback with a placeholder and set total_steps inside training if needed,
    # or let the callback fetch total steps dynamically.
    
    # We will initialize callback with 0 and retrieve max_steps on the fly
    class CliGuiDynamicCallback(CliGuiCallback):
        def __init__(self):
            super().__init__(total_steps=1)
        def on_step_end(self, args, state, control, **kwargs):
            self.total_steps = state.max_steps
            super().on_step_end(args, state, control, **kwargs)

    trainer = SFTTrainer(
        model=model,
        tokenizer=tokenizer,
        train_dataset=train_ds,
        eval_dataset=val_ds,
        args=sft_cfg,
        callbacks=[CliGuiDynamicCallback()],
    )

    print(f"Training on {len(train_records):,} examples ...")
    # Hide the default transformers tqdm bar so it doesn't clash with our CLI GUI
    trainer.train()

    final_dir = adapter_dir / "final"
    final_dir.mkdir(parents=True, exist_ok=True)
    model.save_pretrained(str(final_dir))
    tokenizer.save_pretrained(str(final_dir))

    meta = {
        "base_model": cfg["base_model"],
        "train_examples": len(train_records),
        "val_examples": len(val_records),
        "adapter_path": str(final_dir),
    }
    with open(adapter_dir / "training_meta.json", "w", encoding="utf-8") as f:
        json.dump(meta, f, indent=2)

    print(f"\nTraining complete. Adapter saved -> {final_dir}")
    print("Next: bash convert_to_gguf.sh")


if __name__ == "__main__":
    main()