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

fix: use accelerate load_checkpoint_and_dispatch for proper CPU offload

Browse Source
This commit is contained in:
Christian Medina
2026-07-01 13:39:51 -04:00
parent 528e7002ef
commit 7539e39a93
1 changed files with 12 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
training/scripts/train.py
View File

@@ -69,14 +69,23 @@ def train(config_path):
) )
print("Model loaded with 8-bit CPU offload.") print("Model loaded with 8-bit CPU offload.")
except Exception as e2: except Exception as e2:
print(f"8-bit failed: {e2}, falling back to bf16 with DeepSpeed CPU offload") print(f"8-bit failed: {e2}, falling back to bf16 with accelerate CPU offload")
# Use accelerate to load with CPU offload
from accelerate import load_checkpoint_and_dispatch
model = AutoModelForCausalLM.from_pretrained( model = AutoModelForCausalLM.from_pretrained(
config["base_model"], config["base_model"],
torch_dtype=torch.bfloat16, torch_dtype=torch.bfloat16,
device_map="cpu",
trust_remote_code=True, trust_remote_code=True,
) )
print("Model loaded as bf16 (DeepSpeed CPU offload).") # Load with CPU offload
model = load_checkpoint_and_dispatch(
model,
checkpoint=config["base_model"],
device_map="auto",
dtype=torch.bfloat16,
offload_folder="/tmp/model_offload",
)
print("Model loaded with accelerate CPU offload.")
# 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