使用Metal性能着色器加速Transformer推理

背景介绍

大型Transformer模型以计算昂贵著称，这源于自注意力机制的二次计算复杂度以及模型本身的庞大参数量。例如广泛使用的BERT、RoBERTa和XLM-R基础模型采用12个隐藏层，隐藏表示维度为768，前馈网络块中的表示维度达3072。

性能瓶颈分析

通过性能剖析发现，Transformer模型运行时间主要被矩阵乘法占据——bli_sgemm_armv8a_asm_8x12是单精度矩阵乘法内核。sgemm是线性代数库提供的标准化矩阵乘法函数，实现了BLAS接口。这并不意外，因为矩阵乘法是Transformer模型的主要操作，例如用于计算注意力块中的配对注意力分数和前馈块中的线性投影。

除了矩阵乘法，在前馈块和注意力块中分别使用的GeLU和Softmax非线性函数也通过erff、expf和exp基本函数在性能剖析中可见。这些非线性函数合计占运行时间的19%。

加速方案

为加速Transformer推理，可采用三种方法：

用时间复杂度优于O(N²)的机制替换自注意力机制
加速矩阵乘法运算
加速非线性函数计算

在spaCy transformer管道中，通过限制N的上界已经限制了二次注意力机制的影响。每个文档被分步处理，默认情况下spaCy transformer每次处理96个词符，使用128个词符的窗口创建重叠的上下文表示。

AMX矩阵乘法单元

所有苹果M系列CPU都至少包含一个称为"AMX块"的矩阵乘法协处理器。通过基准测试可以推断AMX块的各种特性：

性能测试显示：

性能不随线程数增加而提升，表明AMX块不属于单个CPU核心
M1、M1 Pro和M1 Ultra分别具有1、2和4个性能核心集群，矩阵乘法性能随性能核心集群数量增加
每个性能集群似乎都有一个AMX块
AMX块速度极快，单个AMX块的矩阵乘法性能相当于9个Ryzen 5950X核心

虽然苹果未公开向AMX块分派计算的指令说明，但第三方应用可通过苹果的Accelerate框架使用AMX块，该框架实现了行业标准BLAS接口。

Metal性能着色器

尽管AMX块在处理矩阵乘法吞吐量方面表现出色，苹果芯片Mac还有另外两个计算子系统：苹果神经引擎（ANE）和GPU。ANE的限制在于需要运行通过Core ML定义的计算图，而GPU可以运行用户定义的计算内核（所谓的"着色器"）。

M1的8核GPU计算性能为2.6 TFLOPS，提供约M1 AMX单元两倍的性能。此外，GPU在M1 Ultra中最多可扩展至64核，理论峰值性能达20.8 TFLOPS。

PyTorch最近通过苹果Metal API引入了对苹果M系列GPU的支持。各种PyTorch操作已实现为自定义Metal着色器，并使用苹果自己的Metal着色器集合。

spaCy和Thinc中的实现

spaCy使用Thinc作为其机器学习库。Thinc 8.1之前的版本使用已弃用的PyTorch设备管理功能，无法支持mps等新Torch设备。为支持苹果硅GPU，新增了MPSOps实现，该实现默认为Torch层使用mps设备。

性能基准测试

使用Thinc 8.1和PyTorch 1.13后，我们可以在苹果硅GPU上执行Transformer推理。基准测试显示：

在M1 Max上，GPU处理速度达到8648词/秒，而AMX块为1821词/秒，推理性能接近NVIDIA RTX 3090的一半。虽然8个M1 GPU核心的计算性能估计约为AMX块的两倍，但实际推理速度在M1 Pro上超过两倍。

原因是AMX仅加速矩阵乘法，而GPU还加速其他内核，包括GELU和Softmax非线性函数。当使用AMX块加速推理时，非线性函数的计算已成为最大的成本中心。

能效比较

基准测试期间的平均功耗数据显示，在GPU上运行spaCy transformer每瓦特提供更高的性能。

实践指南

苹果硅GPU支持在Thinc 8.1.0、spaCy 3.4.2和spacy-transformers 1.1.8及更高版本中可用。使用前需确保安装PyTorch 1.13或更高版本，并通过require_gpu()函数切换到GPU使用。

用户可通过检查当前活动的Ops是否为MPSOps来验证GPU是否确实被使用。要跟踪苹果硅GPU支持的更新，可关注Thinc代码库中的跟踪问题。