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

fix: add accelerate CPU offload to bf16 strategies

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

32
training/scripts/train.py
View File

@@ -98,19 +98,35 @@ def train(config_path):
except Exception as e:
print(f"✗ Failed: {e}")
# Strategy 5: bf16 to CPU
print("\n[5/6] Trying: bf16 to CPU...")
# Strategy 5: bf16 with accelerate CPU offload
print("\n[5/6] Trying: bf16 with accelerate CPU offload...")
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")
from accelerate import Accelerator
model = AutoModelForCausalLM.from_pretrained(
config["base_model"],
torch_dtype=torch.bfloat16,
device_map="cpu",
low_cpu_mem_usage=True,
trust_remote_code=True,
)
accelerator = Accelerator(cpu_offload=True)
model = accelerator.prepare(model)
print("✓ Success: bf16 with accelerate CPU offload")
except Exception as e:
print(f"✗ Failed: {e}")
# Strategy 6: bf16 auto
print("\n[6/6] Trying: bf16 auto...")
# Strategy 6: bf16 to CPU with gradient checkpointing
print("\n[6/6] Trying: bf16 to CPU with gradient checkpointing...")
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")
model = AutoModelForCausalLM.from_pretrained(
config["base_model"],
torch_dtype=torch.bfloat16,
device_map="cpu",
low_cpu_mem_usage=True,
trust_remote_code=True,
)
model.gradient_checkpointing_enable()
print("✓ Success: bf16 to CPU with gradient checkpointing")
except Exception as e:
print(f"✗ Failed: {e}")
raise RuntimeError("All loading strategies failed!")