tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

ensure epsilon is cast to float

Browse Source
This commit is contained in:
Wing Lian
2024-09-18 10:42:26 -07:00
parent 05f61a0ea5
commit 84dad0bd12
1 changed files with 9 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
src/axolotl/core/trainer_builder.py
View File

@@ -506,6 +506,15 @@ class AxolotlTrainer(SchedulerMixin, Trainer):
optim_args.get("use_decoupled_weight_decay")
)
if "epsilon" in self.args.optim_shampoo_grafting_config_kwargs:
if isinstance(
self.args.optim_shampoo_grafting_config_kwargs["epsilon"], str
):
self.args.optim_shampoo_grafting_config_kwargs[
"epsilon"
] = float(
self.args.optim_shampoo_grafting_config_kwargs["epsilon"]
)
if self.args.optim_shampoo_grafting_config_type == "adam":
grafting_config = AdamGraftingConfig(
**self.args.optim_shampoo_grafting_config_kwargs