feat: add CPU offload for optimizer states

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}")
     

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