feat: add CPU offload for optimizer states

2026-07-02 17:53:56 -04:00
parent a51a744d7c
commit 5a84b0eee5
2 changed files with 2 additions and 0 deletions
--- a/train.py
+++ b/train.py
@@ -99,6 +99,7 @@ def train(config_path):
        bf16=True,
        gradient_checkpointing=config.get("gradient_checkpointing", False),
        optim=config["train_params"].get("optim", "adamw_torch"),
+        optim_args=config["train_params"].get("optim_args"),
    )
    print(f"Using optimizer: {training_args.optim}")