Skip to content
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

Jump to bottom

gpt_big_code: make flash attention impl quantization friendly #1282

Merged
merged 1 commit into from
Sep 25, 2024
Merged

gpt_big_code: make flash attention impl quantization friendly #1282

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
9 changes: 5 additions & 4 deletions optimum/habana/transformers/modeling_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,7 @@
GaudiGPT2Block,
GaudiGPT2DoubleHeadsModel,
GaudiGPT2LMHeadModel,
GaudiGPTBigCodeAttention,
GaudiGPTBigCodeForCausalLM,
GaudiGPTJAttention,
GaudiGPTJBlock,
Expand Down Expand Up @@ -148,7 +149,6 @@
gaudi_generate_speech,
gaudi_get_extended_attention_mask,
gaudi_gpt2_forward,
gaudi_gpt_bigcode_attention_forward,
gaudi_gpt_bigcode_block_forward,
gaudi_gpt_bigcode_model_forward,
gaudi_gpt_neox_attention_forward,
Expand Down Expand Up @@ -356,12 +356,13 @@ def adapt_transformers_to_gaudi():
transformers.models.gptj.modeling_gptj.GPTJModel = GaudiGPTJModel

# Optimization for GPTBigCode on Gaudi
transformers.models.gpt_bigcode.modeling_gpt_bigcode.GPTBigCodeAttention.forward = (
gaudi_gpt_bigcode_attention_forward
)
transformers.models.gpt_bigcode.modeling_gpt_bigcode.GPTBigCodeAttention = GaudiGPTBigCodeAttention
transformers.models.gpt_bigcode.modeling_gpt_bigcode.GPTBigCodeForCausalLM = GaudiGPTBigCodeForCausalLM
transformers.models.gpt_bigcode.modeling_gpt_bigcode.GPTBigCodeBlock.forward = gaudi_gpt_bigcode_block_forward
transformers.models.gpt_bigcode.modeling_gpt_bigcode.GPTBigCodeModel.forward = gaudi_gpt_bigcode_model_forward
transformers.models.gpt_bigcode.modeling_gpt_bigcode.GPTBIGCODE_ATTENTION_CLASSES.update(
{"eager": GaudiGPTBigCodeAttention}
)

# Optimization for gpt-neox generation on Gaudi
transformers.models.gpt_neox.modeling_gpt_neox.GPTNeoXForCausalLM = GaudiGPTNeoXForCausalLM
Expand Down
2 changes: 1 addition & 1 deletion optimum/habana/transformers/models/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -79,8 +79,8 @@
gaudi_gpt2_forward,
)
from .gpt_bigcode import (
GaudiGPTBigCodeAttention,
GaudiGPTBigCodeForCausalLM,
gaudi_gpt_bigcode_attention_forward,
gaudi_gpt_bigcode_block_forward,
gaudi_gpt_bigcode_model_forward,
)
Expand Down
2 changes: 1 addition & 1 deletion optimum/habana/transformers/models/gpt_bigcode/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
from .modeling_gpt_bigcode import (
GaudiGPTBigCodeAttention,
GaudiGPTBigCodeForCausalLM,
gaudi_gpt_bigcode_attention_forward,
gaudi_gpt_bigcode_block_forward,
gaudi_gpt_bigcode_model_forward,
)
Loading
Loading