苹果芯片GPU加速Transformer推理：Metal性能显著提升

利用Metal Performance Shaders实现快速Transformer推理

我们很高兴地宣布，Thinc PyTorch层现已支持Metal Performance Shaders。这使得在苹果芯片Mac的GPU上运行基于spaCy Transformer的流水线成为可能，并将推理速度提升了最高4.7倍。在本文中，我们将探讨苹果芯片Mac的硬件加速设施，以及spaCy如何利用它们来加速Transformer模型。文章最后将通过基准测试展示在不同苹果芯片Mac机型上可以预期的加速效果。

众所周知，大型Transformer模型的计算成本高昂。这源于自注意力机制的二次计算复杂度，以及许多Transformer模型的庞大规模。例如，广泛使用的BERT、RoBERTa和XLM-R基础模型使用了12个隐藏层、768维的隐藏表示，以及前馈块中3072维的表示。

下图展示了在使用一个特别编译的PyTorch版本（该版本利用了某Mac M1 Max CPU的CPU核心和通用的ARM64 NEON优化内核）为德语文本标注de_dep_news_trf spaCy Transformer模型时，五个最耗时的部分：

Transformer模型的运行时间主要由矩阵乘法主导——bli_sgemm_armv8a_asm_8x12是一个单精度矩阵乘法内核。sgemm是线性代数库提供的标准化矩阵乘法函数，这些库实现了BLAS接口。这并不令人意外，因为矩阵乘法是Transformer模型使用的主要操作之一，例如：用于计算注意力块中的成对注意力分数以及前馈块中的线性投影。

除了矩阵乘法，分别用于前馈块和注意力块的GeLU和Softmax非线性函数，也通过其对erff、expf和exp初等函数的使用在性能剖析中显现出来。这些非线性函数合计占据了19%的运行时间。

为了加速Transformer推理，我们可以采取三种不同的方法：

1. 用时间复杂度优于O(N²)的机制替换自注意力机制。例如，Longformer的注意力机制具有O(N)的时间复杂度。如果我们希望继续使用现有的预训练Transformer模型，也可以对N设置一个上限。
2. 加速矩阵乘法。
3. 加速非线性计算。

在spaCy Transformer流水线中，二次注意力机制的影响已经通过限制N的上限得到了控制。每个文档被分步处理。默认情况下，spaCy Transformer每次处理96个标记，使用128个标记的窗口来创建重叠的上下文表示。这将注意力机制的时间复杂度上限设定为N=128。

在本文中，我们将重点关注另外两种方法——利用苹果芯片Mac的专用硬件加速矩阵乘法和非线性计算。我们将首先了解苹果硅CPU中称为“AMX块”的矩阵乘法协处理器。之后，我们将探索在GPU上运行计算内核。

AMX矩阵乘法块

所有Apple M系列CPU至少拥有一个称为“AMX块”的矩阵乘法协处理器。AMX在很大程度上没有公开文档。例如，目前尚不清楚Apple M系列CPU的高能效核心集群是否拥有自己的AMX块。然而，我们可以通过基准测试来推断AMX块的各种特性。

下表列出了在不同CPU上使用gemm-benchmark测量得到的768x768矩阵乘法性能（单位：TFLOPS，即每秒万亿次浮点运算）：

从这些数字中，我们可以得出一些有趣的信息：

• 性能并不随线程数增加而提升。因此，AMX块并非单个CPU核心的一部分。
• M1、M1 Pro和M1 Ultra分别拥有1、2和4个性能核心集群。矩阵乘法性能随性能核心集群数量的增加而提高（参见“与M1相比的最大加速”行）。这表明每个性能集群可能都有一个AMX块。
• AMX块速度很快。单个AMX块的矩阵乘法性能相当于9个Ryzen 5950X核心。

