From a58a9e5f6c221e69125ca3fcf81a32c9743d1c52 Mon Sep 17 00:00:00 2001 From: Casper Date: Sun, 10 Dec 2023 19:17:12 +0100 Subject: [PATCH] Only fuse if flash_attn_fuse_mlp is True --- src/axolotl/utils/models.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/axolotl/utils/models.py b/src/axolotl/utils/models.py index 399bd9ff8..4f9fd8ca2 100644 --- a/src/axolotl/utils/models.py +++ b/src/axolotl/utils/models.py @@ -385,8 +385,9 @@ def load_model( **model_kwargs, ) - LOG.info("Mixtral MoE: Replacing experts with SwiGLU") - replace_mixtral_mlp_with_swiglu(model) + if cfg.flash_attn_fuse_mlp: + LOG.info("Mixtral MoE: Replacing experts with SwiGLU") + replace_mixtral_mlp_with_swiglu(model) elif model_type == "MambaLMHeadModel": # FIXME this is janky at best and hacked together to make it work