PyTorch中如何并行实现多头自注意力（Multi Head Self Attention）？

嘿，这个问题我太有发言权了！之前我也踩过用nn.ModuleList循环做多头注意力的坑，速度慢到怀疑人生，后来才发现完全可以通过矩阵维度变换+单步线性投影实现全并行计算，既能跑满GPU性能，代码还更简洁。

一、多头自注意力（MHA）的并行实现思路 #

你用循环逐个处理注意力头的问题在于：GPU的核心优势是并行计算，而循环会把本可以同时完成的计算拆成串行步骤，浪费了硬件性能。正确的做法是把所有注意力头的Q/K/V投影合并到一个大的线性层里，再通过维度拆分实现并行计算。

举个PyTorch的实现例子，这也是Transformer标准MHA的写法：

import torch
import torch.nn as nn
import torch.nn.functional as F

class FastMultiHeadAttention(nn.Module):
    def __init__(self, embed_dim, num_heads=16):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        # 确保嵌入维度能被头数整除，否则需要调整或用padding
        assert self.head_dim * num_heads == embed_dim, "嵌入维度必须是头数的整数倍"

        # 用一个Linear层同时生成所有头的Q、K、V，输出维度是3*embed_dim
        self.qkv_proj = nn.Linear(embed_dim, 3 * embed_dim)
        # 最后把多头结果合并的输出投影层
        self.out_proj = nn.Linear(embed_dim, embed_dim)

    def forward(self, x):
        batch_size, seq_len, _ = x.size()
        
        # 1. 一次性生成所有头的Q、K、V
        qkv = self.qkv_proj(x)  # shape: (batch, seq_len, 3*embed_dim)
        # 拆分成Q、K、V三个独立张量
        q, k, v = qkv.chunk(3, dim=-1)  # 每个shape: (batch, seq_len, embed_dim)

        # 2. 拆分多头：把每个张量拆成(num_heads, head_dim)，并转置让头维度在前
        q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        k = k.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        v = v.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        # 现在Q/K/V的shape都是: (batch, num_heads, seq_len, head_dim)

        # 3. 并行计算所有头的注意力得分
        attn_scores = torch.matmul(q, k.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.head_dim, dtype=torch.float32))
        attn_probs = F.softmax(attn_scores, dim=-1)

        # 4. 并行应用注意力到V上
        attn_output = torch.matmul(attn_probs, v)  # shape: (batch, num_heads, seq_len, head_dim)

        # 5. 把多头结果拼接回原嵌入维度
        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.embed_dim)

        # 6. 最后通过输出投影层
        return self.out_proj(attn_output)

这个实现里没有任何循环，所有16个头的计算都是GPU一次性完成的，速度比循环写法快一个数量级都不止。

二、关于“堆叠Linear层并行计算多个输出” #

当然可以！其实上面的MHA就是这个思路的典型应用：把多个独立Linear层的参数合并到一个大的Linear层里，一次性计算所有输出，再拆分结果。

比如你想并行计算3个不同的Linear输出（比如不同任务的分类头），可以这么写：

class ParallelLinearLayers(nn.Module):
    def __init__(self, in_dim, out_dims):
        super().__init__()
        # out_dims是一个列表，比如[128, 64, 32]代表三个输出的维度
        self.total_out_dim = sum(out_dims)
        self.out_dims = out_dims
        # 用一个Linear层替代多个独立Linear
        self.shared_proj = nn.Linear(in_dim, self.total_out_dim)

    def forward(self, x):
        combined_output = self.shared_proj(x)
        # 拆分出多个输出
        outputs = []
        start_idx = 0
        for dim in self.out_dims:
            outputs.append(combined_output[:, start_idx:start_idx+dim])
            start_idx += dim
        return outputs

这种方式和单独定义3个nn.Linear的效果完全一致，但因为是单步矩阵运算，GPU利用率更高，速度更快，参数存储也更紧凑。

总结一下 #

核心思路就是用维度变换替代循环，把串行的头计算/多Linear计算转化为GPU擅长的大规模矩阵并行运算。PyTorch的view、transpose、chunk这些操作都是零拷贝的维度重排，几乎没有额外开销，放心用就好～

内容的提问来源于stack exchange，提问作者namespace-Pt