diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index b2cdb51342..9dc183e0e7 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -2455,10 +2455,11 @@ This version of the operator has been available since version 1 of the 'com.micr Group Query Self/Cross Attention. - Supports different number of heads for q and kv. Only supports causal or local attention. - Supports rotary position embedding. - Supports k-v cache. - CPU EP supports fp32... CUDA EP supports fp16. + *Highly recommend using k-v cache share buffer for both CPU and CUDA. Enabled through IOBinding past and present kv. + Supports different number of heads for q and kv for CPU and CUDA. + Only supports causal and local attention. + Supports rotary position embedding for CPU and CUDA. + Supports packed input for CPU and CUDA. #### Version diff --git a/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h b/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h index b0ebc50215..9d75fcbab5 100644 --- a/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h +++ b/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h @@ -104,8 +104,8 @@ class GQAAttentionBase : public AttentionBase { int past_buffer_sequence_length, // sequence length of past state int present_buffer_sequence_length, // sequence length of present state int head_size, // head size of self-attention - const T* past_key, // past key only (if not using past state) - T* present_key, // present key only (if not using present state) + const T* past_key, // past key only + T* present_key, // present key only bool past_present_share_buffer, // whether present key and value share the same buffer bool packed_qkv, // whether Q, K, V are packed ThreadPool* tp) const { // thread pool @@ -118,11 +118,13 @@ class GQAAttentionBase : public AttentionBase { const size_t present_buff_chunk_length = static_cast(present_buffer_sequence_length) * head_size; // T x H PrepareMaskGQA(mask_data, batch_size, sequence_length, present_buffer_sequence_length, local_window_size_, seqlens_k); + if (!past_present_share_buffer) { + memset(present_key, 0, batch_size * kv_num_heads_ * present_buffer_sequence_length * head_size * sizeof(T)); + } const int loop_len = batch_size * num_heads_; const float alpha = scale_ == 0.0f ? 1.0f / sqrt(static_cast(head_size)) : scale_; - // TODO: cost might differ for gqa because of right padding and total_sequence_length being sequence dependent TensorOpCost unit_cost; const size_t probs_matrix_bytes = SafeInt(sequence_length) * present_buffer_sequence_length * sizeof(T); unit_cost.compute_cycles = static_cast(2 * sequence_length * head_size * present_buffer_sequence_length); @@ -150,7 +152,6 @@ class GQAAttentionBase : public AttentionBase { T* output = attention_probs + output_offset; // Broadcast mask data: (Bx)SxT -> (BxNx)SxT - // TODO: mask after present_sequence_length memcpy(output, mask_data + mask_offset, probs_matrix_bytes); @@ -202,20 +203,23 @@ class GQAAttentionBase : public AttentionBase { int present_buffer_sequence_length, // sequence length in past state int head_size, // head size of Q, K, V int hidden_size, // hidden size of Output - const T* past_value, // past value only (if not using past state) - T* present_value, // present value only (if not using present state) + const T* past_value, // past value only + T* present_value, // present value only bool past_present_share_buffer, // whether present key and value share the same buffer bool packed_qkv, // whether Q, K, V are packed ThreadPool* tp) const { const bool is_prompt = sequence_length != 1; const int packed_batch_stride = packed_qkv ? (num_heads_ + 2 * kv_num_heads_) * sequence_length * head_size : 0; const int kv_num_heads_factor = num_heads_ / kv_num_heads_; - // TODO: what is with these being ptrdiff_t? - const ptrdiff_t q_input_chunk_length = SafeInt(sequence_length) * head_size; // S x H - const ptrdiff_t kv_input_chunk_length = SafeInt(sequence_length) * head_size; // L x H + const size_t q_input_chunk_length = static_cast(sequence_length) * head_size; // S x H + const size_t kv_input_chunk_length = static_cast(sequence_length) * head_size; // L x H const size_t past_buff_chunk_length = static_cast(past_buffer_sequence_length) * head_size; // L x H const size_t present_buff_chunk_length = static_cast(present_buffer_sequence_length) * head_size; // T x H + if (!past_present_share_buffer) { + memset(present_value, 0, batch_size * kv_num_heads_ * present_buffer_sequence_length * head_size * sizeof(T)); + } + // The cost of Gemm TensorOpCost unit_cost; unit_cost.compute_cycles = static_cast(2 * sequence_length * head_size * present_buffer_sequence_length); @@ -252,15 +256,15 @@ class GQAAttentionBase : public AttentionBase { i / kv_num_heads_factor); } - T* current_tmp_data = reinterpret_cast(tmp_buffer) + q_input_chunk_length * i; - ptrdiff_t attention_probs_offset = SafeInt(sequence_length) * present_buffer_sequence_length * i; + T* current_tmp_data = reinterpret_cast(tmp_buffer) + q_input_chunk_length * static_cast(i); + const int attention_probs_offset = sequence_length * present_buffer_sequence_length * static_cast(i); math::MatMul(sequence_length, head_size, present_buffer_sequence_length, attention_probs + attention_probs_offset, v, current_tmp_data, nullptr); // Transpose: out(B, S, N, H_v) -> out_tmp(B, N, S, H_v) T* src = current_tmp_data; - ptrdiff_t dest_offset = (SafeInt(batch_index) * sequence_length * num_heads_ + head_index) * head_size; + const int dest_offset = (batch_index * sequence_length * num_heads_ + head_index) * head_size; T* dest = output + dest_offset; for (int j = 0; j < sequence_length; j++) { memcpy(dest, src, bytes_to_copy_trans); diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc index 112f609d46..1cd0c69233 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc @@ -163,6 +163,10 @@ Status GroupQueryAttention::ComputeInternal(OpKernelContext* context) const { parameters.sequence_length <= parameters.seqlen_past_kv_cache + parameters.sequence_length && (sizeof(T) == 2 || parameters.sequence_length >= attention::kMinSeqLenForMemoryEfficientAttentionFp32) && has_memory_efficient_attention(sm, sizeof(T) == 2); + if (!use_flash_attention && !use_memory_efficient_attention && local_window_size_ != -1) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Local attention UNSUPPORTED for sm < 80 on CUDA."); + } // allocate buffers size_t kv_buffer_bytes = 0; // need a buffer if we must ungroup kv diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index 422ae068f4..6cac8f9a53 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -1046,10 +1046,11 @@ ONNX_MS_OPERATOR_SET_SCHEMA( constexpr const char* GroupQueryAttention_ver1_doc = R"DOC( Group Query Self/Cross Attention. -Supports different number of heads for q and kv. Only supports causal or local attention. -Supports rotary position embedding. -Supports k-v cache. -CPU EP supports fp32... CUDA EP supports fp16. +*Highly recommend using k-v cache share buffer for both CPU and CUDA. Enabled through IOBinding past and present kv. +Supports different number of heads for q and kv for CPU and CUDA. +Only supports causal and local attention. +Supports rotary position embedding for CPU and CUDA. +Supports packed input for CPU and CUDA. )DOC"; ONNX_MS_OPERATOR_SET_SCHEMA( diff --git a/onnxruntime/test/python/transformers/test_flash_attn_cuda.py b/onnxruntime/test/python/transformers/test_flash_attn_cuda.py index 183d621856..4926512e85 100644 --- a/onnxruntime/test/python/transformers/test_flash_attn_cuda.py +++ b/onnxruntime/test/python/transformers/test_flash_attn_cuda.py @@ -24,6 +24,10 @@ from rotary_flash import apply_rotary_emb 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 @@ -1137,6 +1141,8 @@ def parity_check_mha( 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, @@ -1150,14 +1156,9 @@ def parity_check_mha( config.head_size, " Mean Error:", numpy.mean(numpy.abs(out - out_ref)), - numpy.allclose( - out, - out_ref, - rtol=rtol, - atol=atol, - equal_nan=True, - ), + correct, ) + return all_close def parity_check_gqa_prompt( @@ -1331,6 +1332,8 @@ def parity_check_gqa_prompt( 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:", @@ -1359,15 +1362,26 @@ def parity_check_gqa_prompt( config.head_size, " Mean Error:", numpy.mean(numpy.abs(out - out_ref)), - numpy.allclose( - out, - out_ref, - rtol=rtol, - atol=atol, - equal_nan=True, - ), + correct, ) + 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 + def parity_check_gqa_prompt_no_buff( config, @@ -1517,6 +1531,8 @@ def parity_check_gqa_prompt_no_buff( 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:", @@ -1545,14 +1561,9 @@ def parity_check_gqa_prompt_no_buff( config.head_size, " Mean Error:", numpy.mean(numpy.abs(out - out_ref)), - numpy.allclose( - out, - out_ref, - rtol=rtol, - atol=atol, - equal_nan=True, - ), + correct, ) + return all_close def parity_check_gqa_past( @@ -1723,6 +1734,8 @@ def parity_check_gqa_past( 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:", @@ -1751,14 +1764,9 @@ def parity_check_gqa_past( config.head_size, " Mean Error:", numpy.mean(numpy.abs(out - out_ref)), - numpy.allclose( - out, - out_ref, - rtol=rtol, - atol=atol, - equal_nan=True, - ), + correct, ) + return all_close def parity_check_gqa_past_no_buff( @@ -1931,6 +1939,8 @@ def parity_check_gqa_past_no_buff( 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:", @@ -1959,14 +1969,9 @@ def parity_check_gqa_past_no_buff( config.head_size, " Mean Error:", numpy.mean(numpy.abs(out - out_ref)), - numpy.allclose( - out, - out_ref, - rtol=rtol, - atol=atol, - equal_nan=True, - ), + correct, ) + return all_close class TestMHA(unittest.TestCase): @@ -1986,7 +1991,8 @@ class TestMHA(unittest.TestCase): for n in num_h: for h in h_sizes: config = Config(b, s, s, 0, n, n, h) - parity_check_mha(config, True) + all_close = parity_check_mha(config, True) + self.assertTrue(all_close) def test_mha(self): if not torch.cuda.is_available() or platform.system() != "Linux": @@ -2019,16 +2025,16 @@ class TestMHA(unittest.TestCase): for n in num_h: for h in h_sizes: config = Config(b, s, s2, 0, n, n, h) - parity_check_mha(config, False) + all_close = parity_check_mha(config, False) + self.assertTrue(all_close) class TestGQA(unittest.TestCase): - def test_gqa_no_past(self): + def test_gqa_no_past_memory_efficient(self): if not torch.cuda.is_available(): return major, minor = torch.cuda.get_device_capability() torch.manual_seed(69) - print("-------- TEST GQA NO PAST (PROMPT CASE) ---------") batches = [3] if pipeline_mode else [1, 3, 5] seqs = ( [ @@ -2047,7 +2053,7 @@ class TestGQA(unittest.TestCase): (240, 240), ] ) - num_h = [(32, 32), (9, 3), (4, 4)] if pipeline_mode else [(6, 6), (6, 3), (9, 9), (9, 3)] + num_h = [(32, 8), (9, 3), (4, 4)] if pipeline_mode else [(6, 6), (6, 3), (9, 9), (9, 3)] h_sizes = [16, 128, 256] if pipeline_mode else [32, 40, 64, 80, 96, 128, 160, 192, 224, 256] if major < 5 or (major == 5 and minor < 3): return @@ -2060,24 +2066,52 @@ class TestGQA(unittest.TestCase): for rotary, rotary_interleaved in [(True, False), (True, True), (False, False)]: for packed in [False, True]: config = PromptConfig(b, sq, skv, sq + skv + 8, n, n2, h) - parity_check_gqa_prompt( + all_close = parity_check_gqa_prompt( config, - rtol=2e-3, - atol=2e-3, + rtol=5e-3, + atol=5e-3, past_format=Formats.BNSH, rotary=rotary, rotary_interleaved=rotary_interleaved, packed=packed, ) - parity_check_gqa_prompt_no_buff( + self.assertTrue(all_close) + all_close = parity_check_gqa_prompt_no_buff( config, - rtol=2e-3, - atol=2e-3, + rtol=5e-3, + atol=5e-3, past_format=Formats.BNSH, rotary=rotary, rotary_interleaved=rotary_interleaved, packed=packed, ) + self.assertTrue(all_close) + + def test_gqa_no_past_flash_attention(self): + if not torch.cuda.is_available(): + return + major, _ = torch.cuda.get_device_capability() + torch.manual_seed(69) + batches = [3] if pipeline_mode else [1, 3, 5] + seqs = ( + [ + (127, 127), + (35, 35), + (2000, 2000), + (200, 200), + (240, 240), + ] + if pipeline_mode + else [ + (127, 127), + (35, 35), + (2000, 2000), + (200, 200), + (240, 240), + ] + ) + num_h = [(32, 8), (9, 3), (4, 4)] if pipeline_mode else [(6, 6), (6, 3), (9, 9), (9, 3)] + h_sizes = [16, 128, 256] if pipeline_mode else [32, 40, 64, 80, 96, 128, 160, 192, 224, 256] if major < 8 or platform.system() != "Linux": return print("------- FLASH ATTENTION (PROMPT CASE) --------") @@ -2090,7 +2124,7 @@ class TestGQA(unittest.TestCase): for rotary, rotary_interleaved in [(True, False), (True, True), (False, False)]: for packed in [False, True]: config = PromptConfig(b, sq, skv, sq + skv + 8, n, n2, h) - parity_check_gqa_prompt( + all_close = parity_check_gqa_prompt( config, local=local, past_format=Formats.BNSH, @@ -2098,7 +2132,8 @@ class TestGQA(unittest.TestCase): rotary_interleaved=rotary_interleaved, packed=packed, ) - parity_check_gqa_prompt_no_buff( + self.assertTrue(all_close) + all_close = parity_check_gqa_prompt_no_buff( config, local=local, past_format=Formats.BNSH, @@ -2106,15 +2141,15 @@ class TestGQA(unittest.TestCase): rotary_interleaved=rotary_interleaved, packed=packed, ) + self.assertTrue(all_close) - def test_gqa_past(self): + def test_gqa_past_memory_efficient(self): if not torch.cuda.is_available(): return major, minor = torch.cuda.get_device_capability() if major < 5 or (major == 5 and minor < 3): return os.environ["ORT_DISABLE_FLASH_ATTENTION"] = "1" - print("-------- TEST GQA PAST (TOKEN GEN) ---------") batches = [5] if pipeline_mode else [1, 3, 5] seqs = ( [(1, 128), (1, 1024), (1, 2048)] @@ -2133,7 +2168,7 @@ class TestGQA(unittest.TestCase): # (128, 128), ] ) - num_h = [(32, 32), (9, 3), (4, 4)] if pipeline_mode else [(6, 6), (6, 3), (9, 9), (9, 3)] + num_h = [(32, 8), (9, 3), (4, 4)] if pipeline_mode else [(6, 6), (6, 3), (9, 9), (9, 3)] h_sizes = [16, 128, 256] if pipeline_mode else [32, 40, 64, 80, 96, 128, 160, 192, 224, 256] random.seed(69) print("-------- MEMORY EFFICIENT (TOKEN GEN) --------") @@ -2145,7 +2180,7 @@ class TestGQA(unittest.TestCase): for packed in [False, True]: sp = random.randint(1, s2 - s) if s2 - s > 0 else 0 config = Config(b, s, s2, sp, n, n2, h) - parity_check_gqa_past( + all_close = parity_check_gqa_past( config, past_format=Formats.BNSH, rtol=1e-3, @@ -2154,7 +2189,8 @@ class TestGQA(unittest.TestCase): rotary_interleaved=rotary_interleaved, packed=packed, ) - parity_check_gqa_past_no_buff( + self.assertTrue(all_close) + all_close = parity_check_gqa_past_no_buff( config, past_format=Formats.BNSH, rtol=1e-3, @@ -2163,6 +2199,33 @@ class TestGQA(unittest.TestCase): rotary_interleaved=rotary_interleaved, packed=packed, ) + self.assertTrue(all_close) + + def test_gqa_past_flash_attention(self): + if not torch.cuda.is_available(): + return + major, _ = torch.cuda.get_device_capability() + batches = [5] if pipeline_mode else [1, 3, 5] + seqs = ( + [(1, 128), (1, 1024), (1, 2048)] + if pipeline_mode + else [ + (1, 128), + (1, 339), + (1, 1024), + (1, 5000), + (1, 800), + (1, 256), + (1, 799), + (1, 2048), + # (1, 128 * 512), + # (16, 128 * 512), + # (128, 128), + ] + ) + num_h = [(32, 8), (9, 3), (4, 4)] if pipeline_mode else [(6, 6), (6, 3), (9, 9), (9, 3)] + h_sizes = [16, 128, 256] if pipeline_mode else [32, 40, 64, 80, 96, 128, 160, 192, 224, 256] + random.seed(69) if major < 8 or platform.system() != "Linux": return print("------- FLASH ATTENTION (TOKEN GEN) -------") @@ -2176,7 +2239,7 @@ class TestGQA(unittest.TestCase): for packed in [False, True]: sp = random.randint(1, s2 - s) if s2 - s > 0 else 0 config = Config(b, s, s2, sp, n, n2, h) - parity_check_gqa_past( + all_close = parity_check_gqa_past( config, local=local, past_format=Formats.BNSH, @@ -2186,7 +2249,8 @@ class TestGQA(unittest.TestCase): rotary_interleaved=rotary_interleaved, packed=packed, ) - parity_check_gqa_past_no_buff( + self.assertTrue(all_close) + all_close = parity_check_gqa_past_no_buff( config, local=local, past_format=Formats.BNSH, @@ -2196,6 +2260,7 @@ class TestGQA(unittest.TestCase): rotary_interleaved=rotary_interleaved, packed=packed, ) + self.assertTrue(all_close) if __name__ == "__main__": diff --git a/onnxruntime/test/python/transformers/test_gqa_cpu.py b/onnxruntime/test/python/transformers/test_gqa_cpu.py index d2f3a52d3b..252f6b0965 100644 --- a/onnxruntime/test/python/transformers/test_gqa_cpu.py +++ b/onnxruntime/test/python/transformers/test_gqa_cpu.py @@ -16,14 +16,6 @@ import unittest import numpy import torch from bert_padding import pad_input, unpad_input - -try: - from colorama import Fore, init - - init(autoreset=True) -except ImportError: - print("colorama is not installed, please install it to get prettier output") - Fore = None from einops import rearrange, repeat from onnx import TensorProto, helper @@ -33,6 +25,10 @@ torch.manual_seed(0) pipeline_mode = True # Reduces number of tests so pipeline doesn't time out +RED = "\033[31m" +GREEN = "\033[32m" +RESET = "\033[0m" + class Formats: BSNH = 0 @@ -1163,10 +1159,7 @@ def parity_check_gqa_prompt( # Compare results all_close = numpy.allclose(out, out_ref, rtol=rtol, atol=atol, equal_nan=True) - if Fore is not None: - correct = Fore.GREEN + "True" if all_close else Fore.RED + "False" - else: - correct = "True" if all_close else "False" + correct = GREEN + "True" + RESET if all_close else RED + "False" + RESET print( "KV-buffer", " packed:", @@ -1197,6 +1190,7 @@ def parity_check_gqa_prompt( numpy.mean(numpy.abs(out - out_ref)), correct, ) + return all_close def parity_check_gqa_prompt_no_buff( @@ -1332,10 +1326,7 @@ def parity_check_gqa_prompt_no_buff( # Compare results all_close = numpy.allclose(out, out_ref, rtol=rtol, atol=atol, equal_nan=True) - if Fore is not None: - correct = Fore.GREEN + "True" if all_close else Fore.RED + "False" - else: - correct = "True" if all_close else "False" + correct = GREEN + "True" + RESET if all_close else RED + "False" + RESET print( "No buff", " packed:", @@ -1366,6 +1357,7 @@ def parity_check_gqa_prompt_no_buff( numpy.mean(numpy.abs(out - out_ref)), correct, ) + return all_close def parity_check_gqa_past( @@ -1532,10 +1524,7 @@ def parity_check_gqa_past( # Compare results all_close = numpy.allclose(out, out_ref, rtol=rtol, atol=atol, equal_nan=True) - if Fore is not None: - correct = Fore.GREEN + "True" if all_close else Fore.RED + "False" - else: - correct = "True" if all_close else "False" + correct = GREEN + "True" + RESET if all_close else RED + "False" + RESET print( "KV-buffer", "past kv format:", @@ -1566,6 +1555,7 @@ def parity_check_gqa_past( numpy.mean(numpy.abs(out - out_ref)), correct, ) + return all_close def parity_check_gqa_past_no_buff( @@ -1734,15 +1724,7 @@ def parity_check_gqa_past_no_buff( # Compare results all_close = numpy.allclose(out, out_ref, rtol=rtol, atol=atol, equal_nan=True) - # if not all_close: - # print("seqlens", cache_seqlens) - # print("out", out) - # print("out_ref", out_ref) - # print(out - out_ref) - if Fore is not None: - correct = Fore.GREEN + "True" if all_close else Fore.RED + "False" - else: - correct = "True" if all_close else "False" + correct = GREEN + "True" + RESET if all_close else RED + "False" + RESET print( "NO buff", " packed:", @@ -1773,6 +1755,7 @@ def parity_check_gqa_past_no_buff( numpy.mean(numpy.abs(out - out_ref)), correct, ) + return all_close class TestGQA(unittest.TestCase): @@ -1797,7 +1780,7 @@ class TestGQA(unittest.TestCase): (240, 240), ] ) - num_h = [(32, 32), (9, 3), (4, 4)] if pipeline_mode else [(6, 6), (6, 3), (9, 9), (9, 3)] + num_h = [(32, 8), (9, 3), (4, 4)] if pipeline_mode else [(6, 6), (6, 3), (9, 9), (9, 3)] h_sizes = [16, 128, 256] if pipeline_mode else [32, 40, 64, 80, 96, 128, 160, 192, 224, 256] for b in batches: for sq, skv in seqs: @@ -1808,7 +1791,7 @@ class TestGQA(unittest.TestCase): for packed in [False, True]: config = PromptConfig(b, sq, skv, sq + skv + 8, n, n2, h) past_kv_format = Formats.BNSH - parity_check_gqa_prompt( + all_close = parity_check_gqa_prompt( config, local=local, past_format=past_kv_format, @@ -1816,7 +1799,8 @@ class TestGQA(unittest.TestCase): rotary_interleaved=rotary_interleaved, packed=packed, ) - parity_check_gqa_prompt_no_buff( + self.assertTrue(all_close) + all_close = parity_check_gqa_prompt_no_buff( config, local=local, past_format=past_kv_format, @@ -1824,6 +1808,7 @@ class TestGQA(unittest.TestCase): rotary_interleaved=rotary_interleaved, packed=packed, ) + self.assertTrue(all_close) def test_gqa_past(self): print("-------- TEST GQA PAST (TOKEN GEN) ---------") @@ -1845,7 +1830,7 @@ class TestGQA(unittest.TestCase): # (128, 128), ] ) - num_h = [(16, 16), (9, 3), (4, 4)] if pipeline_mode else [(6, 6), (6, 3), (9, 9), (9, 3)] + num_h = [(32, 8), (9, 3), (4, 4)] if pipeline_mode else [(6, 6), (6, 3), (9, 9), (9, 3)] h_sizes = [16, 64, 256] if pipeline_mode else [32, 40, 64, 80, 96, 128, 160, 192, 224, 256] random.seed(69) for b in batches: @@ -1858,7 +1843,7 @@ class TestGQA(unittest.TestCase): sp = random.randint(1, s2 - s) if s2 - s > 0 else 0 config = Config(b, s, s2, sp, n, n2, h) past_kv_format = Formats.BNSH - parity_check_gqa_past( + all_close = parity_check_gqa_past( config, local=local, past_format=past_kv_format, @@ -1868,7 +1853,8 @@ class TestGQA(unittest.TestCase): rotary_interleaved=rotary_interleaved, packed=packed, ) - parity_check_gqa_past_no_buff( + self.assertTrue(all_close) + all_close = parity_check_gqa_past_no_buff( config, local=local, past_format=past_kv_format, @@ -1878,6 +1864,7 @@ class TestGQA(unittest.TestCase): rotary_interleaved=rotary_interleaved, packed=packed, ) + self.assertTrue(all_close) if __name__ == "__main__": diff --git a/onnxruntime/test/python/transformers/test_parity_neox_attention.py b/onnxruntime/test/python/transformers/test_parity_neox_attention.py index a98bb623be..d0a308987d 100644 --- a/onnxruntime/test/python/transformers/test_parity_neox_attention.py +++ b/onnxruntime/test/python/transformers/test_parity_neox_attention.py @@ -13,6 +13,7 @@ import unittest import numpy as np +import pytest import torch from onnx import TensorProto, helper from torch import nn @@ -422,6 +423,7 @@ class GPTNeoXAttention(nn.Module): class TestGPTNeoXAttention(unittest.TestCase): + @pytest.mark.skip(reason="Test broken: Error Unrecognized attribute: rotary_embedding for operator Attention") def test_gpt_neox_attention(self): for batch_size in [1, 2, 4, 8]: for seq_len in [32, 128, 512, 1024, 2048]: @@ -442,6 +444,7 @@ class TestGPTNeoXAttention(unittest.TestCase): f"Passed: test_gpt_neox_attention: {batch_size}, {seq_len}, {num_head}, {hidden_size}, {rotary_ndims}" ) + @pytest.mark.skip(reason="Test broken: Error Unrecognized attribute: rotary_embedding for operator Attention") def test_gpt_neox_decoder_masked_self_attention(self): for batch_size in [1, 2, 4, 8]: for past_seq_len in [1, 4, 32, 128, 512, 1024]: diff --git a/tools/ci_build/github/azure-pipelines/linux-gpu-ci-pipeline.yml b/tools/ci_build/github/azure-pipelines/linux-gpu-ci-pipeline.yml index c48d56b7bb..09135b465c 100644 --- a/tools/ci_build/github/azure-pipelines/linux-gpu-ci-pipeline.yml +++ b/tools/ci_build/github/azure-pipelines/linux-gpu-ci-pipeline.yml @@ -262,7 +262,7 @@ stages: cd /tmp; \ /tmp/python3 /onnxruntime_src/tools/ci_build/build.py \ --build_dir /build --config Release --test --skip_submodule_sync --build_shared_lib --parallel --use_binskim_compliant_compile_flags --build_wheel --enable_onnx_tests \ - --use_cuda --cuda_version=${{parameters.CudaVersion}} --cuda_home=/usr/local/cuda --cudnn_home=/usr/local/cuda \ + --enable_transformers_tool_test --use_cuda --cuda_version=${{parameters.CudaVersion}} --cuda_home=/usr/local/cuda --cudnn_home=/usr/local/cuda \ --enable_pybind --build_java --ctest_path "" ; \ ' displayName: 'Run Tests'