Scaled_dot_product_attention多头数≥16时内存占用异常问题求助

问题背景 #

当注意力头数≥16时，PyTorch的scaled_dot_product_attention在Tesla V100等Volta及以上架构GPU上出现内存暴增：仅将头数从8翻倍到16，内存占用从405MB飙升至6716MB；但在GTX 1070（Pascal架构）上测试时，内存占用仅小幅增长至406MB，符合预期。

复现代码 #

import torch
length = 10000
dim = 64
head_num1 = 8
head_num2 = 16
batch = 1
shapes = [[batch, head_num1, length, dim//head_num1], [batch, head_num2, length, dim//head_num2]]
for shape in shapes:
    torch.cuda.reset_peak_memory_stats()

    shape2 = [1,1, length, length]
    q = torch.rand(shape, dtype = torch.float16).cuda()
    k= torch.rand(shape, dtype = torch.float16).cuda()
    v = torch.rand(shape, dtype = torch.float16).cuda()

    attn_mask = torch.ones(shape2, dtype = torch.bool).cuda()
    x = torch.nn.functional.scaled_dot_product_attention(
                                q, k, v, attn_mask=attn_mask, dropout_p=0)

    peak_memory = torch.cuda.max_memory_allocated()      
    print(f"head number {shape[1]} case peak memory: {peak_memory / 1e6:.2f} MB")

测试环境对比 #

高内存占用环境（Tesla V100-SXM2-16GB） #

Python Version: 3.9.19
PyTorch Version: 2.1.0+cu121
CUDA Version: 12.1
CUDDN: 8902

输出：

head number 8 case peak memory: 405.17 MB
head number 16 case peak memory: 6716.14 MB

正常内存占用环境（NVIDIA GeForce GTX 1070） #

Python Version: 3.9.19
PyTorch Version: 2.1.0+cu121
CUDA Version: 12.1
CUDDN: 8902

输出：

head number 8 case peak memory: 405.17 MB
head number 16 case peak memory: 406.49 MB

问题原因 #

差异源于PyTorch针对不同GPU架构的scaled_dot_product_attention实现策略：

• Volta及以上架构（如V100）：支持Tensor Core，PyTorch默认启用优化的Tensor Core路径。当注意力头数≥16且使用attn_mask时，会生成完整的[batch, heads, seq_len, seq_len]注意力权重矩阵，seq_len=10000时，16头的该矩阵仅float16格式就占用约3.2GB，叠加其他中间张量后导致内存暴增。
• Pascal架构（如GTX 1070）：不支持Tensor Core，PyTorch使用传统实现，不会存储完整的注意力权重矩阵，而是直接计算输出，因此内存占用稳定。

解决方案 #

• 禁用Tensor Core优化：运行代码前设置环境变量，Linux/macOS执行export TORCH_ATTENTION_DISABLE_TENSOR_CORE=1，Windows执行set TORCH_ATTENTION_DISABLE_TENSOR_CORE=1，强制PyTorch使用低内存的传统实现。
• 升级PyTorch版本：升级至2.2及以上版本，新版本对注意力机制的内存管理有优化，可能修复该问题。
• 替换为FlashAttention：使用第三方库flash-attention，它在所有GPU架构上都能高效控制内存占用且性能更优，安装后替换原注意力调用为对应实现。
• 调整参数：若业务允许，适当减少注意力头数或序列长度，避免触发高内存路径。

内容的提问来源于stack exchange，提问作者Kerry Zhu