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

refactor: cleaner loading strategies with error tracking

Browse Source
This commit is contained in:
Christian Medina
2026-07-01 21:38:38 -04:00
parent 815115e97a
commit 9d8c2b6cf0
1 changed files with 39 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

57
train.py
View File

@@ -32,68 +32,89 @@ def train(config_path):
print(f"Loading model: {config['base_model']}") print(f"Loading model: {config['base_model']}")
# Load model - try multiple strategies # Load model with multiple strategies
print(f"\n[INFO] Loading {config['base_model']}...") print(f"\n[INFO] Loading {config['base_model']}...")
errors = []
# Strategy 1: bf16 model with BnB 4-bit quantization # ------------------------------------------------------------------
print("\n[1/4] Trying: bf16 with BnB 4-bit...") # Strategy 1: QLoRA (preferred)
# ------------------------------------------------------------------
print("\n[1/4] Trying: 4-bit NF4 on GPU...")
try: try:
from transformers import BitsAndBytesConfig
bnb_config = BitsAndBytesConfig( bnb_config = BitsAndBytesConfig(
load_in_4bit=True, load_in_4bit=True,
bnb_4bit_quant_type="nf4", bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_compute_dtype=torch.bfloat16,
bnb_4bit_use_double_quant=True,
) )
model = AutoModelForCausalLM.from_pretrained( model = AutoModelForCausalLM.from_pretrained(
config["base_model"], config["base_model"],
quantization_config=bnb_config, quantization_config=bnb_config,
device_map="auto", device_map="auto",
trust_remote_code=True, trust_remote_code=True,
low_cpu_mem_usage=True,
) )
print("✓ Success: bf16 with BnB 4-bit") print("✓ Success: QLoRA 4-bit")
except Exception as e: except Exception as e:
errors.append(("QLoRA 4-bit", e))
print(f"✗ Failed: {e}") print(f"✗ Failed: {e}")
# Strategy 2: bf16 to CPU # --------------------------------------------------------------
print("\n[2/4] Trying: bf16 to CPU...") # Strategy 2: BF16 GPU
# --------------------------------------------------------------
print("\n[2/4] Trying: bf16 GPU...")
try: try:
model = AutoModelForCausalLM.from_pretrained( model = AutoModelForCausalLM.from_pretrained(
config["base_model"], config["base_model"],
torch_dtype=torch.bfloat16, torch_dtype=torch.bfloat16,
low_cpu_mem_usage=True, device_map="auto",
device_map="cpu",
trust_remote_code=True, trust_remote_code=True,
low_cpu_mem_usage=True,
) )
print("✓ Success: bf16 to CPU") print("✓ Success: bf16 GPU")
except Exception as e: except Exception as e:
errors.append(("bf16 GPU", e))
print(f"✗ Failed: {e}") print(f"✗ Failed: {e}")
# Strategy 3: 4-bit to CPU # ----------------------------------------------------------
print("\n[3/4] Trying: 4-bit to CPU...") # Strategy 3: BF16 CPU
# ----------------------------------------------------------
print("\n[3/4] Trying: bf16 CPU...")
try: try:
model = AutoModelForCausalLM.from_pretrained( model = AutoModelForCausalLM.from_pretrained(
config["base_model"], config["base_model"],
torch_dtype=torch.float16, torch_dtype=torch.bfloat16,
device_map="cpu", device_map="cpu",
trust_remote_code=True, trust_remote_code=True,
low_cpu_mem_usage=True,
) )
print("✓ Success: 4-bit to CPU") print("✓ Success: bf16 CPU")
except Exception as e: except Exception as e:
errors.append(("bf16 CPU", e))
print(f"✗ Failed: {e}") print(f"✗ Failed: {e}")
# Strategy 4: 4-bit auto # ------------------------------------------------------
print("\n[4/4] Trying: 4-bit auto...") # Strategy 4: FP16 GPU
# ------------------------------------------------------
print("\n[4/4] Trying: fp16 GPU...")
try: try:
model = AutoModelForCausalLM.from_pretrained( model = AutoModelForCausalLM.from_pretrained(
config["base_model"], config["base_model"],
torch_dtype=torch.float16, torch_dtype=torch.float16,
device_map="auto", device_map="auto",
trust_remote_code=True, trust_remote_code=True,
low_cpu_mem_usage=True,
) )
print("✓ Success: 4-bit auto") print("✓ Success: fp16 GPU")
except Exception as e: except Exception as e:
errors.append(("fp16 GPU", e))
print(f"✗ Failed: {e}") print(f"✗ Failed: {e}")
raise RuntimeError("All loading strategies failed!") msg = "\n".join(
f"{name}: {err}" for name, err in errors
)
raise RuntimeError(
f"All loading strategies failed:\n\n{msg}"
)
# Add LoRA # Add LoRA
lora_config = LoraConfig( lora_config = LoraConfig(