[Whisper, Bart, MBart] Add Flash Attention 2 (#27203)
* add whisper fa2 * correct * change all * correct * correct * fix more * fix more * fix more * fix more * fix more * fix more * Apply suggestions from code review Co-authored-by: amyeroberts <22614925+amyeroberts@users.noreply.github.com> * fix more * fix more * fix more * fix more * fix more --------- Co-authored-by: amyeroberts <22614925+amyeroberts@users.noreply.github.com>
This commit is contained in:
Patrick von Platen
committed by
GitHub
GitHub
parent
3520e37e86
commit
af3de8d87c
@@ -2856,7 +2856,7 @@ class ModelTesterMixin:
|
||||
if not model_class._supports_flash_attn_2:
|
||||
return
|
||||
|
||||
config, _ = self.model_tester.prepare_config_and_inputs_for_common()
|
||||
config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
|
||||
model = model_class(config)
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdirname:
|
||||
@@ -2871,25 +2871,76 @@ class ModelTesterMixin:
|
||||
)
|
||||
model.to(torch_device)
|
||||
|
||||
dummy_input = torch.LongTensor([[1, 2, 3, 4, 5]]).to(torch_device)
|
||||
dummy_attention_mask = torch.LongTensor([[0, 1, 1, 1, 1]]).to(torch_device)
|
||||
dummy_input = inputs_dict[model.main_input_name][:1]
|
||||
if dummy_input.dtype in [torch.float32, torch.float16]:
|
||||
dummy_input = dummy_input.to(torch.bfloat16)
|
||||
|
||||
logits = model(dummy_input, output_hidden_states=True).hidden_states[-1]
|
||||
logits_fa = model_fa(dummy_input, output_hidden_states=True).hidden_states[-1]
|
||||
dummy_attention_mask = inputs_dict.get("attention_mask", None)
|
||||
|
||||
self.assertTrue(torch.allclose(logits_fa, logits, atol=4e-2, rtol=4e-2))
|
||||
if dummy_attention_mask is not None:
|
||||
dummy_attention_mask = dummy_attention_mask[:1]
|
||||
dummy_attention_mask[:, 1:] = 1
|
||||
dummy_attention_mask[:, :1] = 0
|
||||
|
||||
output_fa = model_fa(dummy_input, attention_mask=dummy_attention_mask, output_hidden_states=True)
|
||||
logits_fa = output_fa.hidden_states[-1]
|
||||
if model.config.is_encoder_decoder:
|
||||
decoder_input_ids = inputs_dict.get("decoder_input_ids", dummy_input)[:1]
|
||||
|
||||
output = model(dummy_input, attention_mask=dummy_attention_mask, output_hidden_states=True)
|
||||
logits = output.hidden_states[-1]
|
||||
outputs = model(dummy_input, decoder_input_ids=decoder_input_ids, output_hidden_states=True)
|
||||
outputs_fa = model_fa(dummy_input, decoder_input_ids=decoder_input_ids, output_hidden_states=True)
|
||||
else:
|
||||
outputs = model(dummy_input, output_hidden_states=True)
|
||||
outputs_fa = model_fa(dummy_input, output_hidden_states=True)
|
||||
|
||||
self.assertTrue(torch.allclose(logits_fa[1:], logits[1:], atol=4e-2, rtol=4e-2))
|
||||
logits = (
|
||||
outputs.hidden_states[-1]
|
||||
if not model.config.is_encoder_decoder
|
||||
else outputs.decoder_hidden_states[-1]
|
||||
)
|
||||
logits_fa = (
|
||||
outputs_fa.hidden_states[-1]
|
||||
if not model.config.is_encoder_decoder
|
||||
else outputs_fa.decoder_hidden_states[-1]
|
||||
)
|
||||
|
||||
assert torch.allclose(logits_fa, logits, atol=4e-2, rtol=4e-2)
|
||||
|
||||
if model.config.is_encoder_decoder:
|
||||
other_inputs = {
|
||||
"decoder_input_ids": decoder_input_ids,
|
||||
"decoder_attention_mask": dummy_attention_mask,
|
||||
"output_hidden_states": True,
|
||||
}
|
||||
if dummy_attention_mask is not None:
|
||||
other_inputs["attention_mask"] = dummy_attention_mask
|
||||
|
||||
outputs = model(dummy_input, **other_inputs)
|
||||
outputs_fa = model_fa(dummy_input, **other_inputs)
|
||||
else:
|
||||
other_inputs = {
|
||||
"output_hidden_states": True,
|
||||
}
|
||||
if dummy_attention_mask is not None:
|
||||
other_inputs["attention_mask"] = dummy_attention_mask
|
||||
|
||||
outputs = model(dummy_input, **other_inputs)
|
||||
outputs_fa = model_fa(dummy_input, **other_inputs)
|
||||
|
||||
logits = (
|
||||
outputs.hidden_states[-1]
|
||||
if not model.config.is_encoder_decoder
|
||||
else outputs.decoder_hidden_states[-1]
|
||||
)
|
||||
logits_fa = (
|
||||
outputs_fa.hidden_states[-1]
|
||||
if not model.config.is_encoder_decoder
|
||||
else outputs_fa.decoder_hidden_states[-1]
|
||||
)
|
||||
|
||||
assert torch.allclose(logits_fa[1:], logits[1:], atol=4e-2, rtol=4e-2)
|
||||
|
||||
# check with inference + dropout
|
||||
model.train()
|
||||
_ = model_fa(dummy_input, attention_mask=dummy_attention_mask, output_hidden_states=True)
|
||||
_ = model_fa(dummy_input, **other_inputs)
|
||||
|
||||
@require_flash_attn
|
||||
@require_torch_gpu
|
||||
@@ -2902,7 +2953,7 @@ class ModelTesterMixin:
|
||||
if not model_class._supports_flash_attn_2:
|
||||
return
|
||||
|
||||
config, _ = self.model_tester.prepare_config_and_inputs_for_common()
|
||||
config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
|
||||
model = model_class(config)
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdirname:
|
||||
@@ -2917,21 +2968,72 @@ class ModelTesterMixin:
|
||||
)
|
||||
model.to(torch_device)
|
||||
|
||||
dummy_input = torch.LongTensor([[1, 2, 3, 4, 5]]).to(torch_device)
|
||||
dummy_attention_mask = torch.LongTensor([[1, 1, 1, 1, 0]]).to(torch_device)
|
||||
dummy_input = inputs_dict[model.main_input_name][:1]
|
||||
if dummy_input.dtype in [torch.float32, torch.float16]:
|
||||
dummy_input = dummy_input.to(torch.bfloat16)
|
||||
|
||||
logits = model(dummy_input, output_hidden_states=True).hidden_states[-1]
|
||||
logits_fa = model_fa(dummy_input, output_hidden_states=True).hidden_states[-1]
|
||||
dummy_attention_mask = inputs_dict.get("attention_mask", None)
|
||||
|
||||
self.assertTrue(torch.allclose(logits_fa, logits, atol=4e-2, rtol=4e-2))
|
||||
if dummy_attention_mask is not None:
|
||||
dummy_attention_mask = dummy_attention_mask[:1]
|
||||
dummy_attention_mask[:, :-1] = 1
|
||||
dummy_attention_mask[:, -1:] = 0
|
||||
|
||||
output_fa = model_fa(dummy_input, attention_mask=dummy_attention_mask, output_hidden_states=True)
|
||||
logits_fa = output_fa.hidden_states[-1]
|
||||
if model.config.is_encoder_decoder:
|
||||
decoder_input_ids = inputs_dict.get("decoder_input_ids", dummy_input)[:1]
|
||||
|
||||
output = model(dummy_input, attention_mask=dummy_attention_mask, output_hidden_states=True)
|
||||
logits = output.hidden_states[-1]
|
||||
outputs = model(dummy_input, decoder_input_ids=decoder_input_ids, output_hidden_states=True)
|
||||
outputs_fa = model_fa(dummy_input, decoder_input_ids=decoder_input_ids, output_hidden_states=True)
|
||||
else:
|
||||
outputs = model(dummy_input, output_hidden_states=True)
|
||||
outputs_fa = model_fa(dummy_input, output_hidden_states=True)
|
||||
|
||||
self.assertTrue(torch.allclose(logits_fa[:-1], logits[:-1], atol=4e-2, rtol=4e-2))
|
||||
logits = (
|
||||
outputs.hidden_states[-1]
|
||||
if not model.config.is_encoder_decoder
|
||||
else outputs.decoder_hidden_states[-1]
|
||||
)
|
||||
logits_fa = (
|
||||
outputs_fa.hidden_states[-1]
|
||||
if not model.config.is_encoder_decoder
|
||||
else outputs_fa.decoder_hidden_states[-1]
|
||||
)
|
||||
|
||||
assert torch.allclose(logits_fa, logits, atol=4e-2, rtol=4e-2)
|
||||
|
||||
if model.config.is_encoder_decoder:
|
||||
other_inputs = {
|
||||
"decoder_input_ids": decoder_input_ids,
|
||||
"decoder_attention_mask": dummy_attention_mask,
|
||||
"output_hidden_states": True,
|
||||
}
|
||||
if dummy_attention_mask is not None:
|
||||
other_inputs["attention_mask"] = dummy_attention_mask
|
||||
|
||||
outputs = model(dummy_input, **other_inputs)
|
||||
outputs_fa = model_fa(dummy_input, **other_inputs)
|
||||
else:
|
||||
other_inputs = {
|
||||
"output_hidden_states": True,
|
||||
}
|
||||
if dummy_attention_mask is not None:
|
||||
other_inputs["attention_mask"] = dummy_attention_mask
|
||||
|
||||
outputs = model(dummy_input, **other_inputs)
|
||||
outputs_fa = model_fa(dummy_input, **other_inputs)
|
||||
|
||||
logits = (
|
||||
outputs.hidden_states[-1]
|
||||
if not model.config.is_encoder_decoder
|
||||
else outputs.decoder_hidden_states[-1]
|
||||
)
|
||||
logits_fa = (
|
||||
outputs_fa.hidden_states[-1]
|
||||
if not model.config.is_encoder_decoder
|
||||
else outputs_fa.decoder_hidden_states[-1]
|
||||
)
|
||||
|
||||
assert torch.allclose(logits_fa[:-1], logits[:-1], atol=4e-2, rtol=4e-2)
|
||||
|
||||
@require_flash_attn
|
||||
@require_torch_gpu
|
||||
@@ -2944,7 +3046,7 @@ class ModelTesterMixin:
|
||||
if not model_class._supports_flash_attn_2:
|
||||
return
|
||||
|
||||
config, _ = self.model_tester.prepare_config_and_inputs_for_common()
|
||||
config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
|
||||
model = model_class(config)
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdirname:
|
||||
@@ -2953,8 +3055,14 @@ class ModelTesterMixin:
|
||||
tmpdirname, torch_dtype=torch.float16, use_flash_attention_2=False, low_cpu_mem_usage=True
|
||||
).to(torch_device)
|
||||
|
||||
dummy_input = torch.LongTensor([[0, 2, 3, 4], [0, 2, 3, 4]]).to(torch_device)
|
||||
dummy_attention_mask = torch.LongTensor([[1, 1, 1, 1], [0, 1, 1, 1]]).to(torch_device)
|
||||
dummy_input = inputs_dict[model.main_input_name]
|
||||
if dummy_input.dtype in [torch.float32, torch.bfloat16]:
|
||||
dummy_input = dummy_input.to(torch.float16)
|
||||
|
||||
dummy_attention_mask = inputs_dict.get("attention_mask", torch.ones_like(dummy_input))
|
||||
# make sure we do left padding
|
||||
dummy_attention_mask[:, :-1] = 0
|
||||
dummy_attention_mask[:, -1:] = 1
|
||||
|
||||
out = model.generate(
|
||||
dummy_input, attention_mask=dummy_attention_mask, max_new_tokens=1, do_sample=False
|
||||
@@ -2981,7 +3089,7 @@ class ModelTesterMixin:
|
||||
if not model_class._supports_flash_attn_2:
|
||||
return
|
||||
|
||||
config, _ = self.model_tester.prepare_config_and_inputs_for_common()
|
||||
config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
|
||||
model = model_class(config)
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdirname:
|
||||
@@ -2990,8 +3098,14 @@ class ModelTesterMixin:
|
||||
tmpdirname, torch_dtype=torch.float16, use_flash_attention_2=False, low_cpu_mem_usage=True
|
||||
).to(torch_device)
|
||||
|
||||
dummy_input = torch.LongTensor([[0, 2, 3, 4], [0, 2, 3, 4]]).to(torch_device)
|
||||
dummy_attention_mask = torch.LongTensor([[1, 1, 1, 1], [1, 1, 1, 0]]).to(torch_device)
|
||||
dummy_input = inputs_dict[model.main_input_name]
|
||||
if dummy_input.dtype in [torch.float32, torch.bfloat16]:
|
||||
dummy_input = dummy_input.to(torch.float16)
|
||||
|
||||
dummy_attention_mask = inputs_dict.get("attention_mask", torch.ones_like(dummy_input))
|
||||
# make sure we do left padding
|
||||
dummy_attention_mask[:, :-1] = 1
|
||||
dummy_attention_mask[:, -1:] = 0
|
||||
|
||||
out = model.generate(
|
||||
dummy_input, attention_mask=dummy_attention_mask, max_new_tokens=1, do_sample=False
|
||||
@@ -3014,26 +3128,39 @@ class ModelTesterMixin:
|
||||
def test_flash_attn_2_generate_use_cache(self):
|
||||
import torch
|
||||
|
||||
max_new_tokens = 30
|
||||
|
||||
for model_class in self.all_generative_model_classes:
|
||||
if not model_class._supports_flash_attn_2:
|
||||
return
|
||||
|
||||
config, _ = self.model_tester.prepare_config_and_inputs_for_common()
|
||||
config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
|
||||
|
||||
dummy_input = inputs_dict[model_class.main_input_name]
|
||||
if dummy_input.dtype in [torch.float32, torch.bfloat16]:
|
||||
dummy_input = dummy_input.to(torch.float16)
|
||||
|
||||
# make sure that all models have enough positions for generation
|
||||
if hasattr(config, "max_position_embeddings"):
|
||||
config.max_position_embeddings = max_new_tokens + dummy_input.shape[1] + 1
|
||||
|
||||
model = model_class(config)
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdirname:
|
||||
model.save_pretrained(tmpdirname)
|
||||
|
||||
dummy_input = torch.LongTensor([[0, 2, 3, 4], [0, 2, 3, 4]]).to(torch_device)
|
||||
dummy_attention_mask = torch.LongTensor([[1, 1, 1, 1], [0, 1, 1, 1]]).to(torch_device)
|
||||
dummy_attention_mask = inputs_dict.get("attention_mask", torch.ones_like(dummy_input))
|
||||
|
||||
model = model_class.from_pretrained(
|
||||
tmpdirname, torch_dtype=torch.float16, use_flash_attention_2=True, low_cpu_mem_usage=True
|
||||
tmpdirname,
|
||||
torch_dtype=torch.float16,
|
||||
use_flash_attention_2=True,
|
||||
low_cpu_mem_usage=True,
|
||||
).to(torch_device)
|
||||
|
||||
# Just test that a large cache works as expected
|
||||
_ = model.generate(
|
||||
dummy_input, attention_mask=dummy_attention_mask, max_new_tokens=30, do_sample=False
|
||||
dummy_input, attention_mask=dummy_attention_mask, max_new_tokens=max_new_tokens, do_sample=False
|
||||
)
|
||||
|
||||
@require_flash_attn
|
||||
@@ -3048,14 +3175,18 @@ class ModelTesterMixin:
|
||||
if not model_class._supports_flash_attn_2:
|
||||
return
|
||||
|
||||
config, _ = self.model_tester.prepare_config_and_inputs_for_common()
|
||||
config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
|
||||
model = model_class(config)
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdirname:
|
||||
model.save_pretrained(tmpdirname)
|
||||
|
||||
dummy_input = torch.LongTensor([[0, 2, 3, 4], [0, 2, 3, 4]]).to(torch_device)
|
||||
dummy_attention_mask = torch.LongTensor([[1, 1, 1, 1], [0, 1, 1, 1]]).to(torch_device)
|
||||
dummy_input = inputs_dict[model.main_input_name]
|
||||
dummy_attention_mask = inputs_dict.get("attention_mask", torch.ones_like(dummy_input))
|
||||
|
||||
if model.config.is_encoder_decoder:
|
||||
dummy_decoder_input_ids = inputs_dict["decoder_input_ids"]
|
||||
dummy_decoder_attention_mask = inputs_dict["decoder_attention_mask"]
|
||||
|
||||
model = model_class.from_pretrained(
|
||||
tmpdirname,
|
||||
@@ -3070,10 +3201,19 @@ class ModelTesterMixin:
|
||||
if (param.dtype == torch.float16) or (param.dtype == torch.bfloat16):
|
||||
param.data = param.data.to(torch.float32)
|
||||
|
||||
_ = model(input_ids=dummy_input)
|
||||
|
||||
# with attention mask
|
||||
_ = model(input_ids=dummy_input, attention_mask=dummy_attention_mask)
|
||||
if model.config.is_encoder_decoder:
|
||||
_ = model(dummy_input, decoder_input_ids=dummy_decoder_input_ids)
|
||||
# with attention mask
|
||||
_ = model(
|
||||
dummy_input,
|
||||
attention_mask=dummy_attention_mask,
|
||||
decoder_input_ids=dummy_decoder_input_ids,
|
||||
decoder_attention_mask=dummy_decoder_attention_mask,
|
||||
)
|
||||
else:
|
||||
_ = model(dummy_input)
|
||||
# with attention mask
|
||||
_ = model(dummy_input, attention_mask=dummy_attention_mask)
|
||||
|
||||
|
||||
global_rng = random.Random()
|
||||
|
||||
Reference in New Issue
Block a user