From 5ee4b7325f3cd09aa4d0c90b5e383fb6c5b1581e Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Fri, 9 Aug 2024 10:54:52 -0400 Subject: [PATCH] fix z3 leaf configuration when not using lists (#1817) [skip ci] --- src/axolotl/utils/models.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/axolotl/utils/models.py b/src/axolotl/utils/models.py index 1e9819c56..5ac66260a 100644 --- 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 ], )