diff --git a/train.py b/train.py index c7175a4..3a96779 100644 --- a/train.py +++ b/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 diff --git a/training/configs/ornith-35b-lora.yaml b/training/configs/ornith-35b-lora.yaml index b43acf5..1a6194f 100644 --- a/training/configs/ornith-35b-lora.yaml +++ b/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