fix: ZeRO3 does not work with MoE models

2024-01-22 15:21:14 +08:00 · 2024-01-22 15:21:14 +08:00 · e1d5c98519
parent 486cc8d360
commit e1d5c98519
1 changed files with 5 additions and 0 deletions
--- a/src/llmtuner/model/loader.py
+++ b/src/llmtuner/model/loader.py
@ -96,6 +96,11 @@ def load_model_and_tokenizer(
            **config_kwargs,
        )

+    if getattr(config, "model_type", None) == "mistral" and is_deepspeed_zero3_enabled():
+        from deepspeed.utils import set_z3_leaf_modules
+        from transformers.models.mixtral.modeling_mixtral import MixtralSparseMoeBlock
+        set_z3_leaf_modules(model, [MixtralSparseMoeBlock])
+
    patch_model(model, tokenizer, model_args, is_trainable)
    register_autoclass(config, model, tokenizer)