使用 AWS Graviton 处理器加速 PyTorch 推理，来自 AWS 机器学习博客

加速 AWS Graviton 处理器上的 PyTorch 推理

关键要点

最初，PyTorch 使用短期运行模式，在此模式中，每个 PyTorch 操作在被调用时即刻执行。PyTorch 20 引入了 torchcompile，以在默认模式的基础上加速 PyTorch 代码。与短期运行模式相对，torchcompile 会将整个模型预编译成一个最佳适用于特定硬件平台的单一图形。AWS 针对 AWS Graviton3 处理器 优化了 torchcompile 功能。这一优化结果使多种自然语言处理 (NLP)、计算机视觉 (CV) 和推荐模型在 AWS Graviton3 基础的 Amazon EC2 实例上，与默认短期运行模式推理相比，性能提升高达 2 倍基于 33 个模型的几何均值提升。

为什么使用 torchcompile？

在短期运行模式下，模型中的操作在遇到时会立即执行，这种模式更便于使用，也更适合机器学习 (ML) 研究者，因此成为默认模式。但是，这种模式的缺点在于会产生运行时开销，尤其是冗余的内核启动和内存读取。而在 torchcompile 模式下，操作首先合成到一个图中，从而合并操作以减少内存读取和内核启动开销。

AWS Graviton 团队的目标是优化 torchcompile 后端，以实现更高效的 Graviton3 处理器性能。PyTorch 的短期运行模式已经对 Graviton3 进行了优化，使用 Arm 计算库 (ACL) 及 oneDNN 也称为 MKLDNN。因此，关键问题在于如何在 torchcompile 模式中重用这些内核，以同时获得图形编译和优化内核性能的最佳效果。

优化结果

AWS Graviton 团队扩展了 torch 的 inductors 和 oneDNN 原语，重用了 ACL 内核，提升了在 Graviton3 处理器上的编译模式性能。从 PyTorch 231 起，优化功能已包含在 torch Python wheels 和 AWS Graviton 深度学习容器中。关于如何运行推理的具体说明，请参见下文的 运行推理 部分。

为了展示性能提升，我们使用了来自 TorchBench 的 NLP、CV 和推荐模型，以及来自 Hugging Face 的最流行的 NLP 模型，涵盖了问答、文本分类、翻译、文本生成等多种客户用例。

我们首先测量了 TorchBench 模型的推理延迟，并与短期模式的性能标记为 10进行比较。经过对 45 个模型的基准测试，我们发现延迟性能改善了 135 倍45 个模型的几何均值。

类似于以上 TorchBench 推理性能图，我们开始对 Hugging Face NLP 模型的推理延迟进行测量。经过对 33 个模型的基准测试，性能提升接近 2 倍33 个模型的几何均值。

运行推理

从 PyTorch 231 开始，优化功能已包含在 torch Python wheels 和 AWS Graviton 深度学习容器中。本节将展示如何使用 torch Python wheels 和 Hugging Face 及 TorchBench 仓库中的基准测试脚本运行推理。

要成功运行脚本并重现本文中提到的速度提升数字，您需要使用 Graviton3 系列硬件的实例如 c7g/r7g/m7g/hpc7g。在本次测试中，我们使用的是 c7g4xl (16 vcpu) 实例，相关实例、AMI 详情及所需的 torch 库版本如下所示：

bash实例：c7g4xl 实例区域：uswest2AMI：ami05cc25bfa725a144a使用 6501017aws 内核的 Ubuntu 2204/Jammy

安装 Python

sudo aptget updatesudo aptget install y python3 python3pip

将 pip3 升级至最新版本

python3 m pip install upgrade pip

安装 PyTorch 及扩展

python3 m pip install torch==231 torchvision==0181 torchaudio==231

为进一步提高 AWS Graviton3 处理器上的 torchcompile 性能，我们设置了以下环境变量：

bash

启用快速数学 GEMM 内核，以加速 fp32 推理与 bfloat16 gemm

export DNNLDEFAULTFPMATHMODE=BF16

启用 Linux 透明大页面 (THP) 分配，以减少张量内存分配延迟

export THPMEMALLOCENABLE=1

设置 LRU 缓存容量以缓存原语并避免冗余的内存分配

export LRUCACHECAPACITY=1024

TorchBench 基准测试脚本

TorchBench 是一个开源基准集，用于评估 PyTorch 性能。我们使用 TorchBench 仓库中的脚本对 45 个模型进行基准测试。以下代码展示了如何在短期模式和 torchcompile 模式下运行这些脚本。

bash

设置 OMPNUMTHREADS 为 vcpu 数量，对应 c7g4xl 实例为 16

export OMPNUMTHREADS=16

安装所需依赖

sudo aptget install y libgl1mesaglxsudo aptget install y libpangocairo100python3 m pip install psutil numpy transformers pynvml numba onnx onnxruntime scikitlearn timm effdet gym doctr opencvpython h5py==3100 pythondoctr

克隆 PyTorch 基准测试仓库

git clone https//githubcom/pytorch/benchmarkgitcd benchmark

PyTorch 基准测试仓库没有发布标签，因此列出我们用于收集性能数据的提交

git checkout 9a5e4137299741e1b6fb7aa7f5a6a853e5dd2295

设置模型

python3 installpy

使用以下命令收集短期模式性能，结果将存储在 userbenchmark/cpu/metricjson。

python3 runbenchmarkpy cpu model BERTpytorchhfBerthfBertlargehfGPT2hfAlberthfBarthfBigBirdhfDistilBerthfGPT2largedlrmhfT5mnasnet10mobilenetv2mobilenetv3largesqueezenet11timmefficientnetshufflenetv2x10timmregnetresnet50softactorcriticphlippedensenetresnet152resnet18resnext5032x4ddensenet121phlipperesnetdoctrdetpredictortimmvovnetalexnetdoctrrecopredictorvgg16dcganyolov3pytorchstarganhfLongformertimmnfnettimmvisiontransformertimmvisiontransformerlargenvidiadeeprecommenderdemucsttsangularhfReformerpytorchCycleGANandpix2pixfunctorchdpcifar10pytorchunet test eval metrics=latenciescpupeakmem

