SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

Fix building alibi tensor when num_heads is not a power of 2 (#28380)

Browse Source 
* Fix building alibi tensor when num_heads is not a power of 2

* Remove print function
This commit is contained in:
Mohamed Abu El-Nasr
2024-01-08 11:39:40 +02:00
committed by GitHub
parent 53cffeb33c
commit 0c2121f99b
2 changed files with 9 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
src/transformers/models/mpt/modeling_mpt.py
View File

@@ -70,10 +70,10 @@ def build_mpt_alibi_tensor(num_heads, sequence_length, alibi_bias_max=8, device=
base = base * (alibi_bias_max / num_heads_power_of_2) base = base * (alibi_bias_max / num_heads_power_of_2)
slopes = 1.0 / torch.pow(2, base) slopes = 1.0 / torch.pow(2, base)
slopes = slopes.view(1, num_heads, 1, 1) slopes = slopes.view(1, num_heads_power_of_2, 1, 1)
if num_heads_power_of_2 != num_heads: if num_heads_power_of_2 != num_heads:
slopes = torch.concat([slopes[1::2], slopes[::2]])[:num_heads] slopes = torch.concat([slopes[:, 1::2, ...], slopes[:, ::2, ...]], dim=1)[:, :num_heads, ...]
alibi = alibi * slopes alibi = alibi * slopes
return alibi.squeeze(0) return alibi.squeeze(0)

8
tests/models/mpt/test_modeling_mpt.py
View File

@@ -53,7 +53,7 @@ class MptModelTester:
use_labels=True, use_labels=True,
use_mc_token_ids=True, use_mc_token_ids=True,
vocab_size=99, vocab_size=99,
hidden_size=32, hidden_size=48,
num_hidden_layers=2, num_hidden_layers=2,
num_attention_heads=4, num_attention_heads=4,
intermediate_size=37, intermediate_size=37,
@@ -385,6 +385,12 @@ class MptModelTest(ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixin,
config_and_inputs = self.model_tester.prepare_config_and_inputs() config_and_inputs = self.model_tester.prepare_config_and_inputs()
self.model_tester.create_and_check_mpt_model(*config_and_inputs) self.model_tester.create_and_check_mpt_model(*config_and_inputs)
def test_mpt_model_alibi_tensor(self):
# test creation of alibi tensor when num heads is not a power of two
config_and_inputs = self.model_tester.prepare_config_and_inputs()
config_and_inputs[0].n_heads = 6
self.model_tester.create_and_check_mpt_model(*config_and_inputs)
def test_mpt_model_past(self): def test_mpt_model_past(self):
config_and_inputs = self.model_tester.prepare_config_and_inputs() config_and_inputs = self.model_tester.prepare_config_and_inputs()
self.model_tester.create_and_check_mpt_model_past(*config_and_inputs) self.model_tester.create_and_check_mpt_model_past(*config_and_inputs)