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

feat: add CPU offload for optimizer states

Browse Source
This commit is contained in:
Christian Medina
2026-07-02 17:53:56 -04:00
parent a51a744d7c
commit 5a84b0eee5
2 changed files with 2 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
train.py
View File

@@ -99,6 +99,7 @@ def train(config_path):
bf16=True, bf16=True,
gradient_checkpointing=config.get("gradient_checkpointing", False), gradient_checkpointing=config.get("gradient_checkpointing", False),
optim=config["train_params"].get("optim", "adamw_torch"), optim=config["train_params"].get("optim", "adamw_torch"),
optim_args=config["train_params"].get("optim_args"),
) )
print(f"Using optimizer: {training_args.optim}") print(f"Using optimizer: {training_args.optim}")

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

@@ -43,6 +43,7 @@ train_params:
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 optim: adamw_bnb_8bit # 8-bit optimizer to save VRAM
optim_args: "offload_optimizer_device=cpu" # Offload optimizer to CPU RAM
# Precision # Precision
mixed_precision: bf16 mixed_precision: bf16