From 8985c43033d8ab778f9c70ebb5f00d1d57e67c89 Mon Sep 17 00:00:00 2001 From: A-Cepheus <60658915+A-Cepheus@users.noreply.github.com> Date: Mon, 22 Jan 2024 15:21:14 +0800 Subject: [PATCH 1/4] fix: ZeRO3 does not work with MoE models Former-commit-id: e1d5c9851922522f45314c3058d4658198631875 --- src/llmtuner/model/loader.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/llmtuner/model/loader.py b/src/llmtuner/model/loader.py index b02a4560..656bfa6d 100644 --- 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) From 39d9aba166c1ffcd420e5321e681769de02be93d 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 2/4] =?UTF-8?q?=F0=9F=90=9E=20fix:=20typo,=20move=20MoE=20?= =?UTF-8?q?fix=20to=20patcher?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Former-commit-id: 319a72b48d79c5887f39fc09c66761d5d230a317 --- 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: From f00ad6b4f8c9efc7ea553417838fe8ecb046bbc0 Mon Sep 17 00:00:00 2001 From: A-Cepheus <60658915+A-Cepheus@users.noreply.github.com> Date: Mon, 22 Jan 2024 16:04:39 +0800 Subject: [PATCH 3/4] =?UTF-8?q?=F0=9F=90=9E=20fix:=20typo?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Former-commit-id: b06a31e76ac623ea89ecc567dd9e9ec88dfdb585 --- src/llmtuner/model/patcher.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/llmtuner/model/patcher.py b/src/llmtuner/model/patcher.py index 52690e68..892a9165 100644 --- a/src/llmtuner/model/patcher.py +++ b/src/llmtuner/model/patcher.py @@ -284,7 +284,7 @@ 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(): + if getattr(model.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]) From 68977c8ca4d8b8a0d85e2398614028e7dcb04150 Mon Sep 17 00:00:00 2001 From: hoshi-hiyouga Date: Mon, 22 Jan 2024 23:27:39 +0800 Subject: [PATCH 4/4] Update patcher.py Former-commit-id: bdc9eff635cbea670722e672fd329ac35a01a812 --- src/llmtuner/model/patcher.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/llmtuner/model/patcher.py b/src/llmtuner/model/patcher.py index 892a9165..c537a6d2 100644 --- a/src/llmtuner/model/patcher.py +++ b/src/llmtuner/model/patcher.py @@ -285,6 +285,7 @@ def patch_model( _prepare_model_for_training(model, model_args) if getattr(model.config, "model_type", None) == "mixtral" and is_deepspeed_zero3_enabled(): + require_version("deepspeed>=0.13.0", "To fix: pip install deepspeed>=0.13.0") from deepspeed.utils import set_z3_leaf_modules from transformers.models.mixtral.modeling_mixtral import MixtralSparseMoeBlock set_z3_leaf_modules(model, [MixtralSparseMoeBlock])