fix lookup and add docs

axolotl-ai-cloud · Dec 26, 2024 · 932fe10 · 932fe10
1 parent 32ab034
commit 932fe10
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/docs/lr_groups.qmd b/docs/lr_groups.qmd
diff --git a/src/axolotl/core/trainer_builder.py b/src/axolotl/core/trainer_builder.py
@@ -495,7 +495,12 @@ def create_optimizer_grouped_parameters(self, opt_model, optimizer_kwargs):
                 if lr_groups_lookup and any(
                     group_modules in name for group_modules in lr_groups_lookup
                 ):
-                    group_name = lr_groups_lookup[name]
+                    lr_group_module = [
+                        group_modules
+                        for group_modules in lr_groups_lookup
+                        if group_modules in name
+                    ][0]
+                    group_name = lr_groups_lookup[lr_group_module]
                     params[f"to_weight_decay_{group_name}"][name] = param
                 else:
                     params["to_weight_decay"][name] = param