From f4b22f89bcfa08d8970f49b1bccf6b6593988c27 Mon Sep 17 00:00:00 2001 From: cloudhan Date: Sat, 15 Jun 2024 01:23:28 +0800 Subject: [PATCH] refactor flash attn test (#21028) Code improvement and allow change to use other ep --- .../transformers/test_flash_attn_cuda.py | 270 ++++-------------- 1 file changed, 62 insertions(+), 208 deletions(-) diff --git a/onnxruntime/test/python/transformers/test_flash_attn_cuda.py b/onnxruntime/test/python/transformers/test_flash_attn_cuda.py index 7dbbb9249c..1d49583b3f 100644 --- a/onnxruntime/test/python/transformers/test_flash_attn_cuda.py +++ b/onnxruntime/test/python/transformers/test_flash_attn_cuda.py @@ -24,10 +24,6 @@ from parameterized import parameterized from onnxruntime import InferenceSession, OrtValue, SessionOptions -RED = "\033[31m" -GREEN = "\033[32m" -RESET = "\033[0m" - torch.manual_seed(0) pipeline_mode = True # Reduces number of tests so pipeline doesn't time out @@ -46,6 +42,7 @@ class Config: num_heads = 0 kv_num_heads = 0 head_size = 0 + ep = "CUDAExecutionProvider" def __init__( self, batch_size, sequence_length, kv_sequence_length, past_sequence_length, num_heads, kv_num_heads, head_size @@ -59,10 +56,11 @@ class Config: self.head_size = head_size def __repr__(self): + short_ep = self.ep[: -len("ExecutionProvider")].lower() return ( f"Config(batch_size={self.batch_size}, sequence_length={self.sequence_length}, " f"kv_sequence_length={self.kv_sequence_length}, past_sequence_length={self.past_sequence_length}, " - f"num_heads={self.num_heads}, kv_num_heads={self.kv_num_heads}, head_size={self.head_size})" + f"num_heads={self.num_heads}, kv_num_heads={self.kv_num_heads}, head_size={self.head_size}, ep={short_ep})" ) @@ -74,6 +72,7 @@ class PromptConfig: num_heads = 0 kv_num_heads = 0 head_size = 0 + ep = "CUDAExecutionProvider" def __init__( self, @@ -94,10 +93,11 @@ class PromptConfig: self.head_size = head_size def __repr__(self): + short_ep = self.ep[: -len("ExecutionProvider")].lower() return ( f"PromptConfig(batch_size={self.batch_size}, q_sequence_length={self.q_sequence_length}, " f"kv_sequence_length={self.kv_sequence_length}, buffer_sequence_length={self.buffer_sequence_length}, " - f"num_heads={self.num_heads}, kv_num_heads={self.kv_num_heads}, head_size={self.head_size})" + f"num_heads={self.num_heads}, kv_num_heads={self.kv_num_heads}, head_size={self.head_size}, ep={short_ep})" ) @@ -738,7 +738,7 @@ def flash_attn_varlen_qkvpacked_func(qkv_unpad, cu_seqlens, token_offset, config "cumulative_sequence_length": cu_seqlens.cpu().numpy(), } sess_options = SessionOptions() - ort_session = InferenceSession(onnx_model_str, sess_options, providers=["CUDAExecutionProvider"]) + ort_session = InferenceSession(onnx_model_str, sess_options, providers=[config.ep]) ort_output = ort_session.run(None, ort_inputs) output = torch.tensor(ort_output) return output @@ -755,7 +755,7 @@ def mha_func(q, k, v, config): "value": v.detach().cpu().numpy(), } sess_options = SessionOptions() - ort_session = InferenceSession(onnx_model_str, sess_options, providers=["CUDAExecutionProvider"]) + ort_session = InferenceSession(onnx_model_str, sess_options, providers=[config.ep]) ort_output = ort_session.run(None, ort_inputs) ort_output = numpy.array(ort_output) output = torch.tensor(ort_output) @@ -807,7 +807,7 @@ def gqa_prompt_func( "total_sequence_length": torch.tensor([config.q_sequence_length], dtype=torch.int32).detach().cpu().numpy(), } sess_options = SessionOptions() - ort_session = InferenceSession(onnx_model_str, sess_options, providers=["CUDAExecutionProvider"]) + ort_session = InferenceSession(onnx_model_str, sess_options, providers=[config.ep]) io_binding = ort_session.io_binding() if new_k is not None: ort_inputs["key"] = new_k.detach().cpu().numpy() @@ -848,7 +848,7 @@ def gqa_prompt_func( "total_sequence_length": torch.tensor([config.q_sequence_length], dtype=torch.int32).detach().cpu().numpy(), } sess_options = SessionOptions() - ort_session = InferenceSession(onnx_model_str, sess_options, providers=["CUDAExecutionProvider"]) + ort_session = InferenceSession(onnx_model_str, sess_options, providers=[config.ep]) io_binding = ort_session.io_binding() if new_k is not None: ort_inputs["key"] = new_k.detach().cpu().numpy() @@ -915,7 +915,7 @@ def gqa_past_func( .numpy(), } sess_options = SessionOptions() - ort_session = InferenceSession(onnx_model_str, sess_options, providers=["CUDAExecutionProvider"]) + ort_session = InferenceSession(onnx_model_str, sess_options, providers=[config.ep]) io_binding = ort_session.io_binding() if new_k is not None: ort_inputs["key"] = new_k.detach().cpu().numpy() @@ -963,7 +963,7 @@ def gqa_past_func( .numpy(), } sess_options = SessionOptions() - ort_session = InferenceSession(onnx_model_str, sess_options, providers=["CUDAExecutionProvider"]) + ort_session = InferenceSession(onnx_model_str, sess_options, providers=[config.ep]) io_binding = ort_session.io_binding() if new_k is not None: ort_inputs["key"] = new_k.detach().cpu().numpy() @@ -1171,25 +1171,9 @@ def parity_check_mha( out_ref, _ = attention_ref(q, k, v, None, None, 0.0, None, causal=False) out_ref = out_ref.detach().cpu().numpy() - # Compare results - all_close = numpy.allclose(out, out_ref, rtol=rtol, atol=atol, equal_nan=True) - correct = GREEN + "True" + RESET if all_close else RED + "False" + RESET - print( - " B:", - config.batch_size, - " S:", - config.sequence_length, - " N:", - config.num_heads, - " kvN:", - config.kv_num_heads, - " h:", - config.head_size, - " Mean Error:", - numpy.mean(numpy.abs(out - out_ref)), - correct, + numpy.testing.assert_allclose( + out, out_ref, rtol=rtol, atol=atol, equal_nan=True, err_msg=f" with {config} packed={packed}" ) - return all_close def rotary_embedding(*args, **kwargs): @@ -1366,60 +1350,19 @@ def parity_check_gqa_prompt( out = torch.reshape(out, (config.batch_size, config.q_sequence_length, config.num_heads, config.head_size)) out = out.detach().cpu().numpy() + err_msg = ( + f" with {config}, causal={causal}, local={local}, past_format={past_format}," + f" rotary={rotary}, rotary_interleaved={rotary_interleaved}, packed={packed}" + ) # Make sure past-present buffer updating correctly - assert numpy.allclose(present_k, k_cache_ref.detach().cpu().numpy(), rtol=rtol, atol=atol, equal_nan=True) - assert numpy.allclose(present_v, v_cache_ref.detach().cpu().numpy(), rtol=rtol, atol=atol, equal_nan=True) - - # Compare results - all_close = numpy.allclose(out, out_ref, rtol=rtol, atol=atol, equal_nan=True) - correct = GREEN + "True" + RESET if all_close else RED + "False" + RESET - print( - "KV-buffer", - " packed:", - packed, - " causal:", - causal, - " local:", - local, - " rotary:", - rotary, - " rotary_interleaved:", - rotary_interleaved, - "past kv format:", - "BSNH" if past_format == Formats.BSNH else "BNSH", - " B:", - config.batch_size, - " S:", - config.q_sequence_length, - " kv S:", - config.kv_sequence_length, - " N:", - config.num_heads, - " kv N:", - config.kv_num_heads, - " h:", - config.head_size, - " Mean Error:", - numpy.mean(numpy.abs(out - out_ref)), - correct, + numpy.testing.assert_allclose( + present_k, k_cache_ref.detach().cpu().numpy(), rtol=rtol, atol=atol, equal_nan=True, err_msg=err_msg + ) + numpy.testing.assert_allclose( + present_v, v_cache_ref.detach().cpu().numpy(), rtol=rtol, atol=atol, equal_nan=True, err_msg=err_msg ) - if not all_close: - close_mask = numpy.isclose(out, out_ref, rtol=rtol, atol=atol, equal_nan=True) - not_close_mask = numpy.logical_not(close_mask) - - # Values that are not close - out_not_close = out[not_close_mask] - out_ref_not_close = out_ref[not_close_mask] - - # Indices that are not close - not_close_indices = numpy.argwhere(not_close_mask) - - print("Values in 'out' that are not close:", out_not_close) - print("Corresponding values in 'out_ref':", out_ref_not_close) - print("Indices of values that are not close:", not_close_indices) - - return all_close + numpy.testing.assert_allclose(out, out_ref, rtol=rtol, atol=atol, equal_nan=True, err_msg=err_msg) def parity_check_gqa_prompt_no_buff( @@ -1566,44 +1509,19 @@ def parity_check_gqa_prompt_no_buff( out = torch.reshape(out, (config.batch_size, config.q_sequence_length, config.num_heads, config.head_size)) out = out.detach().cpu().numpy() - # Make sure past-present buffer updating correctly - assert numpy.allclose(present_k, k_cache_ref.detach().cpu().numpy(), rtol=rtol, atol=atol, equal_nan=True) - assert numpy.allclose(present_v, v_cache_ref.detach().cpu().numpy(), rtol=rtol, atol=atol, equal_nan=True) - - # Compare results - all_close = numpy.allclose(out, out_ref, rtol=rtol, atol=atol, equal_nan=True) - correct = GREEN + "True" + RESET if all_close else RED + "False" + RESET - print( - "No buff", - " packed:", - packed, - " causal:", - causal, - " local:", - local, - " rotary:", - rotary, - " rotary_interleaved:", - rotary_interleaved, - "past kv format:", - "BSNH" if past_format == Formats.BSNH else "BNSH", - " B:", - config.batch_size, - " S:", - config.q_sequence_length, - " kv S:", - config.kv_sequence_length, - " N:", - config.num_heads, - " kv N:", - config.kv_num_heads, - " h:", - config.head_size, - " Mean Error:", - numpy.mean(numpy.abs(out - out_ref)), - correct, + err_msg = ( + f" with {config}, causal={causal}, local={local}, past_format={past_format}," + f" rotary={rotary}, rotary_interleaved={rotary_interleaved}, packed={packed}" ) - return all_close + # Make sure past-present buffer updating correctly + numpy.testing.assert_allclose( + present_k, k_cache_ref.detach().cpu().numpy(), rtol=rtol, atol=atol, equal_nan=True, err_msg=err_msg + ) + numpy.testing.assert_allclose( + present_v, v_cache_ref.detach().cpu().numpy(), rtol=rtol, atol=atol, equal_nan=True, err_msg=err_msg + ) + + numpy.testing.assert_allclose(out, out_ref, rtol=rtol, atol=atol, equal_nan=True, err_msg=err_msg) def parity_check_gqa_past( @@ -1769,44 +1687,19 @@ def parity_check_gqa_past( out = torch.reshape(out, (config.batch_size, config.sequence_length, config.num_heads, config.head_size)) out = out.detach().cpu().numpy() - # Make sure past-present buffer updating correctly - assert numpy.allclose(present_k, k_cache_ref.detach().cpu().numpy(), rtol=rtol, atol=atol, equal_nan=True) - assert numpy.allclose(present_v, v_cache_ref.detach().cpu().numpy(), rtol=rtol, atol=atol, equal_nan=True) - - # Compare results - all_close = numpy.allclose(out, out_ref, rtol=rtol, atol=atol, equal_nan=True) - correct = GREEN + "True" + RESET if all_close else RED + "False" + RESET - print( - "KV-buffer", - "past kv format:", - "BSNH" if past_format == Formats.BSNH else "BNSH", - " packed:", - packed, - " causal:", - causal, - " local:", - local, - " rotary:", - rotary, - " rotary_interleaved:", - rotary_interleaved, - " B:", - config.batch_size, - " S:", - config.sequence_length, - " kv S:", - config.kv_sequence_length, - " N:", - config.num_heads, - " kv N:", - config.kv_num_heads, - " h:", - config.head_size, - " Mean Error:", - numpy.mean(numpy.abs(out - out_ref)), - correct, + err_msg = ( + f" with {config}, causal={causal}, local={local}, past_format={past_format}," + f" rotary={rotary}, rotary_interleaved={rotary_interleaved}, packed={packed}" ) - return all_close + # Make sure past-present buffer updating correctly + numpy.testing.assert_allclose( + present_k, k_cache_ref.detach().cpu().numpy(), rtol=rtol, atol=atol, equal_nan=True, err_msg=err_msg + ) + numpy.testing.assert_allclose( + present_v, v_cache_ref.detach().cpu().numpy(), rtol=rtol, atol=atol, equal_nan=True, err_msg=err_msg + ) + + numpy.testing.assert_allclose(out, out_ref, rtol=rtol, atol=atol, equal_nan=True, err_msg=err_msg) def parity_check_gqa_past_no_buff( @@ -1978,40 +1871,11 @@ def parity_check_gqa_past_no_buff( out = torch.reshape(out, (config.batch_size, config.sequence_length, config.num_heads, config.head_size)) out = out.detach().cpu().numpy() - # Compare results - all_close = numpy.allclose(out, out_ref, rtol=rtol, atol=atol, equal_nan=True) - correct = GREEN + "True" + RESET if all_close else RED + "False" + RESET - print( - "NO buff", - " packed:", - packed, - " causal:", - causal, - " local:", - local, - " rotary:", - rotary, - " rotary_interleaved:", - rotary_interleaved, - "past kv format:", - "BSNH" if past_format == Formats.BSNH else "BNSH", - " B:", - config.batch_size, - " S:", - config.sequence_length, - " kv S:", - config.kv_sequence_length, - " N:", - config.num_heads, - " kv N:", - config.kv_num_heads, - " h:", - config.head_size, - " Mean Error:", - numpy.mean(numpy.abs(out - out_ref)), - correct, + err_msg = ( + f" with {config}, causal={causal}, local={local}, past_format={past_format}," + f" rotary={rotary}, rotary_interleaved={rotary_interleaved}, packed={packed}" ) - return all_close + numpy.testing.assert_allclose(out, out_ref, rtol=rtol, atol=atol, equal_nan=True, err_msg=err_msg) def packed_mha_test_cases(): @@ -2064,8 +1928,7 @@ class TestMHA(unittest.TestCase): if major < 8: return print("-------- TEST PACKED MHA ---------") - all_close = parity_check_mha(config, True) - self.assertTrue(all_close) + parity_check_mha(config, True) @parameterized.expand(mha_test_cases()) def test_mha(self, _, config): @@ -2075,8 +1938,7 @@ class TestMHA(unittest.TestCase): if major < 8: return print("-------- TEST MHA ---------") - all_close = parity_check_mha(config, False) - self.assertTrue(all_close) + parity_check_mha(config, False) def gqa_no_past_memory_efficient_test_cases(): @@ -2252,7 +2114,7 @@ class TestGQA(unittest.TestCase): os.environ["ORT_DISABLE_FLASH_ATTENTION"] = "1" print("------- MEMORY EFFICIENT ATTENTION (PROMPT CASE) ---------") - all_close = parity_check_gqa_prompt( + parity_check_gqa_prompt( config, rtol=5e-3, atol=5e-3, @@ -2261,8 +2123,7 @@ class TestGQA(unittest.TestCase): rotary_interleaved=rotary_interleaved, packed=packed, ) - self.assertTrue(all_close) - all_close = parity_check_gqa_prompt_no_buff( + parity_check_gqa_prompt_no_buff( config, rtol=5e-3, atol=5e-3, @@ -2271,7 +2132,6 @@ class TestGQA(unittest.TestCase): rotary_interleaved=rotary_interleaved, packed=packed, ) - self.assertTrue(all_close) @parameterized.expand(gqa_no_past_flash_attention_test_cases()) def test_gqa_no_past_flash_attention(self, _, config, local, rotary, rotary_interleaved, packed): @@ -2283,7 +2143,7 @@ class TestGQA(unittest.TestCase): print("------- FLASH ATTENTION (PROMPT CASE) --------") os.environ["ORT_DISABLE_FLASH_ATTENTION"] = "0" - all_close = parity_check_gqa_prompt( + parity_check_gqa_prompt( config, local=local, past_format=Formats.BNSH, @@ -2291,8 +2151,7 @@ class TestGQA(unittest.TestCase): rotary_interleaved=rotary_interleaved, packed=packed, ) - self.assertTrue(all_close) - all_close = parity_check_gqa_prompt_no_buff( + parity_check_gqa_prompt_no_buff( config, local=local, past_format=Formats.BNSH, @@ -2300,7 +2159,6 @@ class TestGQA(unittest.TestCase): rotary_interleaved=rotary_interleaved, packed=packed, ) - self.assertTrue(all_close) @parameterized.expand(gqa_past_memory_efficient_test_cases()) def test_gqa_past_memory_efficient(self, _, config, rotary, rotary_interleaved, packed): @@ -2312,7 +2170,7 @@ class TestGQA(unittest.TestCase): os.environ["ORT_DISABLE_FLASH_ATTENTION"] = "1" print("-------- MEMORY EFFICIENT (TOKEN GEN) --------") - all_close = parity_check_gqa_past( + parity_check_gqa_past( config, past_format=Formats.BNSH, rtol=1e-3, @@ -2321,8 +2179,7 @@ class TestGQA(unittest.TestCase): rotary_interleaved=rotary_interleaved, packed=packed, ) - self.assertTrue(all_close) - all_close = parity_check_gqa_past_no_buff( + parity_check_gqa_past_no_buff( config, past_format=Formats.BNSH, rtol=1e-3, @@ -2331,7 +2188,6 @@ class TestGQA(unittest.TestCase): rotary_interleaved=rotary_interleaved, packed=packed, ) - self.assertTrue(all_close) @parameterized.expand(gqa_past_flash_attention_test_cases()) def test_gqa_past_flash_attention(self, _, config, local, rotary, rotary_interleaved, packed): @@ -2343,7 +2199,7 @@ class TestGQA(unittest.TestCase): print("------- FLASH ATTENTION (TOKEN GEN) -------") os.environ["ORT_DISABLE_FLASH_ATTENTION"] = "0" - all_close = parity_check_gqa_past( + parity_check_gqa_past( config, local=local, past_format=Formats.BNSH, @@ -2353,8 +2209,7 @@ class TestGQA(unittest.TestCase): rotary_interleaved=rotary_interleaved, packed=packed, ) - self.assertTrue(all_close) - all_close = parity_check_gqa_past_no_buff( + parity_check_gqa_past_no_buff( config, local=local, past_format=Formats.BNSH, @@ -2364,7 +2219,6 @@ class TestGQA(unittest.TestCase): rotary_interleaved=rotary_interleaved, packed=packed, ) - self.assertTrue(all_close) if __name__ == "__main__":