Profiling in PyTorch (Part 3): Attention is all you profile

查看原文 / View Original 编辑 / Edit

Profiling in PyTorch (Part 3): Attention is all you profile

PyTorch 性能分析（第三部分）：Attention is all you profile

This is the third post of Profiling in PyTorch, a series where we slowly build the skill of reading profiler traces and use it to drive optimization: 这是“PyTorch 性能分析”系列的第三篇文章。在这个系列中，我们将逐步培养阅读性能分析追踪（profiler traces）的技能，并利用它来推动模型优化：

  • Profiling in PyTorch (Part 1): A Beginner’s Guide to torch.profiler PyTorch 性能分析（第一部分）：torch.profiler 初学者指南
  • Profiling in PyTorch (Part 2): From nn.Linear to a Fused MLP PyTorch 性能分析（第二部分）：从 nn.Linear 到融合 MLP
  • Profiling in PyTorch (Part 3): Attention is all you profile (current) PyTorch 性能分析（第三部分）：Attention is all you profile（当前文章）

The series “Profiling in PyTorch” is meant to make you comfortable reading profiler traces and tables. In Part 1 we profiled basic math operations like addition and multiplication. We saw how the profiler table uncovers hotspots, and how the profiler trace shows the order in which an algorithm runs over time. In Part 2 we wrapped that addition and multiplication into a torch linear layer. We then stacked several linear layers on top of each other (a multilayer perceptron) and profiled that. Along the way we also profiled fused and hand-tuned kernels. “PyTorch 性能分析”系列旨在让你能够自如地阅读性能分析追踪和表格。在第一部分中，我们分析了加法和乘法等基础数学运算。我们看到了性能分析表格如何揭示热点，以及追踪图如何展示算法随时间运行的顺序。在第二部分中，我们将这些加法和乘法封装进一个 torch 线性层中。随后，我们将多个线性层堆叠在一起（多层感知机）并进行了分析。在此过程中，我们还分析了融合算子（fused kernels）和手动调优的算子。

From the perspective of the Transformer architecture, the next logical step for us to profile is yet another fundamental algorithm, attention. While being infamous for its quadratic-time complexity, many clever tricks exist to mitigate that issue and make it fast. Our goal here is not to cover every trick in detail. Instead, we want to see how each one looks different under the profiler. 从 Transformer 架构的角度来看，我们接下来要分析的逻辑步骤是另一个基础算法：Attention（注意力机制）。虽然它因其二次时间复杂度而闻名，但存在许多巧妙的技巧来缓解这一问题并提升速度。我们的目标不是详细介绍每一个技巧，而是观察它们在性能分析器下有何不同。

Naive attention

朴素注意力机制 (Naive attention)

Attention works with Queries (q), Keys (k), and Values (v). The interaction between them can be written as a short sequence of steps: 注意力机制处理查询（q）、键（k）和值（v）。它们之间的交互可以写成简短的步骤序列：

  1. Build the attention scores scores: matmul(q, k.T) 构建注意力分数 scores：matmul(q, k.T)
  2. Scale the scores: scores * scale 缩放分数：scores * scale
  3. Apply a causal mask to the scores: scores.masked_fill(mask, “-inf”) 对分数应用因果掩码：scores.masked_fill(mask, “-inf”)
  4. Normalize the scores with softmax to get the attention weights attn: softmax(scores) 使用 softmax 归一化分数以获得注意力权重 attn：softmax(scores)
  5. Reweight the values with those weights: matmul(attn, v) 使用这些权重对值进行加权：matmul(attn, v)

So attention is really a collection of primitive operations. Some of them we already know (the matmuls), and the rest are easy to spot. Let’s write a naive attention module in PyTorch and profile it. 因此，注意力机制实际上是一系列原始操作的集合。其中一些我们已经很熟悉（矩阵乘法），其余的也很容易识别。让我们在 PyTorch 中编写一个朴素的注意力模块并对其进行分析。

class NaiveCausalAttention(nn.Module):
    def __init__(self, head_dim):
        super().__init__()
        self.scale = 1.0 / math.sqrt(head_dim)

    def forward(self, q, k, v, mask):
        scores = torch.matmul(q, k.transpose(-2, -1))
        scores = scores * self.scale
        scores = scores.masked_fill(mask, float("-inf"))
        attn = torch.softmax(scores, dim=-1)
        out = torch.matmul(attn, v)
        return out

Before opening the trace, let’s do our usual exercise and guess what we should see. Tracing the forward of this module, we expect: 在打开追踪图之前，让我们进行常规练习，猜测一下我们会看到什么。追踪该模块的前向传播时，我们预期会看到：

  • a matmul kernel (q . k.T) 一个矩阵乘法算子 (q . k.T)
  • a mul kernel (the scaling) 一个乘法算子（缩放）
  • an operation for the masking 一个掩码操作
  • a softmax kernel 一个 softmax 算子
  • a matmul kernel (atten . v) 一个矩阵乘法算子 (atten . v)

Figure 1 shows the CPU lane of the profile (the GPU lane is folded so it does not overwhelm us). Inside attn_fwd (our annotated forward call) we can see exactly the operations we guessed. The matmul is an old friend by now, and the new operations are easy to spot: 图 1 展示了性能分析的 CPU 通道（GPU 通道已折叠，以免信息过载）。在 attn_fwd（我们标注的前向调用）内部，我们可以清楚地看到我们猜测的操作。矩阵乘法已经是老朋友了，而新的操作也很容易识别：

  • mul: the scaling mul：缩放
  • masked_fill: the causal masking masked_fill：因果掩码
  • softmax: the softmax kernel softmax：softmax 算子

Naive attention with inplace causal masking

使用原地（In-place）因果掩码的朴素注意力机制

All we change is masked_fill to masked_fill_ (note the trailing underscore, PyTorch’s convention for in-place operations), and we run the same script. 我们所做的唯一改变是将 masked_fill 改为 masked_fill_（注意末尾的下划线，这是 PyTorch 中原地操作的惯例），然后运行相同的脚本。

    def forward(self, q, k, v, mask):
        # ...
        # scores = scores.masked_fill(mask, float("-inf"))
        scores = scores.masked_fill_(mask, float("-inf"))
        # ...

The in-place version (Figure 5) wraps far fewer CPU ops inside the masking step than the out-of-place version (Figure 4). This is an encouraging signal. Let’s unfold the GPU lane to confirm what happened there. 原地操作版本（图 5）在掩码步骤中包含的 CPU 操作远少于非原地版本（图 4）。这是一个令人鼓舞的信号。让我们展开 GPU 通道来确认那里发生了什么。

← 返回新闻列表 / Back to news list