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

fix: add BitsAndBytesConfig import + remove bf16 GPU strategy

Browse Source
This commit is contained in:
Christian Medina
2026-07-02 08:54:41 -04:00
parent 70a2bf0dd6
commit 4c9522072e
1 changed files with 18 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

35
train.py
View File

@@ -22,6 +22,7 @@ def train(config_path):
from transformers import ( from transformers import (
AutoModelForCausalLM, AutoModelForCausalLM,
AutoTokenizer, AutoTokenizer,
BitsAndBytesConfig,
TrainingArguments, TrainingArguments,
) )
from peft import ( from peft import (
@@ -60,54 +61,54 @@ def train(config_path):
print(f"✗ Failed: {e}") print(f"✗ Failed: {e}")
# -------------------------------------------------------------- # --------------------------------------------------------------
# Strategy 2: BF16 GPU # Strategy 2: BF16 CPU (model too large for single GPU)
# -------------------------------------------------------------- # --------------------------------------------------------------
print("\n[2/4] Trying: bf16 GPU...") print("\n[2/4] Trying: bf16 CPU...")
try: try:
model = AutoModelForCausalLM.from_pretrained( model = AutoModelForCausalLM.from_pretrained(
config["base_model"], config["base_model"],
torch_dtype=torch.bfloat16, torch_dtype=torch.bfloat16,
device_map="auto", device_map="cpu",
trust_remote_code=True,
low_cpu_mem_usage=True, low_cpu_mem_usage=True,
trust_remote_code=True,
) )
print("✓ Success: bf16 GPU") print("✓ Success: bf16 CPU")
except Exception as e: except Exception as e:
errors.append(("bf16 GPU", e)) errors.append(("bf16 CPU", e))
print(f"✗ Failed: {e}") print(f"✗ Failed: {e}")
# ---------------------------------------------------------- # ----------------------------------------------------------
# Strategy 3: BF16 CPU # Strategy 3: FP16 CPU (fallback)
# ---------------------------------------------------------- # ----------------------------------------------------------
print("\n[3/4] Trying: bf16 CPU...") print("\n[3/4] Trying: fp16 CPU...")
try: try:
model = AutoModelForCausalLM.from_pretrained( model = AutoModelForCausalLM.from_pretrained(
config["base_model"], config["base_model"],
torch_dtype=torch.bfloat16, torch_dtype=torch.float16,
device_map="cpu", device_map="cpu",
trust_remote_code=True, trust_remote_code=True,
low_cpu_mem_usage=True, low_cpu_mem_usage=True,
) )
print("✓ Success: bf16 CPU") print("✓ Success: fp16 CPU")
except Exception as e: except Exception as e:
errors.append(("bf16 CPU", e)) errors.append(("fp16 CPU", e))
print(f"✗ Failed: {e}") print(f"✗ Failed: {e}")
# ------------------------------------------------------ # ------------------------------------------------------
# Strategy 4: FP16 GPU # Strategy 4: 4-bit AS-IS (already quantized)
# ------------------------------------------------------ # ----------------------------------------------------------
print("\n[4/4] Trying: fp16 GPU...") print("\n[4/4] Trying: 4-bit AS-IS...")
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="cpu",
trust_remote_code=True, trust_remote_code=True,
low_cpu_mem_usage=True, low_cpu_mem_usage=True,
) )
print("✓ Success: fp16 GPU") print("✓ Success: 4-bit AS-IS")
except Exception as e: except Exception as e:
errors.append(("fp16 GPU", e)) errors.append(("4-bit AS-IS", e))
print(f"✗ Failed: {e}") print(f"✗ Failed: {e}")
msg = "\n".join( msg = "\n".join(
f"{name}: {err}" for name, err in errors f"{name}: {err}" for name, err in errors