training/train_sft.py

"""Unsloth + Hugging Face TRL SFT for CORP-ENV action-format warm start.

This is the hackathon SFT training script. It uses:

- `unsloth.FastLanguageModel` for efficient 4-bit LoRA/QLoRA loading.
- `trl.SFTTrainer` / `trl.SFTConfig` for supervised fine-tuning.
- `messages`-format JSONL with TRL 0.2x conversational SFT (optional `--assistant-only` if the
  model chat template supports assistant token masks; Qwen2.5 Instruct defaults to off).

Run on Colab, Lightning AI H100, or another GPU machine:

  python training/train_sft.py \
    --model Qwen/Qwen2.5-7B-Instruct \
    --data data/sft/e1_m1_h1_examples.jsonl \
    --output outputs/sft_adapter \
    --max-steps 30 \
    --push-to-hub your-org/corp-env-sft-adapter
"""

from __future__ import annotations

import argparse
import inspect
import json
from dataclasses import fields
from pathlib import Path
from typing import Any, Dict, List

import torch  # imported before unsloth in main for dtype hooks


def _sft_config_field_names() -> set[str]:
    from trl import SFTConfig

    if hasattr(SFTConfig, "__dataclass_fields__"):
        return set(SFTConfig.__dataclass_fields__.keys())
    return {f.name for f in fields(SFTConfig)}


def load_conversation_rows(path: Path) -> List[Dict[str, Any]]:
    """TRL conversational format: each row has a `messages` list."""
    rows: List[Dict[str, Any]] = []
    with path.open("r", encoding="utf-8") as f:
        for line in f:
            if not line.strip():
                continue
            obj = json.loads(line)
            if "messages" not in obj:
                raise SystemExit(
                    f"{path}: SFT example missing 'messages' (use --legacy-text for old format)."
                )
            rows.append({"messages": obj["messages"]})
    return rows


def load_text_rows(path: Path, tokenizer: object) -> List[Dict[str, str]]:
    rows: List[Dict[str, str]] = []
    with path.open("r", encoding="utf-8") as f:
        for line in f:
            if not line.strip():
                continue
            obj = json.loads(line)
            messages = obj["messages"]
            text = tokenizer.apply_chat_template(
                messages,
                tokenize=False,
                add_generation_prompt=False,
            )
            rows.append({"text": text})
    return rows


def _build_sft_config(
    allowed: set[str],
    output_dir: str,
    max_seq: int,
    args: argparse.Namespace,
) -> Any:
    from trl import SFTConfig

    candidate: Dict[str, Any] = {
        "output_dir": output_dir,
        "max_length": max_seq,
        "per_device_train_batch_size": args.batch_size,
        "gradient_accumulation_steps": args.grad_accum,
        "num_train_epochs": args.epochs,
        "learning_rate": args.lr,
        "warmup_ratio": 0.05,
        "lr_scheduler_type": "cosine",
        "logging_steps": 5,
        "save_steps": args.save_steps,
        "save_total_limit": 3,
        "max_steps": args.max_steps,
        "optim": args.optim,
        "bf16": (not args.fp16) and torch.cuda.is_available(),
        "fp16": bool(args.fp16) and torch.cuda.is_available(),
        "packing": args.packing,
        "report_to": "none",
        "push_to_hub": bool(args.push_to_hub),
        "hub_model_id": args.push_to_hub or None,
        "assistant_only_loss": args.assistant_only,
    }
    if args.dataset_num_proc is not None and "dataset_num_proc" in allowed:
        candidate["dataset_num_proc"] = args.dataset_num_proc
    if args.dataloader_num_workers and "dataloader_num_workers" in allowed:
        candidate["dataloader_num_workers"] = args.dataloader_num_workers
    if args.packing and "padding_free" in allowed and args.padding_free:
        candidate["padding_free"] = True
    if args.legacy_text and "dataset_text_field" in allowed:
        candidate["dataset_text_field"] = "text"

    kwargs = {k: v for k, v in candidate.items() if k in allowed}
    if "output_dir" not in kwargs:
        kwargs["output_dir"] = output_dir
    return SFTConfig(**kwargs)


def _build_trainer(
    model: object,
    tokenizer: object,
    config: object,
    dataset: object,
) -> Any:
    from trl import SFTTrainer

    sig = inspect.signature(SFTTrainer.__init__)
    if "processing_class" in sig.parameters:
        return SFTTrainer(
            model=model,
            args=config,
            train_dataset=dataset,
            processing_class=tokenizer,
        )
    if "tokenizer" in sig.parameters:
        return SFTTrainer(
            model=model,
            args=config,
            train_dataset=dataset,
            tokenizer=tokenizer,
        )
    raise SystemExit("SFTTrainer: expected processing_class or tokenizer in __init__.")


