cmedina/agenx-lora-training
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add fallback to bf16 with CPU offload if 4-bit fails

Browse Source
This commit is contained in:
Christian Medina
2026-07-01 13:29:23 -04:00
parent b41d99f956
commit 59ddfb45ed
1 changed files with 25 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

46
training/scripts/train.py
View File

@@ -32,30 +32,34 @@ def train(config_path):
print(f"Loading model: {config['base_model']}")
# Load model (skip broken quantization config)
# Load model - try 4-bit first, fall back to bf16 with CPU offload
print(f"Loading model: {config['base_model']}")
from transformers import BitsAndBytesConfig, AutoConfig
# Remove broken quantization config
model_config = AutoConfig.from_pretrained(config["base_model"])
model_config.quantization_config = None
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16,
bnb_4bit_use_double_quant=True,
)
# Load with fresh BnB config
model = AutoModelForCausalLM.from_pretrained(
config["base_model"],
quantization_config=bnb_config,
torch_dtype=torch.bfloat16,
device_map="auto",
trust_remote_code=True,
)
print("Model loaded with QLoRA (4-bit).")
try:
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16,
bnb_4bit_use_double_quant=True,
)
model = AutoModelForCausalLM.from_pretrained(
config["base_model"],
quantization_config=bnb_config,
dtype=torch.bfloat16,
device_map="auto",
trust_remote_code=True,
)
print("Model loaded with QLoRA (4-bit).")
except Exception as e:
print(f"4-bit failed: {e}, falling back to bf16 with CPU offload")
model = AutoModelForCausalLM.from_pretrained(
config["base_model"],
torch_dtype=torch.bfloat16,
device_map="cpu",
trust_remote_code=True,
)
print("Model loaded as bf16 (CPU offload).")
# Prepare model for k-bit training
from peft import prepare_model_for_kbit_training