From 1b3f678b50e5ee60912b39a96f50ab5ae0ecd944 Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Thu, 2 Jul 2026 11:35:09 -0400
Subject: [PATCH] feat: add Strategy 1 - 4-bit QLoRA with device_map=auto
 (distributed across GPUs)

---
 train.py | 45 ++++++++++++++++++++++++++++++++++-----------
 1 file changed, 34 insertions(+), 11 deletions(-)

diff --git a/train.py b/train.py
index 165c813..79390e4 100644
--- a/train.py
+++ b/train.py
@@ -38,34 +38,57 @@ def train(config_path):
     errors = []
     
     # ------------------------------------------------------------------
-    # Strategy 1: QLoRA 4-bit with FSDP (load to CPU, FSDP shards across GPUs)
+    # Strategy 1: QLoRA 4-bit with device_map="auto" (distributed across GPUs, no FSDP)
     # ------------------------------------------------------------------
-    print("\n[1/4] Trying: 4-bit QLoRA (FSDP, load to CPU)...")
+    print("\n[1/4] Trying: 4-bit QLoRA (distributed across GPUs, no FSDP)...")
     try:
         bnb_config = BitsAndBytesConfig(
             load_in_4bit=True,
             bnb_4bit_quant_type="nf4",
             bnb_4bit_compute_dtype=torch.bfloat16,
             bnb_4bit_use_double_quant=True,
-            bnb_4bit_quant_storage=torch.bfloat16,  # Enable FSDP sharding of 4-bit weights
         )
         model = AutoModelForCausalLM.from_pretrained(
             config["base_model"],
             quantization_config=bnb_config,
-            device_map="cpu",  # Load to CPU, FSDP shards later
+            device_map="auto",  # Distribute layers across GPUs automatically
             trust_remote_code=True,
             low_cpu_mem_usage=True,
         )
-        print("✓ Success: QLoRA 4-bit loaded to CPU (FSDP will shard across GPUs)")
+        print("✓ Success: QLoRA 4-bit distributed across GPUs (no FSDP)")
     except Exception as e:
         errors.append(("QLoRA 4-bit", e))
         print(f"✗ Failed: {e}")
         
         # --------------------------------------------------------------
-        # Strategy 2: BF16 CPU (model too large for single GPU)
+        # Strategy 2: QLoRA 4-bit with FSDP (load to CPU, FSDP shards across GPUs)
         # --------------------------------------------------------------
-        print("\n[2/4] Trying: bf16 CPU...")
+        print("\n[2/4] Trying: 4-bit QLoRA (FSDP, load to CPU)...")
         try:
+            bnb_config = BitsAndBytesConfig(
+                load_in_4bit=True,
+                bnb_4bit_quant_type="nf4",
+                bnb_4bit_compute_dtype=torch.bfloat16,
+                bnb_4bit_use_double_quant=True,
+                bnb_4bit_quant_storage=torch.bfloat16,  # Enable FSDP sharding of 4-bit weights
+            )
+            model = AutoModelForCausalLM.from_pretrained(
+                config["base_model"],
+                quantization_config=bnb_config,
+                device_map="cpu",  # Load to CPU, FSDP shards later
+                trust_remote_code=True,
+                low_cpu_mem_usage=True,
+            )
+            print("✓ Success: QLoRA 4-bit loaded to CPU (FSDP will shard across GPUs)")
+        except Exception as e:
+            errors.append(("QLoRA 4-bit FSDP", e))
+            print(f"✗ Failed: {e}")
+            
+            # --------------------------------------------------------------
+            # Strategy 3: BF16 CPU (model too large for single GPU)
+            # --------------------------------------------------------------
+            print("\n[3/4] Trying: bf16 CPU...")
+            try:
             model = AutoModelForCausalLM.from_pretrained(
                 config["base_model"],
                 torch_dtype=torch.bfloat16,
@@ -79,9 +102,9 @@ def train(config_path):
             print(f"✗ Failed: {e}")
             
             # ----------------------------------------------------------
-            # Strategy 3: FP16 CPU (fallback)
+            # Strategy 4: FP16 CPU (fallback)
             # ----------------------------------------------------------
-            print("\n[3/4] Trying: fp16 CPU...")
+            print("\n[4/4] Trying: fp16 CPU...")
             try:
                 model = AutoModelForCausalLM.from_pretrained(
                     config["base_model"],
@@ -96,9 +119,9 @@ def train(config_path):
                 print(f"✗ Failed: {e}")
                 
                 # ------------------------------------------------------
-                # Strategy 4: 4-bit AS-IS (already quantized)
+                # Strategy 5: 4-bit AS-IS (already quantized)
                 # ----------------------------------------------------------
-                print("\n[4/4] Trying: 4-bit AS-IS...")
+                print("\n[5/5] Trying: 4-bit AS-IS...")
                 try:
                     model = AutoModelForCausalLM.from_pretrained(
                         config["base_model"],