From cd58ffd31ee25e1661a19a05e266a0feca229c73 Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Tue, 30 Jun 2026 18:23:34 -0400
Subject: [PATCH] fix: use FSDP for 2-GPU training, reduce batch size

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

diff --git a/training/configs/llama2-7b-lora.yaml b/training/configs/llama2-7b-lora.yaml
index d951891..6b70c91 100644
--- a/training/configs/llama2-7b-lora.yaml
+++ b/training/configs/llama2-7b-lora.yaml
@@ -31,8 +31,8 @@ dataset:
 # Training Parameters
 train_params:
   num_train_epochs: 3
-  per_device_train_batch_size: 4
-  gradient_accumulation_steps: 4
+  per_device_train_batch_size: 1
+  gradient_accumulation_steps: 8
   learning_rate: 2e-4
   lr_scheduler_type: cosine
   weight_decay: 0.01
@@ -46,6 +46,13 @@ 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
+
 # Evaluation
 eval_strategy: steps
 eval_steps: 100
diff --git a/training/scripts/train.py b/training/scripts/train.py
index 83424a6..97e4b47 100755
--- a/training/scripts/train.py
+++ b/training/scripts/train.py
@@ -34,12 +34,16 @@ def train(config_path):
     
     print(f"Loading model: {config['base_model']}")
     
-    # Load model - let the model's own quantization config handle it
-    # (Ornith uses CompressedTensors, not BitsAndBytes)
+    # 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
     model = AutoModelForCausalLM.from_pretrained(
         config["base_model"],
-        device_map="auto",
         torch_dtype=torch.bfloat16,
+        device_map="cpu",  # Load on CPU first
     )
     
     # Add LoRA
@@ -89,6 +93,12 @@ def train(config_path):
         gradient_checkpointing=config.get("gradient_checkpointing", True),
     )
     
+    # Use FSDP for multi-GPU training
+    from trl import SFTTrainer
+    
+    # Prepare model for FSDP
+    model = accelerator.prepare(model)
+    
     # SFT Trainer
     trainer = SFTTrainer(
         model=model,
@@ -99,6 +109,9 @@ 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()