fix: use accelerate load_checkpoint_and_dispatch for proper CPU offload

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