worlddisasterlm/training/train_qlora.py

"""
Production QLoRA fine-tuning pipeline for WorldDisasterLM.

Hardware requirements
---------------------
Minimum  : 1× NVIDIA A100 40 GB (or 2× RTX 3090/4090 24 GB with gradient checkpointing)
Preferred: 2–8× A100/H100 80 GB for full batch sizes
Consumer : 1× RTX 4090 24 GB  →  set per_device_train_batch_size=1, gradient_accumulation_steps=16

Prerequisites
-------------
1. Accept Meta Llama 3.1 license on Hugging Face:
   https://huggingface.co/meta-llama/Llama-3.1-8B-Instruct
2. Set HF_TOKEN in your .env file:
   huggingface-cli login  (or export HF_TOKEN=hf_...)

Run
---
python scripts/train_production.py --dataset data/processed/instruction_dataset.jsonl
"""

from __future__ import annotations

import json
import logging
import os
from dataclasses import dataclass, field
from pathlib import Path

from worlddisasterlm.training.chat_format import apply_template

logger = logging.getLogger(__name__)


@dataclass
class QLoRAConfig:
    # ── Model ─────────────────────────────────────────────────────────────────
    base_model: str = "meta-llama/Llama-3.1-8B-Instruct"
    output_dir: str = "checkpoints/worlddisasterlm-qlora"

    # ── Data ──────────────────────────────────────────────────────────────────
    dataset_path: str = "data/processed/instruction_dataset.jsonl"
    max_seq_length: int = 2048

    # ── QLoRA / PEFT ──────────────────────────────────────────────────────────
    use_4bit: bool = True
    bnb_4bit_quant_type: str = "nf4"          # "nf4" or "fp4"
    use_double_quant: bool = True
    compute_dtype: str = "bfloat16"           # "bfloat16" or "float16"

    lora_r: int = 16
    lora_alpha: int = 32
    lora_dropout: float = 0.05
    lora_target_modules: list[str] = field(
        default_factory=lambda: [
            "q_proj", "k_proj", "v_proj", "o_proj",
            "gate_proj", "up_proj", "down_proj",
        ]
    )

    # ── Optimisation ──────────────────────────────────────────────────────────
    epochs: int = 3
    learning_rate: float = 2e-4
    per_device_train_batch_size: int = 2
    gradient_accumulation_steps: int = 8
    warmup_ratio: float = 0.03
    max_grad_norm: float = 0.3
    lr_scheduler_type: str = "cosine"
    weight_decay: float = 0.0

    bf16: bool = True
    fp16: bool = False
    gradient_checkpointing: bool = True

    # ── Logging & saving ──────────────────────────────────────────────────────
    logging_steps: int = 10
    save_steps: int = 200
    save_total_limit: int = 3
    report_to: str = "none"      # "mlflow" | "wandb" | "none"

    # ── Misc ──────────────────────────────────────────────────────────────────
    seed: int = 42


def load_jsonl(path: str) -> list[dict]:
    records = []
    with open(path, encoding="utf-8") as handle:
        for line in handle:
            line = line.strip()
            if line:
                records.append(json.loads(line))
    return records


def build_dataset(path: str, tokenizer):
    from datasets import Dataset

    raw = load_jsonl(path)
    logger.info("Loaded %d instruction samples from %s", len(raw), path)

    texts = []
    for sample in raw:
        try:
            text = apply_template(
                tokenizer,
                instruction=sample.get("instruction", ""),
                context=sample.get("input", ""),
                output=sample.get("output", ""),
            )
            texts.append(text)
        except Exception as exc:
            logger.warning("Failed to format sample: %s", exc)

    logger.info("Formatted %d samples with chat template", len(texts))
    return Dataset.from_dict({"text": texts})


