feat: use DeepSpeed ZeRO-3 for proper model sharding across GPUs

train.py

@@ -37,9 +37,9 @@ def train(config_path):
     errors = []
     
     # ------------------------------------------------------------------
-    # Strategy 1: QLoRA (preferred)
+    # Strategy 1: QLoRA with DeepSpeed ZeRO-3 (preferred)
     # ------------------------------------------------------------------
-    print("\n[1/4] Trying: 4-bit NF4 on GPU...")
+    print("\n[1/4] Trying: 4-bit QLoRA (DeepSpeed ZeRO-3)...")
     try:
         bnb_config = BitsAndBytesConfig(
             load_in_4bit=True,
@@ -50,11 +50,10 @@ def train(config_path):
         model = AutoModelForCausalLM.from_pretrained(
             config["base_model"],
             quantization_config=bnb_config,
-            device_map="auto",
             trust_remote_code=True,
             low_cpu_mem_usage=True,
         )
-        print("✓ Success: QLoRA 4-bit")
+        print("✓ Success: QLoRA 4-bit (DeepSpeed will handle placement)")
     except Exception as e:
         errors.append(("QLoRA 4-bit", e))
         print(f"✗ Failed: {e}")
@@ -150,7 +149,7 @@ def train(config_path):
         },
     )
     
-    # Training arguments (max_seq_length removed - passed to SFTTrainer instead)
+    # Training arguments with DeepSpeed ZeRO-3
     training_args = TrainingArguments(
         output_dir=config["train_params"]["output_dir"],
         num_train_epochs=config["train_params"]["num_train_epochs"],
@@ -165,9 +164,9 @@ def train(config_path):
         save_total_limit=config["train_params"]["save_total_limit"],
         eval_strategy=config.get("eval_strategy", "steps"),
         eval_steps=config.get("eval_steps", 100),
-        bf16=config.get("mixed_precision", "bf16") == "bf16",
-        fp16=config.get("mixed_precision", "bf16") == "fp16",
+        bf16=True,
         gradient_checkpointing=config.get("gradient_checkpointing", True),
+        deepspeed=config.get("deepspeed_config_path", "training/configs/ds_zero3.json"),
     )
     
     # SFT Trainer (DeepSpeed handles distributed training via config)