[optim] add support to APOLLO (#6617)

Former-commit-id: 5a252e5a458457adbd19da3b68a3897ad2962824
This commit is contained in:
zhuHQ
2025-01-14 10:24:56 -06:00
committed by GitHub
parent 66184762e8
commit c2120432db
10 changed files with 351 additions and 5 deletions

View File

@@ -147,6 +147,7 @@ class Runner:
shift_attn=get("train.shift_attn"),
report_to="all" if get("train.report_to") else "none",
use_galore=get("train.use_galore"),
use_apollo=get("train.use_apollo"),
use_badam=get("train.use_badam"),
use_swanlab=get("train.use_swanlab"),
output_dir=get_save_dir(model_name, finetuning_type, get("train.output_dir")),
@@ -223,6 +224,13 @@ class Runner:
args["galore_update_interval"] = get("train.galore_update_interval")
args["galore_scale"] = get("train.galore_scale")
args["galore_target"] = get("train.galore_target")
# apollo config
if args["use_apollo"]:
args["apollo_rank"] = get("train.apollo_rank")
args["apollo_update_interval"] = get("train.apollo_update_interval")
args["apollo_scale"] = get("train.apollo_scale")
args["apollo_target"] = get("train.apollo_target")
# badam config
if args["use_badam"]: