From 5a84b0eee57e62f08f79c4f97f0bf25aef508a69 Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Thu, 2 Jul 2026 17:53:56 -0400
Subject: [PATCH] feat: add CPU offload for optimizer states

---
 train.py                              | 1 +
 training/configs/ornith-35b-lora.yaml | 1 +
 2 files changed, 2 insertions(+)

diff --git a/train.py b/train.py
index 3a96779..9215c46 100644
--- a/train.py
+++ b/train.py
@@ -99,6 +99,7 @@ def train(config_path):
         bf16=True,
         gradient_checkpointing=config.get("gradient_checkpointing", False),
         optim=config["train_params"].get("optim", "adamw_torch"),
+        optim_args=config["train_params"].get("optim_args"),
     )
     print(f"Using optimizer: {training_args.optim}")
     
diff --git a/training/configs/ornith-35b-lora.yaml b/training/configs/ornith-35b-lora.yaml
index cfec004..5088306 100644
--- a/training/configs/ornith-35b-lora.yaml
+++ b/training/configs/ornith-35b-lora.yaml
@@ -43,6 +43,7 @@ train_params:
   save_total_limit: 3
   output_dir: ../../output/ornith-35b-lora
   optim: adamw_bnb_8bit  # 8-bit optimizer to save VRAM
+  optim_args: "offload_optimizer_device=cpu"  # Offload optimizer to CPU RAM
 
 # Precision
 mixed_precision: bf16