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

train-on-this-server.sh

@@ -18,9 +18,9 @@ pip install --upgrade pip
 echo "Installing training dependencies..."
 pip install transformers datasets trl peft accelerate bitsandbytes deepspeed
 
-# Run training with CUDA memory optimization
+# Run training with DeepSpeed ZeRO-3
-echo "Starting training with accelerate..."
+echo "Starting training with DeepSpeed ZeRO-3..."
 export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
-accelerate launch --multi_gpu --num_processes 2 train.py --config training/configs/ornith-35b-lora.yaml
+torchrun --nproc_per_node=2 train.py --config training/configs/ornith-35b-lora.yaml
 
 echo "Training completed!"

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",
+        bf16=True,
-        fp16=config.get("mixed_precision", "bf16") == "fp16",
         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)

training/configs/ds_zero3.json

@@ -0,0 +1,22 @@
+{
+  "bf16": {
+    "enabled": true
+  },
+
+  "zero_optimization": {
+    "stage": 3,
+
+    "overlap_comm": true,
+    "contiguous_gradients": true,
+
+    "reduce_bucket_size": 50000000,
+    "stage3_prefetch_bucket_size": 50000000,
+    "stage3_param_persistence_threshold": 100000
+  },
+
+  "gradient_clipping": 1.0,
+
+  "train_micro_batch_size_per_gpu": 1,
+
+  "steps_per_print": 100
+}