File size: 4,669 Bytes
071ba6b 8787bd3 7f150a7 8787bd3 7f150a7 071ba6b 8787bd3 071ba6b 8787bd3 071ba6b 8787bd3 071ba6b 8787bd3 071ba6b 8787bd3 071ba6b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 | """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"]
|