fix: use 4-bit model first, then bf16 with CPU offload

training/scripts/train.py

@@ -59,14 +59,48 @@ 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.float16,  # Load as fp16 (already 4-bit)
+                device_map="auto",
+                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},
+                },
+                "fp16": {"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: 4-bit model with DeepSpeed ZeRO-3")
+        except Exception as e:
+            print(f"✗ Failed: {e}")
+            
+            # 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=None,
+                    device_map="cpu",
                     low_cpu_mem_usage=True,
                     trust_remote_code=True,
                 )
@@ -91,9 +125,9 @@ def train(config_path):
                     optimizer=optimizer,
                     config=ds_config,
                 )
-            print("✓ Success: DeepSpeed ZeRO-3 loaded")
-        except Exception as e:
-            print(f"✗ Failed: {e}")
+                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