探索Transformer中的多头注意力机制：如何利用GPU并发

news2024/9/22 19:40:01

什么是多头注意力机制？

首先，什么是多头注意力机制？简单来说，它是Transformer模型的核心组件之一。它通过并行计算多个注意力头（attention heads），使模型能够从不同的表示子空间中捕捉不同的特征。想象一下，你有八只眼睛，每只眼睛都能看到不同的东西，这样你就能更全面地理解世界。

为什么要用多头注意力？

你可能会问，为什么要用多头注意力？单头注意力不够吗？答案是，不够！单头注意力只能关注输入的某一部分，而多头注意力可以同时关注多个部分，从而捕捉到更多的信息。这就像你在看电影时，不仅能看到主角的表演，还能注意到背景中的细节。

多头注意力如何支持GPU并发？

好了，重点来了：多头注意力是如何支持GPU并发的？答案在于并行计算。每个注意力头的计算是独立的，因此可以分配到不同的GPU上并行处理。这不仅提高了计算效率，还能充分利用GPU的计算能力。

实例证明

让我们通过一个实例来证明这一点。以下是一个简化的多头注意力机制的实现，并展示了如何将不同的注意力头分配到不同的GPU上进行并行计算。

import torch
import torch.nn as nn

class MultiHeadAttention(nn.Module):
    def __init__(self, embed_size, heads, devices):
        super().__init__()
        self.embed_size = embed_size
        self.heads = heads
        self.head_dim = embed_size // heads
        self.devices = devices

        assert (
            self.head_dim * heads == embed_size
        ), "Embedding size needs to be divisible by heads"

        self.values = nn.ModuleList([nn.Linear(self.head_dim, self.head_dim, bias=False).to(devices[i % len(devices)]) for i in range(heads)])
        self.keys = nn.ModuleList([nn.Linear(self.head_dim, self.head_dim, bias=False).to(devices[i % len(devices)]) for i in range(heads)])
        self.queries = nn.ModuleList([nn.Linear(self.head_dim, self.head_dim, bias=False).to(devices[i % len(devices)]) for i in range(heads)])
        self.fc_out = nn.Linear(heads * self.head_dim, embed_size).to(devices[0])

    def forward(self, values, keys, query, mask):
        N = query.shape[0]
        value_len, key_len, query_len = values.shape[1], keys.shape[1], query.shape[1]

        values = values.reshape(N, value_len, self.heads, self.head_dim)
        keys = keys.reshape(N, key_len, self.heads, self.head_dim)
        queries = query.reshape(N, query_len, self.heads, self.head_dim)

        attention_heads = []
        for i in range(self.heads):
            device = self.devices[i % len(self.devices)]
            v = values[:, :, i].to(device)
            k = keys[:, :, i].to(device)
            q = queries[:, :, i].to(device)

            v = self.values[i](v)
            k = self.keys[i](k)
            q = self.queries[i](q)

            energy = torch.einsum("nqd,nkd->nqk", [q, k])
            if mask is not None:
                energy = energy.masked_fill(mask == 0, float("-1e20"))

            attention = torch.softmax(energy / (self.embed_size ** (1 / 2)), dim=2)
            out = torch.einsum("nqk,nvd->nqd", [attention, v])
            attention_heads.append(out.to(self.devices[0]))

        out = torch.cat(attention_heads, dim=2)
        out = self.fc_out(out)
        return out

# 创建模型实例
devices = [torch.device("cuda:0"), torch.device("cuda:1")]
model = MultiHeadAttention(embed_size=512, heads=8, devices=devices)

# 输入数据
values = torch.randn(64, 10, 512).to(devices[0])
keys = torch.randn(64, 10, 512).to(devices[0])
query = torch.randn(64, 10, 512).to(devices[0])
mask = None

# 前向传播
output = model(values, keys, query, mask)
print(output.shape)  # 输出 (64, 10, 512)

代码解析

初始化
- init 方法中，我们定义了多头注意力机制的各个部分，并将它们分配到不同的GPU上
前向传播
- 在 forward 方法中，我们将输入数据分割成多个注意力头，并将每个注意力头的数据分配到相应的GPU上。
- 然后，我们在每个GPU上独立计算注意力分数和加权和。
- 最后，将所有注意力头的输出拼接在一起，并通过一个线性变换生成最终的输出
运行效果