scripts/train_qlora.py

#!/usr/bin/env python3
"""QLoRA SFT training for TMF921 intent-to-config research dataset.

Designed for a single RTX 6000 Ada 48/50GB server. Uses TRL SFTTrainer with PEFT QLoRA.
"""
import argparse
import math
import os
import re
from pathlib import Path

import torch
from datasets import load_dataset
from peft import LoraConfig
from transformers import AutoTokenizer, BitsAndBytesConfig, TrainerCallback, set_seed
from trl import SFTConfig, SFTTrainer

from tmf921_train.utils import load_config, write_json

try:
    import trackio
except Exception:  # pragma: no cover
    trackio = None


class TrackioAlertCallback(TrainerCallback):
    def on_log(self, args, state, control, logs=None, **kwargs):
        if not state.is_world_process_zero or not logs or trackio is None:
            return
        loss = logs.get("loss")
        grad_norm = logs.get("grad_norm")
        if loss is not None and (math.isnan(float(loss)) or math.isinf(float(loss))):
            trackio.alert(
                title="NaN/Inf training loss",
                text=f"step={state.global_step} loss={loss} — stop run and reduce learning_rate by 10x.",
                level="ERROR",
            )
        if grad_norm is not None and float(grad_norm) > 10.0:
            trackio.alert(
                title="Gradient norm spike",
                text=f"step={state.global_step} grad_norm={float(grad_norm):.3f} — consider lower lr or max_grad_norm.",
                level="WARN",
            )

    def on_evaluate(self, args, state, control, metrics=None, **kwargs):
        if not state.is_world_process_zero or not metrics or trackio is None:
            return
        loss = metrics.get("eval_loss")
        if loss is not None and float(loss) > 1.0:
            trackio.alert(
                title="High validation loss",
                text=f"step={state.global_step} eval_loss={float(loss):.4f} — check convergence and rare-class oversampling.",
                level="WARN",
            )


def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("--config", default="configs/rtx6000ada_qwen3_8b_qlora.yaml")
    p.add_argument("--model_name_or_path")
    p.add_argument("--dataset_name")
    p.add_argument("--train_split")
    p.add_argument("--eval_split")
    p.add_argument("--output_dir")
    p.add_argument("--hub_model_id")
    p.add_argument("--max_steps", type=int, default=None, help="Debug/short run override")
    p.add_argument("--no_push", action="store_true")
    p.add_argument("--packing", action="store_true", help="Override config and enable packing. Requires compatible attention setup.")
    p.add_argument("--flash_attn", action="store_true", help="Use flash_attention_2 in model_init_kwargs. Install flash-attn first.")
    p.add_argument("--resume_from_checkpoint", default=None, help="Path to checkpoint dir, or 'true' to auto-resume latest checkpoint in output_dir")
    p.add_argument("--seed", type=int, default=42)
    return p.parse_args()


def require_cuda():
    print("=== CUDA CHECK ===")
    print(f"torch={torch.__version__} torch.version.cuda={torch.version.cuda} CUDA_VISIBLE_DEVICES={os.environ.get('CUDA_VISIBLE_DEVICES')}")
    if not torch.cuda.is_available():
        raise RuntimeError(
            "CUDA is not available to PyTorch. Refusing to train on CPU. "
            "Run `bash scripts/install_rtx6000ada.sh`, verify `nvidia-smi`, and set CUDA_VISIBLE_DEVICES=0."
        )
    print(f"cuda device_count={torch.cuda.device_count()} gpu0={torch.cuda.get_device_name(0)}")


def valid_hf_repo_id(repo_id):
    if not repo_id or not isinstance(repo_id, str):
        return False
    if repo_id.endswith("/") or repo_id.startswith("/") or "//" in repo_id:
        return False
    pattern = r"^[A-Za-z0-9][A-Za-z0-9._-]{0,95}/[A-Za-z0-9][A-Za-z0-9._-]{0,95}$"
    return re.match(pattern, repo_id) is not None


def sanitize_trackio_config(cfg):
    # Environment variable takes precedence only if valid. Invalid values like "nraptisss/"
    # crash Trackio before training starts, so ignore them and continue without a Space.
    env_space = os.environ.get("TRACKIO_SPACE_ID", "").strip()
    cfg_space = str(cfg.get("trackio_space_id") or "").strip()
    chosen = env_space or cfg_space
    if chosen and valid_hf_repo_id(chosen):
        cfg["trackio_space_id"] = chosen
        print(f"Trackio Space: {chosen}")
    else:
        if chosen:
            print(f"WARNING: ignoring invalid Trackio Space ID: {chosen!r}. Expected format: namespace/space-name")
        cfg["trackio_space_id"] = None
        os.environ.pop("TRACKIO_SPACE_ID", None)
    # Set DISABLE_TRACKIO=1 to bypass Trackio completely if desired.
    if os.environ.get("DISABLE_TRACKIO", "0") == "1":
        print("Trackio disabled via DISABLE_TRACKIO=1")
        cfg["project"] = None
        cfg["trackio_space_id"] = None
    return cfg