收集 torchcompile 模式的性能，启用 inductors 后端及权重预打包。结果将存储在 userbenchmark/cpu/metricjson

python3 runbenchmarkpy cpu model BERTpytorchhfBerthfBertlargehfGPT2hfAlberthfBarthfBigBirdhfDistilBerthfGPT2largedlrmhfT5mnasnet10mobilenetv2mobilenetv3largesqueezenet11timmefficientnetshufflenetv2x10timmregnetresnet50softactorcriticphlippedensenetresnet152resnet18resnext5032x4ddensenet121phlipperesnetdoctrdetpredictortimmvovnetalexnetdoctrrecopredictorvgg16dcganyolov3pytorchstarganhfLongformertimmnfnettimmvisiontransformertimmvisiontransformerlargenvidiadeeprecommenderdemucsttsangularhfReformerpytorchCycleGANandpix2pixfunctorchdpcifar10pytorchunet test eval torchdynamo inductor freezeprepackweights metrics=latenciescpupeakmem

当推理运行成功完成后，脚本将以 JSON 格式存储结果，样本输出如下：

json{ name cpu environ { pytorchgitversion d44533f9d073df13895333e70b66f81c513c1889 } metrics { BERTpytorchevallatency 563769865 BERTpytorchevalcmem 04169921875 }}

Hugging Face 基准测试脚本

谷歌 T5 小型文本翻译模型是我们基准测试的约 30 个 Hugging Face 模型之一。我们使用该模型作为示例，演示如何在短期和编译模式下运行推理。运行编译模式所需的额外配置和 API 用 粗体 高亮显示。将以下脚本保存为 googlet5smalltexttranslationpy。

pythonimport argparsefrom transformers import T5Tokenizer T5Modelimport torchfrom torchprofiler import profile recordfunction ProfilerActivityimport torchinductorconfig as config configcppweightprepack=True configfreezing=True

def testinference(mode numiter) tokenizer = T5Tokenizerfrompretrained(t5small) model = T5Modelfrompretrained(t5small)

inputids = tokenizer(    Studies have been shown that owning a dog is good for you returntensors=pt)inputids  # Batch size 1decoderinputids = tokenizer(Studies show that returntensors=pt)inputids  # Batch size 1if (mode == compile) model = torchcompile(model)with torchnograd()    for  in range(50)        outputs = model(inputids=inputids decoderinputids=decoderinputids)with profile(activities=[ProfilerActivityCPU]) as prof    with recordfunction(modelinference)        for  in range(numiter)            outputs = model(inputids=inputids decoderinputids=decoderinputids)print(profkeyaverages()table(sortby=selfcputimetotal))

def main() gt None global m args parser = argparseArgumentParser(doc) parseraddargument( m mode choices=[eager compile] default=eager help=Which test to run ) parseraddargument( n number type=int default=100 help=how many iterations to run ) args = parserparseargs() testinference(argsmode argsnumber)

if name == main main()

执行脚本的步骤如下：

bash

设置 OMPNUMTHREADS 为 vcpu 数量为 4，因为

脚本按顺序运行推理，不需要大数量的 vcpu

export OMPNUMTHREADS=4

安装所需依赖

python3 m pip install transformers

在短期模式下运行推理脚本

迭代次数设为 1，仅用于展示 torch profiler 输出

但我们基准测试时使用了 1000 次迭代

python3 googlet5smalltexttranslationpy n 1 m eager

在 torch 编译模式下运行推理脚本

python3 googlet5smalltexttranslationpy n 1 m compile

成功完成推理运行后，脚本将打印出具有延迟细分的 torch profiler 输出。下面是来自 torch profiler 的样本输出：

bash

Eager 模式在 c7gxl (4vcpu) 上运行的 torch profiler 输出

Name Self CPU Self CPU CPU total CPU total CPU time avg # of Calls

atenmm 4071 12502ms 4071 12502ms 130229us 96modelinference 2644 8118ms 10000 30708ms 30708ms 1atenbmm 685 2102ms 947 2908ms 80778us 36atenmatmul 373 1146ms 5726 17583ms 133205us 132atenselect 188 576000us 190 583000us 0998us 584atentranspose 151 464000us 183 563000us 3027us 186

Self CPU time total 30708ms

Compile 模式在同一实例上相同模型的 torch profiler 输出

Name Self CPU Self CPU CPU total CPU total CPU time avg # of Calls

mkldnnlinearpointwise 3798 5461ms 4591 6602ms 68771us 96TorchCompiled Region 2956 4251ms 9853 14168ms 14168ms 1atenbmm 1490 2143ms 2173 3124ms 86778us 36atenselect 451 648000us 462 665000us 1155us 576atenview 329 473000us 329 473000us 1642us 288atenempty 253 364000us 253 364000us 3165us 115

Self CPU time total 14379ms

接下来

接下来，我们将扩展 torch 的 inductor CPU 后端支持，以编译 Llama 模型，并为 AWS Graviton3 处理器添加对融合 GEMM 内核的支持，以启用 torch inductor 操作符融合优化。

结论

在本教程中，我们讨论了如何优化 AWS Graviton3 基础 EC2 实例上的 torchcompile 性能，如何使用这些优化来提升 PyTorch 模型推理性能，并演示了由此带来的速度提升。我们希望您也能够尝试一下！如果您在 Graviton 上的 ML 软件方面需要任何帮助