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"],
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