Skip to content

Update Mixtral-8x7B Optimization #836

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
jychen21 wants to merge 23 commits into from
Closed

Update Mixtral-8x7B Optimization #836

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
23 commits
Select commit Hold shift + click to select a range
600c735
add reuse_cache support
Mar 26, 2024
7690614
make style
Mar 26, 2024
67535f0
make style
Mar 26, 2024
838b55f
remove debug code
Mar 26, 2024
458375a
add fp8 support of non-sdpa attn
Mar 26, 2024
ac3f004
add bucket_internal support of Mixtral
Mar 26, 2024
d7e6877
pick changes
Apr 1, 2024
6816af0
fit to r1.15 and also fp8 sdpa
Apr 3, 2024
de9d07f
make style
Apr 3, 2024
4fa8741
support long sequence prompt
Apr 9, 2024
25356a4
Merge branch 'huggingface:main' into update-mixtral-optimizations
Apr 9, 2024
0a3b7ac
make style
Apr 9, 2024
cb74e20
make style
Apr 9, 2024
ef96dc1
update long seq support
Apr 10, 2024
09fcfed
fix bug for w/o reuse_kvcache
Apr 10, 2024
5fb5354
make style
Apr 10, 2024
8d41302
fix accuracy of NaiveFA
Apr 11, 2024
416720d
update thresh of long seq
Apr 11, 2024
6cdedcf
update bucket size for naive fa
Apr 11, 2024
745bb4b
make style
Apr 11, 2024
3e0ca19
tune bucket size
Apr 11, 2024
9b0bad6
update naive fa
Apr 12, 2024
e9fba18
make style
Apr 12, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions examples/text-generation/quantization_config/maxabs_quant_mixtral.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,8 @@
"mode": "QUANTIZE",
"observer": "maxabs",
"scale_method": "maxabs_hw",
"whitelist": {"types": [], "names": ["gate","w1","w3","w2"]},
"blacklist": {"types": [], "names": [
"allowlist": {"types": [], "names": ["gate","w1","w3","w2"]},
"blocklist": {"types": [], "names": [
"model.layers.1.block_sparse_moe.experts.(3|4).w2",
"model.layers.[29-31].block_sparse_moe.experts.[0-7].w2"
]},
Expand Down
3 changes: 2 additions & 1 deletion optimum/habana/transformers/generation/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -584,8 +584,9 @@ def generate(
assert self.config.model_type in [
"llama",
"mistral",
"mixtral",
"falcon",
], "reuse_cache only supported by llama, mistral and falcon at the moment"
], "reuse_cache only supported by llama, mistral, falcon and mixtral at the moment"
if not generation_config.bucket_internal:
assert (
generation_config.bucket_size <= 0
Expand Down
12 changes: 6 additions & 6 deletions optimum/habana/transformers/modeling_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,10 @@
GaudiMistralDecoderLayer,
GaudiMistralForCausalLM,
GaudiMistralModel,
GaudiMixtralAttention,
GaudiMixtralDecoderLayer,
GaudiMixtralForCausalLM,
GaudiMixtralModel,
GaudiMptForCausalLM,
GaudiMptModel,
GaudiOPTForCausalLM,
Expand Down Expand Up @@ -104,10 +107,7 @@
gaudi_invert_attention_mask,
gaudi_llama_rmsnorm_forward,
gaudi_mistral_rmsnorm_forward,
gaudi_mixtral_attention_forward,
gaudi_mixtral_block_sparse_moe_forward,
gaudi_mixtral_decoder_layer_forward,
gaudi_mixtral_model_forward,
gaudi_mixtral_rmsnorm_forward,
gaudi_mpt_attention_forward,
gaudi_mpt_block_forward,
Expand Down Expand Up @@ -353,10 +353,10 @@ def adapt_transformers_to_gaudi():

# Optimization for mixtral on Gaudi
transformers.models.mixtral.modeling_mixtral.MixtralForCausalLM = GaudiMixtralForCausalLM
transformers.models.mixtral.modeling_mixtral.MixtralModel.forward = gaudi_mixtral_model_forward
transformers.models.mixtral.modeling_mixtral.MixtralAttention.forward = gaudi_mixtral_attention_forward
transformers.models.mixtral.modeling_mixtral.MixtralModel = GaudiMixtralModel
transformers.models.mixtral.modeling_mixtral.MixtralAttention = GaudiMixtralAttention
transformers.models.mixtral.modeling_mixtral.MixtralSparseMoeBlock.forward = gaudi_mixtral_block_sparse_moe_forward
transformers.models.mixtral.modeling_mixtral.MixtralDecoderLayer.forward = gaudi_mixtral_decoder_layer_forward
transformers.models.mixtral.modeling_mixtral.MixtralDecoderLayer = GaudiMixtralDecoderLayer
transformers.models.mixtral.modeling_mixtral.MixtralRMSNorm.forward = gaudi_mixtral_rmsnorm_forward

# Optimization for speecht5 on Gaudi
Expand Down
6 changes: 3 additions & 3 deletions optimum/habana/transformers/models/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -89,11 +89,11 @@
gaudi_mistral_rmsnorm_forward,
)
from .mixtral import (
GaudiMixtralAttention,
GaudiMixtralDecoderLayer,
GaudiMixtralForCausalLM,
gaudi_mixtral_attention_forward,
GaudiMixtralModel,
gaudi_mixtral_block_sparse_moe_forward,
gaudi_mixtral_decoder_layer_forward,
gaudi_mixtral_model_forward,
gaudi_mixtral_rmsnorm_forward,
)
from .modeling_all_models import (
Expand Down
6 changes: 3 additions & 3 deletions optimum/habana/transformers/models/mixtral/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
from .modeling_mixtral import (
GaudiMixtralAttention,
GaudiMixtralDecoderLayer,
GaudiMixtralForCausalLM,
gaudi_mixtral_attention_forward,
GaudiMixtralModel,
gaudi_mixtral_block_sparse_moe_forward,
gaudi_mixtral_decoder_layer_forward,
gaudi_mixtral_model_forward,
gaudi_mixtral_rmsnorm_forward,
)
Loading