From e967bd74f1589bc61a9b14999cc7d9fd26cc75c6 Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Wed, 1 Jul 2026 13:53:06 -0400
Subject: [PATCH] fix: skip quantization, use DeepSpeed ZeRO-3 CPU offload
 directly

---
 training/scripts/train.py | 59 ++++++---------------------------------
 1 file changed, 9 insertions(+), 50 deletions(-)

diff --git a/training/scripts/train.py b/training/scripts/train.py
index fa32219..0c51dab 100755
--- a/training/scripts/train.py
+++ b/training/scripts/train.py
@@ -36,56 +36,15 @@ def train(config_path):
     print(f"Loading model: {config['base_model']}")
     from transformers import BitsAndBytesConfig, AutoConfig
     
-    try:
-        # Try 4-bit QLoRA first
-        bnb_config = BitsAndBytesConfig(
-            load_in_4bit=True,
-            bnb_4bit_quant_type="nf4",
-            bnb_4bit_compute_dtype=torch.bfloat16,
-            bnb_4bit_use_double_quant=True,
-        )
-        model = AutoModelForCausalLM.from_pretrained(
-            config["base_model"],
-            quantization_config=bnb_config,
-            dtype=torch.bfloat16,
-            device_map="auto",
-            trust_remote_code=True,
-        )
-        print("Model loaded with QLoRA (4-bit).")
-    except Exception as e:
-        print(f"4-bit failed: {e}")
-        try:
-            # Try 8-bit with CPU offload
-            print("Trying 8-bit with CPU offload...")
-            bnb_config_8bit = BitsAndBytesConfig(
-                load_in_8bit=True,
-                llm_int8_enable_fp32_cpu_offload=True,
-            )
-            model = AutoModelForCausalLM.from_pretrained(
-                config["base_model"],
-                quantization_config=bnb_config_8bit,
-                device_map="auto",
-                trust_remote_code=True,
-            )
-            print("Model loaded with 8-bit CPU offload.")
-        except Exception as e2:
-            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,
-                trust_remote_code=True,
-            )
-            # 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.")
+    # Skip quantization - use DeepSpeed ZeRO-3 with CPU offload
+    print(f"Loading model: {config['base_model']}")
+    model = AutoModelForCausalLM.from_pretrained(
+        config["base_model"],
+        torch_dtype=torch.bfloat16,
+        device_map="cpu",  # Load to CPU first
+        trust_remote_code=True,
+    )
+    print("Model loaded to CPU. DeepSpeed will distribute.")
     
     # Prepare model for k-bit training
     from peft import prepare_model_for_kbit_training