cmedina/agenx-lora-training
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: use 8-bit AdamW optimizer to save VRAM

Browse Source
This commit is contained in:
Christian Medina
2026-07-02 17:53:05 -04:00
parent 92999b7783
commit 691e5d3611
2 changed files with 5 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
train.py
View File

@@ -95,11 +95,12 @@ def train(config_path):
logging_steps=config["train_params"]["logging_steps"], logging_steps=config["train_params"]["logging_steps"],
save_steps=config["train_params"]["save_steps"], save_steps=config["train_params"]["save_steps"],
save_total_limit=config["train_params"]["save_total_limit"], save_total_limit=config["train_params"]["save_total_limit"],
eval_strategy=config.get("eval_strategy", "steps"), eval_strategy=config.get("eval_strategy", "no"),
eval_steps=config.get("eval_steps", 100),
bf16=True, 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 # SFT Trainer
from trl import SFTTrainer from trl import SFTTrainer

1
training/configs/ornith-35b-lora.yaml
View File

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