fix z3 leaf configuration when not using lists (#1817) [skip ci]

2024-08-09 10:54:52 -04:00
parent 70978467a0
commit 5ee4b7325f
1 changed files with 3 additions and 1 deletions
--- a/src/axolotl/utils/models.py
+++ b/src/axolotl/utils/models.py
@@ -817,11 +817,13 @@ def load_model(
        )

        if cfg.model_config_type in MOE_ARCH_BLOCK:
+            moe_blocks = MOE_ARCH_BLOCK[cfg.model_config_type]
+            moe_blocks = [moe_blocks] if isinstance(moe_blocks, str) else moe_blocks
            set_z3_leaf_modules(
                model,
                [
                    get_module_class_from_name(model, module_name)
-                    for module_name in MOE_ARCH_BLOCK[cfg.model_config_type]
+                    for module_name in moe_blocks
                ],
            )