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

fix: use aggressive CPU loading for memory optimization

Browse Source
This commit is contained in:
Christian Medina
2026-07-01 18:00:55 -04:00
parent 83b45c2a5c
commit f18a9f0973
1 changed files with 9 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files

59
train.py
View File

@@ -32,63 +32,22 @@ 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 memory optimization
print(f"\n[INFO] Loading {config['base_model']}...") print(f"\n[INFO] Loading {config['base_model']} with memory optimization...")
# Strategy 1: 4-bit AS-IS (already quantized) # Direct low-memory load for large model + DeepSpeed
print("\n[1/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.bfloat16,
# device_map="auto", low_cpu_mem_usage=True,
device_map="cpu", # Force CPU first
trust_remote_code=True, trust_remote_code=True,
) )
print("✓ Success: 4-bit AS-IS") print("✓ Success: Loaded on CPU with low memory usage")
except Exception as e: except Exception as e:
print(f"✗ Failed: {e}") print(f"✗ Failed initial load: {e}")
raise
# Strategy 2: 4-bit to CPU
print("\n[2/4] 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 3: bf16 to CPU
print("\n[3/4] 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 4: bf16 auto
print("\n[4/4] 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!")
# Add LoRA # Add LoRA
lora_config = LoraConfig( lora_config = LoraConfig(