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

feat: restore all 6 loading strategies

Browse Source
This commit is contained in:
Christian Medina
2026-07-01 15:11:45 -04:00
parent 313b44381f
commit 4ed80ad3a3
1 changed files with 42 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

56
training/scripts/train.py
View File

@@ -59,8 +59,8 @@ def train(config_path):
except Exception as e: except Exception as e:
print(f"✗ Failed: {e}") print(f"✗ Failed: {e}")
# Strategy 1: Load 4-bit model AS-IS (no quantization, no DeepSpeed) # Strategy 1: 4-bit AS-IS with auto device_map
print("\n[1/2] Trying: 4-bit model AS-IS...") print("\n[1/6] Trying: 4-bit AS-IS...")
try: try:
model = AutoModelForCausalLM.from_pretrained( model = AutoModelForCausalLM.from_pretrained(
config["base_model"], config["base_model"],
@@ -68,24 +68,52 @@ def train(config_path):
device_map="auto", device_map="auto",
trust_remote_code=True, trust_remote_code=True,
) )
print("✓ Success: 4-bit model loaded") print("✓ Success: 4-bit AS-IS")
except Exception as e: except Exception as e:
print(f"✗ Failed: {e}") print(f"✗ Failed: {e}")
# Strategy 2: bf16 to CPU # Strategy 2: 4-bit with bf16 compute
print("\n[2/2] Trying: bf16 model to CPU...") print("\n[2/6] Trying: 4-bit with bf16 compute...")
try: try:
model = AutoModelForCausalLM.from_pretrained( from transformers import BitsAndBytesConfig
config["base_model"], bnb_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=True)
torch_dtype=torch.bfloat16, model = AutoModelForCausalLM.from_pretrained(config["base_model"], quantization_config=bnb_config, device_map="auto", trust_remote_code=True)
device_map="cpu", print("✓ Success: 4-bit with bf16 compute")
low_cpu_mem_usage=True,
trust_remote_code=True,
)
print("✓ Success: bf16 model loaded to CPU")
except Exception as e: except Exception as e:
print(f"✗ Failed: {e}") print(f"✗ Failed: {e}")
raise RuntimeError("All loading strategies failed!")
# Strategy 3: 4-bit to CPU
print("\n[3/6] Trying: 4-bit to CPU...")
try:
model = AutoModelForCausalLM.from_pretrained(config["base_model"], torch_dtype=torch.float16, device_map="cpu", trust_remote_code=True)
print("✓ Success: 4-bit to CPU")
except Exception as e:
print(f"✗ Failed: {e}")
# Strategy 4: 4-bit fp32
print("\n[4/6] Trying: 4-bit fp32...")
try:
model = AutoModelForCausalLM.from_pretrained(config["base_model"], torch_dtype=torch.float32, device_map="auto", trust_remote_code=True)
print("✓ Success: 4-bit fp32")
except Exception as e:
print(f"✗ Failed: {e}")
# Strategy 5: bf16 to CPU
print("\n[5/6] Trying: bf16 to CPU...")
try:
model = AutoModelForCausalLM.from_pretrained(config["base_model"], torch_dtype=torch.bfloat16, device_map="cpu", low_cpu_mem_usage=True, trust_remote_code=True)
print("✓ Success: bf16 to CPU")
except Exception as e:
print(f"✗ Failed: {e}")
# Strategy 6: bf16 auto
print("\n[6/6] Trying: bf16 auto...")
try:
model = AutoModelForCausalLM.from_pretrained(config["base_model"], torch_dtype=torch.bfloat16, device_map="auto", low_cpu_mem_usage=True, trust_remote_code=True)
print("✓ Success: bf16 auto")
except Exception as e:
print(f"✗ Failed: {e}")
raise RuntimeError("All loading strategies failed!")
# Prepare model for k-bit training # Prepare model for k-bit training
from peft import prepare_model_for_kbit_training from peft import prepare_model_for_kbit_training