From ce758f7eb58a0508f074a65fe23a541282f5906c Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Wed, 1 Jul 2026 14:00:41 -0400
Subject: [PATCH] feat: improve loading strategies with low_cpu_mem_usage and
 proper DeepSpeed config

---
 training/scripts/train.py | 67 +++++++++++++++++++++++++++------------
 1 file changed, 46 insertions(+), 21 deletions(-)

diff --git a/training/scripts/train.py b/training/scripts/train.py
index 016717e..a432f72 100755
--- a/training/scripts/train.py
+++ b/training/scripts/train.py
@@ -39,12 +39,19 @@ def train(config_path):
     # Try multiple loading strategies in order
     print(f"Loading model: {config['base_model']}")
     
-    # Strategy 1: Load 4-bit AS-IS
-    print("\n[1/4] Trying: 4-bit model AS-IS...")
+    # Strategy 1: Load 4-bit with BitsAndBytes
+    print("\n[1/4] Trying: 4-bit with BitsAndBytes...")
     try:
+        from transformers import BitsAndBytesConfig
+        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"],
-            torch_dtype=torch.float16,
+            quantization_config=bnb_config,
             device_map="auto",
             trust_remote_code=True,
         )
@@ -52,49 +59,67 @@ def train(config_path):
     except Exception as e:
         print(f"✗ Failed: {e}")
         
-        # Strategy 2: Load bf16 to CPU
+        # Strategy 2: Load bf16 to CPU with low_cpu_mem_usage
         print("\n[2/4] Trying: bf16 model to CPU...")
         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 model loaded to CPU")
         except Exception as e:
             print(f"✗ Failed: {e}")
             
-            # Strategy 3: Load with accelerate CPU offload
-            print("\n[3/4] Trying: bf16 with accelerate CPU offload...")
+            # Strategy 3: Load fp16 with auto placement
+            print("\n[3/4] Trying: fp16 model with auto placement...")
             try:
-                from accelerate import load_checkpoint_and_dispatch
-                base_model = AutoModelForCausalLM.from_pretrained(
+                model = AutoModelForCausalLM.from_pretrained(
                     config["base_model"],
-                    torch_dtype=torch.bfloat16,
+                    torch_dtype=torch.float16,
+                    device_map="auto",
+                    low_cpu_mem_usage=True,
                     trust_remote_code=True,
                 )
-                model = load_checkpoint_and_dispatch(
-                    base_model,
-                    checkpoint=config["base_model"],
-                    device_map="auto",
-                    dtype=torch.bfloat16,
-                    offload_folder="/tmp/model_offload",
-                )
-                print("✓ Success: bf16 with accelerate offload")
+                print("✓ Success: fp16 model loaded")
             except Exception as e:
                 print(f"✗ Failed: {e}")
                 
-                # Strategy 4: Use bf16 with DeepSpeed ZeRO-3
-                print("\n[4/4] Trying: bf16 with DeepSpeed ZeRO-3 CPU offload...")
+                # Strategy 4: DeepSpeed ZeRO-3 with CPU offload
+                print("\n[4/4] Trying: DeepSpeed ZeRO-3 CPU offload...")
                 try:
+                    import deepspeed
                     model = AutoModelForCausalLM.from_pretrained(
                         config["base_model"],
                         torch_dtype=torch.bfloat16,
-                        device_map="cpu",
+                        device_map=None,
+                        low_cpu_mem_usage=True,
                         trust_remote_code=True,
                     )
-                    print("✓ Success: bf16 model (DeepSpeed will handle offload)")
+                    
+                    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=1e-5)
+                    model, optimizer, _, _ = deepspeed.initialize(
+                        model=model,
+                        model_parameters=model.parameters(),
+                        optimizer=optimizer,
+                        config=ds_config,
+                    )
+                    print("✓ Success: DeepSpeed ZeRO-3 loaded")
                 except Exception as e:
                     print(f"✗ Failed: {e}")
                     raise RuntimeError("All loading strategies failed!")