From 22d1bfd573fa77461cc4f92fbbd7094c457dad33 Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Wed, 1 Jul 2026 14:43:30 -0400
Subject: [PATCH] fix: use 4-bit model first, then bf16 with CPU offload

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

diff --git a/training/scripts/train.py b/training/scripts/train.py
index 95fef47..8205fa4 100755
--- a/training/scripts/train.py
+++ b/training/scripts/train.py
@@ -59,15 +59,14 @@ def train(config_path):
     except Exception as e:
         print(f"✗ Failed: {e}")
         
-        # Strategy 2: DeepSpeed ZeRO-3 (distribute across GPUs)
-        print("\n[2/4] Trying: DeepSpeed ZeRO-3...")
+        # Strategy 1: Load 4-bit model AS-IS with DeepSpeed ZeRO-3
+        print("\n[1/3] Trying: 4-bit model with DeepSpeed ZeRO-3...")
         try:
             import deepspeed
             model = AutoModelForCausalLM.from_pretrained(
                 config["base_model"],
-                torch_dtype=torch.bfloat16,
-                device_map=None,
-                low_cpu_mem_usage=True,
+                torch_dtype=torch.float16,  # Load as fp16 (already 4-bit)
+                device_map="auto",
                 trust_remote_code=True,
             )
             
@@ -81,7 +80,7 @@ def train(config_path):
                     "offload_optimizer": {"device": "cpu", "pin_memory": True},
                     "offload_param": {"device": "cpu", "pin_memory": True},
                 },
-                "bf16": {"enabled": True},
+                "fp16": {"enabled": True},
             }
             
             optimizer = torch.optim.AdamW(model.parameters(), lr=float(config["train_params"]["learning_rate"]))
@@ -91,10 +90,45 @@ def train(config_path):
                 optimizer=optimizer,
                 config=ds_config,
             )
-            print("✓ Success: DeepSpeed ZeRO-3 loaded")
+            print("✓ Success: 4-bit model with DeepSpeed ZeRO-3")
         except Exception as e:
             print(f"✗ Failed: {e}")
-            raise RuntimeError("All loading strategies failed!")
+            
+            # Strategy 2: bf16 with DeepSpeed ZeRO-3 + CPU offload
+            print("\n[2/3] Trying: bf16 with DeepSpeed ZeRO-3 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,
+                )
+                
+                ds_config = {
+                    "train_micro_batch_size_per_gpu": 1,
+                    "gradient_accumulation_steps": 1,
+                    "zero_optimization": {
+                        "stage": 3,
+                        "contiguous_gradients": True,
+                        "overlap_comm": True,
+                        "offload_optimizer": {"device": "cpu", "pin_memory": True},
+                        "offload_param": {"device": "cpu", "pin_memory": True},
+                    },
+                    "bf16": {"enabled": True},
+                }
+                
+                optimizer = torch.optim.AdamW(model.parameters(), lr=float(config["train_params"]["learning_rate"]))
+                model, optimizer, _, _ = deepspeed.initialize(
+                    model=model,
+                    model_parameters=model.parameters(),
+                    optimizer=optimizer,
+                    config=ds_config,
+                )
+                print("✓ Success: bf16 with DeepSpeed ZeRO-3 CPU offload")
+            except Exception as e2:
+                print(f"✗ Failed: {e2}")
+                raise RuntimeError("All loading strategies failed!")
     
     # Prepare model for k-bit training
     from peft import prepare_model_for_kbit_training