From e64515e3ca557243a7513abc9d19a4b146eb65e8 Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Wed, 1 Jul 2026 22:33:06 -0400
Subject: [PATCH] refactor: switch from DeepSpeed ZeRO-3 to FSDP for QLoRA
 compatibility

---
 train-on-this-server.sh |  6 +++---
 train.py                | 16 +++++++++++-----
 2 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/train-on-this-server.sh b/train-on-this-server.sh
index 34ccce4..1728bed 100755
--- a/train-on-this-server.sh
+++ b/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 DeepSpeed ZeRO-3
-echo "Starting training with DeepSpeed ZeRO-3..."
+# Run training with FSDP
+echo "Starting training with FSDP..."
 export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
-torchrun --nproc_per_node=2 train.py --config training/configs/ornith-35b-lora.yaml
+accelerate launch --multi_gpu --num_processes 2 train.py --config training/configs/ornith-35b-lora.yaml
 
 echo "Training completed!"
\ No newline at end of file
diff --git a/train.py b/train.py
index 8f40d20..931b116 100644
--- a/train.py
+++ b/train.py
@@ -37,9 +37,9 @@ def train(config_path):
     errors = []
     
     # ------------------------------------------------------------------
-    # Strategy 1: QLoRA with DeepSpeed ZeRO-3 (preferred)
+    # Strategy 1: QLoRA with FSDP (preferred)
     # ------------------------------------------------------------------
-    print("\n[1/4] Trying: 4-bit QLoRA (DeepSpeed ZeRO-3)...")
+    print("\n[1/4] Trying: 4-bit QLoRA (FSDP)...")
     try:
         bnb_config = BitsAndBytesConfig(
             load_in_4bit=True,
@@ -50,10 +50,11 @@ 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 (DeepSpeed will handle placement)")
+        print("✓ Success: QLoRA 4-bit")
     except Exception as e:
         errors.append(("QLoRA 4-bit", e))
         print(f"✗ Failed: {e}")
@@ -149,7 +150,7 @@ def train(config_path):
         },
     )
     
-    # Training arguments with DeepSpeed ZeRO-3
+    # Training arguments with FSDP
     training_args = TrainingArguments(
         output_dir=config["train_params"]["output_dir"],
         num_train_epochs=config["train_params"]["num_train_epochs"],
@@ -166,7 +167,12 @@ def train(config_path):
         eval_steps=config.get("eval_steps", 100),
         bf16=True,
         gradient_checkpointing=config.get("gradient_checkpointing", True),
-        deepspeed=config.get("deepspeed_config_path", "training/configs/ds_zero3.json"),
+        fsdp="full_shard auto_wrap",
+        fsdp_config={
+            "backward_prefetch": "backward_pre",
+            "forward_prefetch": "true",
+            "cpu_offload": "false",
+        },
     )
     
     # SFT Trainer (DeepSpeed handles distributed training via config)