train/model_load.py

"""FATHOM model loader — TRN-01.

Wraps Unsloth FastLanguageModel.from_pretrained + get_peft_model with a
mandatory Instruct-variant gate (STACK §3.3 — chat template required).
All configuration comes from cfg.model.* — zero hardcoded strings here.
"""
from __future__ import annotations

import logging
from typing import Tuple

from omegaconf import DictConfig

log = logging.getLogger("fathom.train")


def load_model_and_tokenizer(cfg: DictConfig) -> Tuple:
    """Load Unsloth-patched model + tokenizer with LoRA attached.

    Args:
        cfg: OmegaConf DictConfig with cfg.model.{name, lora_rank, lora_alpha,
             max_seq_length, load_in_4bit, target_modules} and cfg.seed.

    Returns:
        (model, tokenizer) — model is Unsloth-patched PeftModel, tokenizer has chat_template.

    Raises:
        ValueError: if cfg.model.name does not contain 'Instruct' (TRN-01 gate).
    """
    # TRN-01 gate: base model must be Instruct variant (STACK §3.3)
    name = str(cfg.model.name)
    if "Instruct" not in name:
        raise ValueError(
            f"TRN-01 gate: base model must be Instruct variant "
            f"(chat template required); got {name}"
        )

    try:
        # Lazy import: Unsloth is expensive + CUDA-side-effectful at import time
        from unsloth import FastLanguageModel  # type: ignore

        model, tokenizer = FastLanguageModel.from_pretrained(
            model_name=name,
            max_seq_length=int(cfg.model.max_seq_length),
            load_in_4bit=bool(cfg.model.load_in_4bit),
            dtype=None,
        )
        model = FastLanguageModel.get_peft_model(
            model,
            r=int(cfg.model.lora_rank),
            lora_alpha=int(cfg.model.lora_alpha),
            target_modules=str(cfg.model.target_modules),  # "all-linear" string works
            lora_dropout=0.0,
            bias="none",
            use_gradient_checkpointing="unsloth",
            random_state=int(cfg.seed),
        )
        log.info(
            "TRN-01 loaded with Unsloth: %s lora_rank=%d alpha=%d 4bit=%s",
            name,
            cfg.model.lora_rank,
            cfg.model.lora_alpha,
            cfg.model.load_in_4bit,
        )
        return model, tokenizer
    except Exception as e:
        # HF Jobs can hit version skew between Unsloth/TRl/Transformers.
        # Keep smoke and baseline training alive by falling back to vanilla HF+PEFT.
        log.warning("TRN-01 Unsloth path unavailable (%s) — falling back to HF+PEFT", e)

    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig  # type: ignore
    from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training  # type: ignore
    import torch  # type: ignore

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

    tokenizer = AutoTokenizer.from_pretrained(name, use_fast=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    # IMPORTANT: 4-bit/8-bit models must live on a SINGLE device to be trained
    # by accelerate. `device_map="auto"` shards across multi-GPU boxes (and
    # silently spills to CPU when its allocator under-counts), which causes
    # `accelerator.prepare()` to raise: "You can't train a model that has been
    # loaded in 8-bit or 4-bit precision on a different device than the one
    # you're training on." Pin to current CUDA device.
    if torch.cuda.is_available():
        device_map = {"": torch.cuda.current_device()}
    else:
        device_map = None

    model = AutoModelForCausalLM.from_pretrained(
        name,
        quantization_config=quantization_config,
        torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
        device_map=device_map,
    )

    if bool(cfg.model.load_in_4bit):
        model = prepare_model_for_kbit_training(model)

    peft_config = LoraConfig(
        r=int(cfg.model.lora_rank),
        lora_alpha=int(cfg.model.lora_alpha),
        target_modules=str(cfg.model.target_modules),
        lora_dropout=0.0,
        bias="none",
        task_type="CAUSAL_LM",
    )
    model = get_peft_model(model, peft_config)

    log.info(
        "TRN-01 loaded with HF+PEFT fallback: %s lora_rank=%d alpha=%d 4bit=%s",
        name,
        cfg.model.lora_rank,
        cfg.model.lora_alpha,
        cfg.model.load_in_4bit,
    )
    return model, tokenizer


__all__ = ["load_model_and_tokenizer"]