def main():
    args = parse_args()
    require_cuda()
    cfg = load_config(args.config)
    cfg = sanitize_trackio_config(cfg)
    for k in ["model_name_or_path", "dataset_name", "train_split", "eval_split", "output_dir", "hub_model_id"]:
        v = getattr(args, k)
        if v is not None:
            cfg[k] = v
    if args.max_steps is not None:
        cfg["max_steps"] = args.max_steps
        cfg["num_train_epochs"] = 1
    if args.no_push:
        cfg["push_to_hub"] = False
    if args.packing:
        cfg["packing"] = True

    set_seed(args.seed)
    Path(cfg["output_dir"]).mkdir(parents=True, exist_ok=True)
    write_json(Path(cfg["output_dir"]) / "resolved_config.json", cfg)

    print("Loading dataset", cfg["dataset_name"])
    ds = load_dataset(cfg["dataset_name"])
    train_dataset = ds[cfg.get("train_split", "train_sota")]
    eval_dataset = ds[cfg.get("eval_split", "validation")]
    print(train_dataset)
    print(eval_dataset)

    # TRL infers dataset type from column names. This research dataset includes both
    # `messages` and convenience `prompt`/`completion` columns; passing all columns can
    # make TRL classify it as prompt-completion instead of conversational and reject
    # assistant_only_loss=True. For SFT we intentionally train from ChatML `messages`.
    train_dataset = train_dataset.select_columns(["messages"])
    eval_dataset = eval_dataset.select_columns(["messages"])
    print("SFT train columns:", train_dataset.column_names)
    print("SFT eval columns:", eval_dataset.column_names)

    tokenizer = AutoTokenizer.from_pretrained(cfg["model_name_or_path"], trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    bnb_config = None
    if cfg.get("load_in_4bit", True):
        bnb_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type=cfg.get("bnb_4bit_quant_type", "nf4"),
            bnb_4bit_use_double_quant=bool(cfg.get("bnb_4bit_use_double_quant", True)),
            bnb_4bit_compute_dtype=torch.bfloat16,
        )

    model_init_kwargs = {
        "trust_remote_code": True,
        "device_map": {"": 0},
        "dtype": torch.bfloat16 if cfg.get("bf16", True) else torch.float16,
    }
    if bnb_config is not None:
        model_init_kwargs["quantization_config"] = bnb_config
    if args.flash_attn:
        model_init_kwargs["attn_implementation"] = "flash_attention_2"

    target_modules = cfg.get("lora_target_modules", "all-linear")
    peft_config = LoraConfig(
        r=int(cfg.get("lora_r", 64)),
        lora_alpha=int(cfg.get("lora_alpha", 16)),
        lora_dropout=float(cfg.get("lora_dropout", 0.05)),
        bias="none",
        task_type="CAUSAL_LM",
        target_modules=target_modules,
    )

    report_to = "trackio" if cfg.get("project") else "none"
    sft_args = SFTConfig(
        output_dir=cfg["output_dir"],
        model_init_kwargs=model_init_kwargs,
        max_length=int(cfg.get("max_length", 2048)),
        packing=bool(cfg.get("packing", False)),
        assistant_only_loss=bool(cfg.get("assistant_only_loss", True)),
        dataset_num_proc=int(cfg.get("dataset_num_proc", 8)),
        learning_rate=float(cfg.get("learning_rate", 2e-4)),
        lr_scheduler_type=cfg.get("lr_scheduler_type", "constant"),
        warmup_steps=int(cfg.get("warmup_steps", 0)),
        weight_decay=float(cfg.get("weight_decay", 0.0)),
        max_grad_norm=float(cfg.get("max_grad_norm", 0.3)),
        num_train_epochs=float(cfg.get("epochs", 2)),
        max_steps=int(cfg["max_steps"]) if cfg.get("max_steps") is not None else -1,
        per_device_train_batch_size=int(cfg.get("per_device_train_batch_size", 2)),
        gradient_accumulation_steps=int(cfg.get("gradient_accumulation_steps", 8)),
        per_device_eval_batch_size=int(cfg.get("per_device_eval_batch_size", 2)),
        bf16=bool(cfg.get("bf16", True)),
        gradient_checkpointing=bool(cfg.get("gradient_checkpointing", True)),
        gradient_checkpointing_kwargs={"use_reentrant": False},
        optim=cfg.get("optim", "paged_adamw_32bit"),
        eval_strategy="steps",
        eval_steps=int(cfg.get("eval_steps", 250)),
        save_strategy="steps",
        save_steps=int(cfg.get("save_steps", 250)),
        save_total_limit=int(cfg.get("save_total_limit", 3)),
        load_best_model_at_end=True,
        metric_for_best_model="eval_loss",
        greater_is_better=False,
        logging_strategy="steps",
        logging_steps=int(cfg.get("logging_steps", 10)),
        logging_first_step=True,
        disable_tqdm=True,
        report_to=report_to,
        run_name=cfg.get("run_name"),
        project=cfg.get("project"),
        trackio_space_id=cfg.get("trackio_space_id"),
        push_to_hub=bool(cfg.get("push_to_hub", True)),
        hub_model_id=cfg.get("hub_model_id"),
    )

    trainer = SFTTrainer(
        model=cfg["model_name_or_path"],
        args=sft_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        processing_class=tokenizer,
        peft_config=peft_config,
        callbacks=[TrackioAlertCallback()],
    )

    resume_arg = args.resume_from_checkpoint
    if resume_arg is not None and str(resume_arg).lower() == "true":
        resume_arg = True
    trainer.train(resume_from_checkpoint=resume_arg)
    metrics = trainer.evaluate()
    write_json(Path(cfg["output_dir"]) / "final_eval_metrics.json", metrics)
    trainer.save_model(cfg["output_dir"])
    tokenizer.save_pretrained(cfg["output_dir"])

    if bool(cfg.get("push_to_hub", True)):
        trainer.push_to_hub(
            commit_message="Qwen TMF921 QLoRA SFT",
            dataset_name=cfg["dataset_name"],
        )
        print(f"Pushed model/adapters to https://huggingface.co/{cfg.get('hub_model_id')}")


if __name__ == "__main__":
    main()