def main() -> None:
    parser = argparse.ArgumentParser(description="Train CORP-ENV SFT LoRA adapter.")
    parser.add_argument("--model", default="Qwen/Qwen2.5-7B-Instruct")
    parser.add_argument("--data", default="data/sft/e1_m1_h1_examples.jsonl")
    parser.add_argument("--output", default="outputs/sft_adapter")
    parser.add_argument("--max-seq-length", type=int, default=8192)
    parser.add_argument("--epochs", type=float, default=2.0)
    parser.add_argument("--batch-size", type=int, default=1)
    parser.add_argument("--grad-accum", type=int, default=8)
    parser.add_argument("--lr", type=float, default=2e-4)
    parser.add_argument("--max-steps", type=int, default=-1, help="Positive value for quick judge/Colab smoke runs.")
    parser.add_argument("--save-steps", type=int, default=50)
    parser.add_argument("--optim", default="adamw_8bit")
    parser.add_argument("--push-to-hub", default="")
    parser.add_argument(
        "--fp16",
        action="store_true",
        help="Use fp16 instead of bf16 (bf16 is default on CUDA when available).",
    )
    parser.add_argument(
        "--legacy-text",
        action="store_true",
        help="Pre-tokenize to a single 'text' column (older path). Default: TRL messages format.",
    )
    parser.add_argument(
        "--assistant-only",
        action="store_true",
        help=(
            "Loss only on assistant tokens (requires tokenizer chat template with assistant masks). "
            "Qwen2.5 usually needs this off; use with templates that support {% generation %}."
        ),
    )
    parser.add_argument(
        "--packing",
        action="store_true",
        help="Sequence packing (faster; avoid for very long or uneven traces unless you know the tradeoffs).",
    )
    parser.add_argument(
        "--padding-free",
        action="store_true",
        help="With packing+bfd, TRL can use padding-free; requires FlashAttention-friendly setup.",
    )
    parser.add_argument(
        "--dataset-num-proc",
        type=int,
        default=4,
        help="Multiprocess dataset preprocessing; set 0 to disable where supported (default 4 for H100).",
    )
    parser.add_argument(
        "--dataloader-num-workers",
        type=int,
        default=2,
        help="DataLoader workers (default 2; increase on H100 if CPU allows).",
    )
    args = parser.parse_args()

    try:
        from unsloth import FastLanguageModel
        from datasets import Dataset
    except ImportError as exc:
        raise SystemExit(
            "SFT training requires datasets, trl, and unsloth. On Lightning AI, install with:\n"
            "  pip install -e \".[training]\""
        ) from exc

    allowed = _sft_config_field_names()
    if args.dataset_num_proc == 0 and "dataset_num_proc" in allowed:
        args = argparse.Namespace(**{**vars(args), "dataset_num_proc": None})

    load_dtype = torch.float16 if args.fp16 else None
    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name=args.model,
        max_seq_length=args.max_seq_length,
        dtype=load_dtype,
        load_in_4bit=True,
    )
    if getattr(tokenizer, "pad_token", None) is None and getattr(
        tokenizer, "eos_token", None
    ) is not None:
        tokenizer.pad_token = tokenizer.eos_token

    model = FastLanguageModel.get_peft_model(
        model,
        r=32,
        target_modules=[
            "q_proj",
            "k_proj",
            "v_proj",
            "o_proj",
            "gate_proj",
            "up_proj",
            "down_proj",
        ],
        lora_alpha=32,
        lora_dropout=0.0,
        bias="none",
        use_gradient_checkpointing="unsloth",
        random_state=3407,
    )

    if args.legacy_text:
        rows = load_text_rows(Path(args.data), tokenizer)
    else:
        rows = load_conversation_rows(Path(args.data))
    if not rows:
        raise SystemExit(f"No training rows in {args.data!r}.")

    dataset = Dataset.from_list(rows)
    config = _build_sft_config(allowed, args.output, args.max_seq_length, args)
    trainer = _build_trainer(model, tokenizer, config, dataset)
    trainer.train()
    trainer.save_model(args.output)
    tokenizer.save_pretrained(args.output)

    if args.push_to_hub:
        trainer.push_to_hub()


if __name__ == "__main__":
    main()