From 5713bf02857c6f7f39c417d89851ea55c87eab97 Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Tue, 30 Jun 2026 22:47:19 -0400
Subject: [PATCH] fix: add CPU offload back to DeepSpeed ZeRO-3

---
 training/configs/ornith-35b-lora.yaml | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/training/configs/ornith-35b-lora.yaml b/training/configs/ornith-35b-lora.yaml
index 06e9f13..22b8583 100644
--- a/training/configs/ornith-35b-lora.yaml
+++ b/training/configs/ornith-35b-lora.yaml
@@ -53,6 +53,12 @@ huggingface_hub:
 deepspeed_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