IamCreateAI

/

Ruyi-Mini-7B

like 458

Follow

IamCreateAI 148

I think this issue is similar to the one mentioned in this GitHub issue.

The current model requires the bfloat16 data type, which is only supported by NVIDIA's Ampere architecture (RTX 30 series) and later. Therefore, GPUs based on architectures prior to Ampere may encounter undesirable issues.

I think the main cause is related to the backend used in scaled_dot_product_attention. Could you help us identify which backend is being utilized? The following code snippets can be used to test the backend:

from torch.nn.attention import sdpa_kernel, SDPBackend

backends = []
backends.append(SDPBackend.CUDNN_ATTENTION)
# backends.append(SDPBackend.EFFICIENT_ATTENTION)
# backends.append(SDPBackend.MATH)
# backends_flash.append(SDPBackend.FLASH_ATTENTION)

...

with sdpa_kernel(backends):
    out = F.scaled_dot_product_attention(
        query, key, value,
        attn_mask=None,
        dropout_p=0.0,
        is_causal=False
    )

Thanks for getting back to me.
The confusion is that it’s not failing on bf16 but internally runs f32 according to this post:
https://discuss.pytorch.org/t/bfloat16-on-nvidia-v100-gpu/201629
This has confused me.

Sure I’ll run the test code and get back to you soon.

Thankfully I have another machine with RTX 3090s
I’ll try on that

Thanks again for the model!

I've run this:

import torch
import torch.nn.functional as F
from torch.nn.attention import sdpa_kernel, SDPBackend

# Define input tensors
query = torch.randn(2, 4, 8, device='cuda')  # (batch_size, seq_len, embed_dim)
key = torch.randn(2, 4, 8, device='cuda')    # (batch_size, seq_len, embed_dim)
value = torch.randn(2, 4, 8, device='cuda')  # (batch_size, seq_len, embed_dim)

# Select backends for testing
backends = [SDPBackend.CUDNN_ATTENTION]  # You can add other backends here

# Use the chosen backend with the scaled_dot_product_attention function
with sdpa_kernel(backends):
    out = F.scaled_dot_product_attention(
        query, key, value,
        attn_mask=None,    # Optional mask
        dropout_p=0.0,     # No dropout for testing
        is_causal=False    # Non-causal attention
    )

print("Scaled Dot-Product Attention Output:")
print(out)

STDOUT:

/home/kecso/Documents/workspace/Ruyi-Models/test1.py:15: UserWarning: Memory efficient kernel not used because: (Triggered internally at ../aten/src/ATen/native/transformers/cuda/sdp_utils.cpp:773.)
  out = F.scaled_dot_product_attention(
/home/kecso/Documents/workspace/Ruyi-Models/test1.py:15: UserWarning: Memory Efficient attention has been runtime disabled. (Triggered internally at ../aten/src/ATen/native/transformers/sdp_utils_cpp.h:558.)
  out = F.scaled_dot_product_attention(
/home/kecso/Documents/workspace/Ruyi-Models/test1.py:15: UserWarning: Flash attention kernel not used because: (Triggered internally at ../aten/src/ATen/native/transformers/cuda/sdp_utils.cpp:775.)
  out = F.scaled_dot_product_attention(
/home/kecso/Documents/workspace/Ruyi-Models/test1.py:15: UserWarning: Flash attention has been runtime disabled. (Triggered internally at ../aten/src/ATen/native/transformers/sdp_utils_cpp.h:546.)
  out = F.scaled_dot_product_attention(
/home/kecso/Documents/workspace/Ruyi-Models/test1.py:15: UserWarning: CuDNN attention kernel not used because: (Triggered internally at ../aten/src/ATen/native/transformers/cuda/sdp_utils.cpp:777.)
  out = F.scaled_dot_product_attention(
/home/kecso/Documents/workspace/Ruyi-Models/test1.py:15: UserWarning: All fused kernels requires query, key and value to be 4 dimensional, but got Query dim: 3, Key dim: 3, Value dim: 3 instead. (Triggered internally at ../aten/src/ATen/native/transformers/sdp_utils_cpp.h:317.)
  out = F.scaled_dot_product_attention(
Traceback (most recent call last):
  File "/home/kecso/Documents/workspace/Ruyi-Models/test1.py", line 15, in <module>
    out = F.scaled_dot_product_attention(
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: No available kernel. Aborting execution.

Hope this helps

1. Thank you for your testing. Based on the outputs, I believe that scaled_dot_product_attention could be the main cause. Since the other backends are not available, PyTorch might consider using SDPBackend.MATH, which does require a significant amount of GPU memory. I have listed my results below:

Profiling with shape (32, 8, 4096, 128)...

Results:
----------------------------------------------------------------------
Implementation    Avg Time (ms)    Std Dev (ms)     Peak Memory (GB)
----------------------------------------------------------------------
sdpa_cudnn               12.503           0.340                1.258
sdpa_memory              19.870           0.269                1.008
sdpa_math               154.977         181.029               54.759
sdpa_flash               11.492           0.241                1.258

As shown in the table above, sdpa_math indeed uses a substantial amount of GPU memory.

2. Regarding bfloat16 and float32, I think that having the same number of exponent bits (8 bits) may be the reason why PyTorch internally uses float32. This should be able to execute without any loss of precision, but it will require more GPU memory.