mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-29 20:14:01 +00:00
fix gqa cpu nan bug (#20521)
### Description There was a bug with gqa on cpu where on token case, with batch_size > 1, and with past_present_share_buffer off, the output would occasionally contain nans. this pr fixes that. it also updates documentation and fixes posid gen for rotary in cuda in prompt case. ### Motivation and Context this pr solves the GQA CPU bug as well as updates the documentation and makes seqlens_k irrelevant for prompt case, which is useful to prevent user error.
This commit is contained in:
aciddelgado
GitHub
parent
aff04ba08a
commit
4e27841bdb
8 changed files with 176 additions and 111 deletions
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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<size_t>(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<float>(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<size_t>(sequence_length) * present_buffer_sequence_length * sizeof(T);
|
||||
unit_cost.compute_cycles = static_cast<double>(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<ptrdiff_t>(sequence_length) * head_size; // S x H
|
||||
const ptrdiff_t kv_input_chunk_length = SafeInt<ptrdiff_t>(sequence_length) * head_size; // L x H
|
||||
const size_t q_input_chunk_length = static_cast<size_t>(sequence_length) * head_size; // S x H
|
||||
const size_t kv_input_chunk_length = static_cast<size_t>(sequence_length) * head_size; // L x H
|
||||
const size_t past_buff_chunk_length = static_cast<size_t>(past_buffer_sequence_length) * head_size; // L x H
|
||||
const size_t present_buff_chunk_length = static_cast<size_t>(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<double>(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<T*>(tmp_buffer) + q_input_chunk_length * i;
|
||||
ptrdiff_t attention_probs_offset = SafeInt<ptrdiff_t>(sequence_length) * present_buffer_sequence_length * i;
|
||||
T* current_tmp_data = reinterpret_cast<T*>(tmp_buffer) + q_input_chunk_length * static_cast<int>(i);
|
||||
const int attention_probs_offset = sequence_length * present_buffer_sequence_length * static_cast<int>(i);
|
||||
math::MatMul<T>(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<ptrdiff_t>(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);
|
||||
|
|
|
|||
|
|
@ -163,6 +163,10 @@ Status GroupQueryAttention<T>::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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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__":
|
||||
|
|
|
|||
|
|
@ -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__":
|
||||
|
|
|
|||
|
|
@ -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]:
|
||||
|
|
|
|||
|
|
@ -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'
|
||||
|
|
|
|||
Loading…
Reference in a new issue