def setup_model_tokenizer(config: QLoRAConfig):
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
    from peft import LoraConfig, TaskType, get_peft_model, prepare_model_for_kbit_training

    # ── Tokenizer ─────────────────────────────────────────────────────────────
    tokenizer = AutoTokenizer.from_pretrained(
        config.base_model,
        trust_remote_code=True,
        padding_side="right",
    )
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
        tokenizer.pad_token_id = tokenizer.eos_token_id

    # ── Quantization ──────────────────────────────────────────────────────────
    torch_dtype = torch.bfloat16 if config.bf16 else torch.float16
    bnb_config = None
    if config.use_4bit:
        bnb_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type=config.bnb_4bit_quant_type,
            bnb_4bit_compute_dtype=torch_dtype,
            bnb_4bit_use_double_quant=config.use_double_quant,
        )

    # ── Base model ────────────────────────────────────────────────────────────
    model = AutoModelForCausalLM.from_pretrained(
        config.base_model,
        quantization_config=bnb_config,
        device_map="auto",
        torch_dtype=torch_dtype,
        trust_remote_code=True,
        attn_implementation="flash_attention_2" if _flash_attn_available() else "eager",
    )

    if config.use_4bit:
        model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=config.gradient_checkpointing)
    elif config.gradient_checkpointing:
        model.gradient_checkpointing_enable()

    model.config.use_cache = False

    # ── LoRA ──────────────────────────────────────────────────────────────────
    peft_config = LoraConfig(
        task_type=TaskType.CAUSAL_LM,
        r=config.lora_r,
        lora_alpha=config.lora_alpha,
        target_modules=config.lora_target_modules,
        lora_dropout=config.lora_dropout,
        bias="none",
        inference_mode=False,
    )
    model = get_peft_model(model, peft_config)
    model.print_trainable_parameters()

    return model, tokenizer


def _flash_attn_available() -> bool:
    try:
        import flash_attn  # noqa: F401
        return True
    except ImportError:
        return False


def train(config: QLoRAConfig) -> None:
    from trl import SFTConfig, SFTTrainer

    # Seed
    from worlddisasterlm.utils.seed import seed_everything
    seed_everything(config.seed)

    Path(config.output_dir).mkdir(parents=True, exist_ok=True)

    # MLflow / W&B config
    if config.report_to == "mlflow" and os.getenv("MLFLOW_TRACKING_URI"):
        import mlflow
        mlflow.set_tracking_uri(os.environ["MLFLOW_TRACKING_URI"])
        mlflow.set_experiment("worlddisasterlm-training")

    model, tokenizer = setup_model_tokenizer(config)
    train_dataset = build_dataset(config.dataset_path, tokenizer)

    report_to_list = [config.report_to] if config.report_to != "none" else []

    sft_config = SFTConfig(
        output_dir=config.output_dir,
        num_train_epochs=config.epochs,
        per_device_train_batch_size=config.per_device_train_batch_size,
        gradient_accumulation_steps=config.gradient_accumulation_steps,
        gradient_checkpointing=config.gradient_checkpointing,
        learning_rate=config.learning_rate,
        max_grad_norm=config.max_grad_norm,
        warmup_ratio=config.warmup_ratio,
        lr_scheduler_type=config.lr_scheduler_type,
        weight_decay=config.weight_decay,
        bf16=config.bf16,
        fp16=config.fp16,
        logging_steps=config.logging_steps,
        save_steps=config.save_steps,
        save_total_limit=config.save_total_limit,
        load_best_model_at_end=False,
        max_seq_length=config.max_seq_length,
        dataset_text_field="text",
        packing=False,
        report_to=report_to_list,
        seed=config.seed,
    )

    trainer = SFTTrainer(
        model=model,
        tokenizer=tokenizer,
        train_dataset=train_dataset,
        args=sft_config,
    )

    logger.info("Starting QLoRA fine-tuning: %d samples, %d epochs", len(train_dataset), config.epochs)
    trainer.train()

    logger.info("Saving adapter to %s", config.output_dir)
    trainer.save_model(config.output_dir)
    tokenizer.save_pretrained(config.output_dir)
    logger.info("Training complete.")