利用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_trfspaCy Transformer模型时五个最耗时的部分Transformer模型的运行时间主要由矩阵乘法主导——bli_sgemm_armv8a_asm_8x12是一个单精度矩阵乘法内核。sgemm是线性代数库提供的标准化矩阵乘法函数这些库实现了BLAS接口。这并不令人意外因为矩阵乘法是Transformer模型使用的主要操作之一例如用于计算注意力块中的成对注意力分数以及前馈块中的线性投影。除了矩阵乘法分别用于前馈块和注意力块的GeLU和Softmax非线性函数也通过其对erff、expf和exp初等函数的使用在性能剖析中显现出来。这些非线性函数合计占据了19%的运行时间。为了加速Transformer推理我们可以采取三种不同的方法用时间复杂度优于O(N²)的机制替换自注意力机制。例如Longformer的注意力机制具有O(N)的时间复杂度。如果我们希望继续使用现有的预训练Transformer模型也可以对N设置一个上限。加速矩阵乘法。加速非线性计算。在spaCy Transformer流水线中二次注意力机制的影响已经通过限制N的上限得到了控制。每个文档被分步处理。默认情况下spaCy Transformer每次处理96个标记使用128个标记的窗口来创建重叠的上下文表示。这将注意力机制的时间复杂度上限设定为N128。在本文中我们将重点关注另外两种方法——利用苹果芯片Mac的专用硬件加速矩阵乘法和非线性计算。我们将首先了解苹果硅CPU中称为“AMX块”的矩阵乘法协处理器。之后我们将探索在GPU上运行计算内核。AMX矩阵乘法块所有Apple M系列CPU至少拥有一个称为“AMX块”的矩阵乘法协处理器。AMX在很大程度上没有公开文档。例如目前尚不清楚Apple M系列CPU的高能效核心集群是否拥有自己的AMX块。然而我们可以通过基准测试来推断AMX块的各种特性。下表列出了在不同CPU上使用gemm-benchmark测量得到的768x768矩阵乘法性能单位TFLOPS即每秒万亿次浮点运算线程数M1M2M1 Pro/MaxM1 UltraRyzen 5950X11.31.52.12.20.121.21.62.63.40.341.01.72.73.80.681.31.62.54.31.0121.21.52.44.31.6161.21.42.44.41.9与M1相比的最大加速1.01.32.13.41.5从这些数字中我们可以得出一些有趣的信息性能并不随线程数增加而提升。因此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核心上进行矩阵乘法importtorchutorch.rand((10,20),dtypetorch.float,devicetorch.device(mps))vtorch.rand((20,10),dtypetorch.float,devicetorch.device(mps))torch.matmul(u,v).device device(typemps,index0)在撰写本文时某些操作尚未实现但在这种情况下当环境变量PYTORCH_ENABLE_MPS_FALLBACK设置为1时PyTorch将回退到CPU内核。spaCy和Thinc中的Metal Performance ShadersspaCy使用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_trfTransformer模型为德语文本做标注时的速度单位每秒字数机器CPU核心GPU核心AMX (WPS)GPU (WPS)加速比Mac Mini M14P/4E8118022021.9MacBook Air M24P/4E10124233622.7MacBook Pro 14” M1 Pro6P/2E14163146612.9MacBook Pro 14” M1 Max8P/2E32182186484.7Mac Studio M1 Ultra16P/4E482197120735.5Ryzen 5950X RTX 309016328 (Tensor cores)1879 (CPU)1884510.0基准测试显示与使用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提供了更高的每瓦性能。机器CPU核心GPU核心AMX (W)GPU (W)Mac Mini M14P/4E81110MacBook Air M24P/4E10139MacBook Pro 14” M1 Pro6P/2E141617MacBook Pro 14” M1 Max8P/2E321731Mac Studio M1 Ultra16P/4E483470在苹果芯片GPU上试用spaCy Transformer流水线对苹果芯片GPU的支持已在Thinc 8.1.0、spaCy 3.4.2及更高版本以及spacy-transformers 1.1.8及更高版本中提供。要使用对苹果芯片GPU的支持请首先确保已安装PyTorch 1.13或更高版本pipinstallspacytorch1.13.0然后您可以安装想要使用的Transformer模型这也会安装spacy-transformers包spacy download de_dep_news_trf之后您可以在切换到使用GPU通过require_gpu函数后像往常一样使用spaCyimportspacyspacy.require_gpu()nlpspacy.load(de_dep_news_trf)docslist(nlp.pipe([Marbach am Neckar ist eine Stadt etwa 20 Kilometer nördlich von Stuttgart.]))[(t.text,t.pos_)fortindocs[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是否为MPSOpsfromthinc.apiimportget_current_opsget_current_ops()thinc.backends.mps_ops.MPSOpsobjectat0x1010b6e90要跟踪对苹果芯片GPU支持的更新您可以关注我们在Thinc仓库中的跟踪问题。FINISHED更多精彩内容 请关注我的个人公众号 公众号办公AI智能小助手或者 我的个人博客 https://blog.qife122.com/对网络安全、黑客技术感兴趣的朋友可以关注我的安全公众号网络安全技术点滴分享