尽管苹果公司没有公开将计算任务分派给AMX块的指令文档，但第三方应用程序可以通过苹果的Accelerate框架使用AMX块，该框架实现了行业标准的BLAS接口。因此，BLAS矩阵乘法函数——例如我们在性能剖析中看到的sgemm函数——会自动得到加速。

由于Transformer使用矩阵乘法作为其主要操作，AMX单元为Transformer提供了显著的加速。PyTorch在苹果平台上使用Accelerate进行矩阵乘法，因此PyTorch默认使用AMX块。

Metal Performance Shaders

尽管AMX块在处理矩阵乘法吞吐量时表现出令人印象深刻的速度，但苹果芯片Mac还有另外两个用于计算的子系统，即苹果神经引擎（ANE）和GPU。ANE的限制相对较多，因为它需要运行通过Core ML定义的计算图，而GPU可以运行用户定义的计算内核（即所谓的“着色器”）。这使得GPU足够灵活，可以运行各种机器学习模型。

M1的8核GPU计算性能为2.6 TFLOPS，这大概能提供M1 AMX单元两倍的性能。此外，GPU在M1 Ultra中最多可扩展至64核，理论峰值性能达到20.8 TFLOPS。因此，苹果芯片GPU有可能将Transformer性能推至超出AMX块所能提供的水平。

PyTorch最近通过苹果的Metal API引入了对苹果M系列GPU的支持。各种PyTorch操作已被实现为自定义Metal着色器，并使用了苹果自己的Metal着色器集合，这些着色器包含在Metal Performance Shaders框架中。对于受支持的操作，在PyTorch中使用苹果芯片GPU非常简单，只需将张量或模块放在新的mps设备上即可。例如，可以按以下方式在GPU核心上进行矩阵乘法：

>>> import torch
>>> u = torch.rand((10, 20), dtype=torch.float,  device=torch.device("mps"))
>>> v = torch.rand((20, 10), dtype=torch.float,  device=torch.device("mps"))
>>> torch.matmul(u, v).device
device(type='mps', index=0)

在撰写本文时，某些操作尚未实现，但在这种情况下，当环境变量PYTORCH_ENABLE_MPS_FALLBACK设置为1时，PyTorch将回退到CPU内核。

spaCy和Thinc中的Metal Performance Shaders

spaCy使用Thinc作为其机器学习库。Thinc是一个轻量级的深度学习库，也支持在其他框架（如PyTorch和TensorFlow）中定义的层。spacy-transformers包利用Thinc的这种互操作性，使Huggingface PyTorch Transformer模型可在spaCy流水线中使用。现在，由于PyTorch支持苹果芯片GPU，原则上，基于Transformer的spaCy流水线中的Transformer模型可以在苹果芯片机器的GPU核心上执行。

不幸的是，Thinc 8.1之前的版本使用了一个已弃用的PyTorch设备管理功能，这使得它无法支持像mps这样的新Torch设备。Thinc在各个设备特定的Ops类中实现自己的操作。在Thinc 8.1之前，提供了以下Ops实现：

• NumpyOps：在CPU上执行操作。使用NumPy和额外的C++内核。
• CupyOps：在支持CUDA的GPU上执行操作。使用CuPy和额外的CUDA C++内核。
• AppleOps：继承自NumpyOps，通过利用苹果的Accelerate框架，覆盖矩阵乘法以在AMX块上运行。
• BigEndianOps：继承自NumpyOps，覆盖特定操作以支持大端平台。

每个Thinc层都与一个Ops类的实例相关联。该层使用Ops实例来分配参数、执行计算等。PyTorch Thinc层与常规Thinc层的不同之处在于，它们使用PyTorch自身的操作，而不是关联的Ops实例。然而，当我们在使用CupyOps时包装一个PyTorch层，我们希望PyTorch层在CUDA设备上运行，而不是默认的CPU设备。Thinc过去通过使用现已弃用的torch.set_default_tensor_type函数，根据当前活动的Ops实例，将默认张量类型设置为torch.cuda.FloatTensor或torch.FloatTensor来实现这一点。

