feat: use 8-bit AdamW optimizer to save VRAM

train.py

@@ -95,11 +95,12 @@ def train(config_path):
         logging_steps=config["train_params"]["logging_steps"],
         save_steps=config["train_params"]["save_steps"],
         save_total_limit=config["train_params"]["save_total_limit"],
-        eval_strategy=config.get("eval_strategy", "steps"),
-        eval_steps=config.get("eval_steps", 100),
+        eval_strategy=config.get("eval_strategy", "no"),
         bf16=True,
-        gradient_checkpointing=config.get("gradient_checkpointing", True),
+        gradient_checkpointing=config.get("gradient_checkpointing", False),
+        optim=config["train_params"].get("optim", "adamw_torch"),
     )
+    print(f"Using optimizer: {training_args.optim}")
     
     # SFT Trainer
     from trl import SFTTrainer

training/configs/ornith-35b-lora.yaml

@@ -42,6 +42,7 @@ train_params:
   save_steps: 100
   save_total_limit: 3
   output_dir: ../../output/ornith-35b-lora
+  optim: adamw_bnb_8bit  # 8-bit optimizer to save VRAM
 
 # Precision
 mixed_precision: bf16