p04798526
/
LLaMA-Factory-310P3
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: ZeRO3 does not work with MoE models

This commit is contained in:
A-Cepheus 2024-01-22 15:21:14 +08:00 committed by GitHub
parent 486cc8d360
commit e1d5c98519
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 5 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
src/llmtuner/model/loader.py
View File

@ -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)