但是，set_default_tensor_type函数不允许我们将默认设备设置为mps。因此，出于这个原因（以及其他原因），我们必须用某种使用Torch设备标识符的机制来替换它，就像上面的矩阵乘法示例那样。从Thinc 8.1开始，PyTorch包装器添加了一个关键字参数来指定层应放置于哪个Torch设备上。如果未指定此参数，Thinc将根据当前活动的Ops使用适当的设备。

为了支持苹果芯片GPU，我们添加了一个新的Ops实现——MPSOps，它默认将Torch层置于mps设备上。当您安装Thinc 8.1并请求Thinc或spaCy使用GPU时，MPSOps会自动被使用。

速度有多快？⏱️

随着Thinc 8.1和PyTorch 1.13的推出，所有环节都已就绪，我们可以在苹果芯片GPU上进行Transformer推理。下表显示了在不同苹果芯片Mac上使用de_dep_news_trf Transformer模型为德语文本做标注时的速度（单位：每秒字数）：

基准测试显示，与使用AMX块相比，使用苹果芯片GPU时速度有显著提升，在M1 Max上，GPU达到每秒8648字，而AMX块为每秒1821字。M1 Max的推理性能几乎是NVIDIA RTX 3090的一半。

8个M1 GPU核心的计算性能估计大约是AMX块的两倍，但结果显示，在M1 Pro上推理速度比AMX块快两倍多，尽管该特定型号只有两个带有AMX块的性能集群和14个GPU核心。原因是AMX只加速矩阵乘法，而GPU也加速其他内核，包括GELU和Softmax非线性函数。下图显示了使用AMX块加速推理时五个最耗时的部分：

由于AMX只加速矩阵乘法，非线性计算已成为最大的耗时部分。这对于GPU推理来说不是问题，因为非线性计算是在GPU上并行进行的。

另一个有趣的问题是，苹果芯片GPU提高的吞吐量是否以更高的功耗为代价。下表显示了基准测试期间的平均功耗（单位：瓦特）。在GPU上运行spaCy Transformer提供了更高的每瓦性能。

在苹果芯片GPU上试用spaCy Transformer流水线

对苹果芯片GPU的支持已在Thinc 8.1.0、spaCy 3.4.2及更高版本以及spacy-transformers 1.1.8及更高版本中提供。要使用对苹果芯片GPU的支持，请首先确保已安装PyTorch 1.13或更高版本：

pip install spacy "torch>=1.13.0"

然后，您可以安装想要使用的Transformer模型，这也会安装spacy-transformers包：

spacy download de_dep_news_trf

之后，您可以在切换到使用GPU（通过require_gpu函数）后像往常一样使用spaCy：

>>> import spacy
>>> spacy.require_gpu()
>>> nlp = spacy.load('de_dep_news_trf')
>>> docs = list(nlp.pipe(["Marbach am Neckar ist eine Stadt etwa 20 Kilometer nördlich von Stuttgart."]))
>>> [(t.text, t.pos_) for t in docs[0]]
[('Marbach', 'PROPN'), ('am', 'ADP'), ('Neckar', 'PROPN'), ('ist', 'AUX'), ('eine', 'DET'), ('Stadt', 'NOUN'), ('etwa', 'ADV'), ('20', 'NUM'), ('Kilometer', 'NOUN'), ('nördlich', 'ADV'), ('von', 'ADP'), ('Stuttgart', 'PROPN'), ('.', 'PUNCT')]

如果您想确认是否确实使用了GPU，可以检查当前活动的Ops是否为MPSOps：

>>> from thinc.api import get_current_ops
>>> get_current_ops()
<thinc.backends.mps_ops.MPSOps object at 0x1010b6e90>

要跟踪对苹果芯片GPU支持的更新，您可以关注我们在Thinc仓库中的跟踪问题。FINISHED