From 712ab4ae7aa3fff3085d4f439f5ecfee7a4099d3 Mon Sep 17 00:00:00 2001 From: A-Cepheus <60658915+A-Cepheus@users.noreply.github.com> Date: Mon, 22 Jan 2024 16:01:58 +0800 Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9E=20fix:=20typo,=20move=20MoE=20fix?= =?UTF-8?q?=20to=20patcher?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Former-commit-id: 4ff28e99ff9b48df7150591c6bbd3723f22b7715 --- src/llmtuner/model/loader.py | 5 ----- src/llmtuner/model/patcher.py | 5 +++++ 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/src/llmtuner/model/loader.py b/src/llmtuner/model/loader.py index 656bfa6d..b02a4560 100644 --- a/src/llmtuner/model/loader.py +++ b/src/llmtuner/model/loader.py @@ -96,11 +96,6 @@ 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) diff --git a/src/llmtuner/model/patcher.py b/src/llmtuner/model/patcher.py index 5f67f618..52690e68 100644 --- a/src/llmtuner/model/patcher.py +++ b/src/llmtuner/model/patcher.py @@ -284,6 +284,11 @@ def patch_model( if is_trainable: _prepare_model_for_training(model, model_args) + if getattr(config, "model_type", None) == "mixtral" 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]) + def patch_valuehead_model(model: "AutoModelForCausalLMWithValueHead") -> None: def tie_weights(self: "AutoModelForCausalLMWithValueHead") -> None: