From d9fd079686913fb3a513fa1671d839aae8e57734 Mon Sep 17 00:00:00 2001 From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com> Date: Wed, 1 Jul 2026 15:14:00 -0400 Subject: [PATCH] fix: add accelerate CPU offload to bf16 strategies --- training/scripts/train.py | 32 ++++++++++++++++++++++++-------- 1 file changed, 24 insertions(+), 8 deletions(-) diff --git a/training/scripts/train.py b/training/scripts/train.py index 423afd3..e393685 100755 --- a/training/scripts/train.py +++ b/training/scripts/train.py @@ -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!")