3D视觉Transformer中4D张量相乘时CUDA内存不足问题求助

问题描述 #

我正在开发基于Vision Transformer的3D图像模型，Q和K的张量形状为(1, 4, 2097152, 32)，计算Q与K的乘积时触发CUDA内存不足错误，试过torch.matmul和torch.bmm都没解决。使用的是Colab Pro的A100 GPU。

报错信息 #

OutOfMemoryError: CUDA out of memory. Tried to allocate 65536.00 GiB. GPU 0 has a total capacity of 15.77 GiB of which 13.47 GiB is free. Process 56962 has 2.30 GiB memory in use. Of the allocated memory 2.00 GiB is allocated by PyTorch, and 0 bytes is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)

使用的CPA3D模块代码 #

class CPA3D(nn.Module):
    def __init__(self, dim, dim_, heads=2, dim_head=64, dropout=0., p_one = True):
        super().__init__()

        self.dim = dim
        self.heads = heads
        self.scale = (dim_head ** -0.5)

        self.attend = nn.Softmax(dim=-1)
        self.to_q = nn.Linear(dim_, heads * dim_head, bias=False)
        self.to_k = nn.Linear(dim, heads * dim_head, bias=False)
        self.to_v = nn.Linear(dim, heads * dim_head, bias=False)
        self.p_one = p_one
        self.to_out = nn.Sequential(
            nn.Linear(heads * dim_head, dim),
            nn.Dropout(dropout)
        )


    def forward(self, x1, x2):
        B, D, H, W, C = x1.shape

        x1_flat = x1.view(B, -1, C)  # B, D*H*W, C
        x2_flat = x2.view(B, -1, C)  # B, D*H*W, C

        k = self.to_k(x1_flat)

        v = self.to_v(x1_flat)

        q = self.to_q(x2_flat)

        q = rearrange(q, 'b n (h d) -> b h n d', h=self.heads)
     
        k = rearrange(k, 'b n (h d) -> b h n d', h=self.heads)
        v = rearrange(v, 'b n (h d) -> b h n d', h=self.heads)

        dots = torch.bmm(q, k.permute(0,1,3,2)) * self.scale # line
        attn = self.attend(dots)
        out = torch.bmm(attn, v)
        out = rearrange(out, 'b h n d -> b n (h d)')

        if self.p_one:
          out = self.to_out(out)

        out = out.view(B, D, H, W, -1)

        return out

解决方案 #

核心问题分析 #

当前Q/K的序列长度是2097152，计算注意力矩阵时会生成(1,4,2097152,2097152)的张量，单精度下需要约65536GB内存，远超A100的16GB显存，这是内存溢出的根本原因。

具体优化方案 #

1. 替换为FlashAttention

PyTorch 2.0+自带的scaled_dot_product_attention实现了FlashAttention，通过分块计算和内存重排大幅降低显存占用，直接替换手动实现的注意力逻辑：

# 替换原代码中dots、attn、out的计算部分
out = torch.nn.functional.scaled_dot_product_attention(
    q, k, v, dropout_p=dropout if self.training else 0.0
)

你的代码已经将Q/K/V调整为(B, h, n, d)的格式，符合FlashAttention的输入要求。

2. 降低序列长度

• 下采样输入特征：用3D卷积池化层对输入的D、H、W维度进行下采样，比如将128x128x128的特征降到64x64x64，序列长度会从2097152降到262144，内存需求直接变为原来的1/64。
• 调整输入分辨率：如果原始输入尺寸过大，直接缩小输入的3D图像分辨率，从根源减少序列长度。

3. 使用局部/稀疏注意力

放弃全局注意力，只计算局部区域的注意力：

• 窗口注意力：将3D特征划分为固定大小的窗口（比如32x32x32），每个窗口内独立计算注意力，避免生成全局的超大注意力矩阵。
• 轴向注意力：分别在D、H、W三个轴上依次计算注意力，拆分计算步骤降低内存峰值。

4. 调整注意力参数

• 减少注意力头数：比如将heads从4改为2，或者降低dim_head从64改为32，虽然可能影响模型性能，但能直接降低注意力矩阵的内存占用。

5. 启用混合精度训练

在Colab中开启自动混合精度，将大部分张量转为FP16，显存占用直接减半：

from torch.cuda.amp import autocast, GradScaler
scaler = GradScaler()

# 训练循环中
with autocast():
    outputs = model(inputs)
    loss = criterion(outputs, labels)

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

6. 优化显存分配

设置环境变量减少显存碎片：

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

或者在代码中手动清理显存：

import torch
torch.cuda.empty_cache()

内容的提问来源于stack exchange，提问作者Oshan Devinda