make sure to patch for z3

axolotl-ai-cloud · Jan 24, 2024 · 05be258 · 05be258
1 parent 0fc6645
commit 05be258
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/src/axolotl/utils/models.py b/src/axolotl/utils/models.py
@@ -333,7 +333,10 @@ def load_model(
         )
 
         LOG.info("patching mixtral with flash attention")
-        replace_mixtral_attn_with_multipack_flash_attn()
+        mixtral_patch_kwargs = {}
+        if is_deepspeed_zero3_enabled():
+            mixtral_patch_kwargs["for_zero3"] = True
+        replace_mixtral_attn_with_multipack_flash_attn(**mixtral_patch_kwargs)
 
     if cfg.model_config_type == "falcon" and cfg.flash_attention and cfg.sample_packing:
         from axolotl.monkeypatch.falcon import (