feat: add CPU offload for optimizer states

2026-07-02 17:53:56 -04:00
parent a51a744d7c
commit 5a84b0eee5
2 changed files with 2 additions and 0 deletions
--- 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