From 7539e39a933271b7ac2a78441ab9d94199428d21 Mon Sep 17 00:00:00 2001 From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com> Date: Wed, 1 Jul 2026 13:39:51 -0400 Subject: [PATCH] fix: use accelerate load_checkpoint_and_dispatch for proper CPU offload --- training/scripts/train.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/training/scripts/train.py b/training/scripts/train.py index 7231ad9..fa32219 100755 --- a/training/scripts/train.py +++ b/training/scripts/train.py @@ -69,14 +69,23 @@ def train(config_path): ) print("Model loaded with 8-bit CPU offload.") 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( config["base_model"], torch_dtype=torch.bfloat16, - device_map="cpu", 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 from peft import prepare_model_for_kbit_training