From 4f2f8ef03f44fb9fd2a8903b6d979dfa93253983 Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Tue, 30 Jun 2026 18:26:48 -0400
Subject: [PATCH] feat: use DeepSpeed ZeRO-3 for optimal 2-GPU training

---
 training/configs/llama2-7b-lora.yaml | 22 ++++++++++++++++------
 training/scripts/train.py            | 19 ++++---------------
 2 files changed, 20 insertions(+), 21 deletions(-)

diff --git a/training/configs/llama2-7b-lora.yaml b/training/configs/llama2-7b-lora.yaml
index 6b70c91..91c7e29 100644
--- a/training/configs/llama2-7b-lora.yaml
+++ b/training/configs/llama2-7b-lora.yaml
@@ -46,12 +46,22 @@ train_params:
 # Precision
 mixed_precision: bf16
 
-# Distributed training (2x RTX 5090)
-fsdp: full_shard
-fsdp_config:
-  limit_all_gathers: true
-  offload_optimizer: true
-  offload_model: false
+# Distributed training (2x RTX 5090) - DeepSpeed ZeRO-3
+plugin: deepspeed
+huggingface_hub:
+  token: null
+deeepspeed_config:
+  zero_optimization:
+    stage: 3
+    offload_optimizer:
+      device: cpu
+      pin_memory: true
+    offload_param:
+      device: cpu
+      pin_memory: true
+  gradient_clipping: 1.0
+  train_batch_size: auto
+  train_micro_batch_size_per_gpu: auto
 
 # Evaluation
 eval_strategy: steps
diff --git a/training/scripts/train.py b/training/scripts/train.py
index 97e4b47..64183fc 100755
--- a/training/scripts/train.py
+++ b/training/scripts/train.py
@@ -34,16 +34,12 @@ def train(config_path):
     
     print(f"Loading model: {config['base_model']}")
     
-    # Load model with distributed training support
-    # Use FSDP for multi-GPU training
-    from accelerate import Accelerator
-    accelerator = Accelerator()
-    
-    # Load model on CPU first, then distribute
+    # Load model - let the model's own quantization config handle it
+    # (Ornith uses CompressedTensors, not BitsAndBytes)
     model = AutoModelForCausalLM.from_pretrained(
         config["base_model"],
         torch_dtype=torch.bfloat16,
-        device_map="cpu",  # Load on CPU first
+        device_map="auto",  # Let transformers distribute across GPUs
     )
     
     # Add LoRA
@@ -93,13 +89,9 @@ def train(config_path):
         gradient_checkpointing=config.get("gradient_checkpointing", True),
     )
     
-    # Use FSDP for multi-GPU training
+    # SFT Trainer (DeepSpeed handles distributed training via config)
     from trl import SFTTrainer
     
-    # Prepare model for FSDP
-    model = accelerator.prepare(model)
-    
-    # SFT Trainer
     trainer = SFTTrainer(
         model=model,
         tokenizer=tokenizer,
@@ -109,9 +101,6 @@ def train(config_path):
         max_seq_length=config["train_params"]["max_seq_length"],
     )
     
-    # Prepare trainer for distributed training
-    trainer = accelerator.prepare(trainer)
-    
     # Train
     print("Starting training...")
     trainer.train()