使用AWS Graviton處理器優化的PyTorch 2.0推理

作者：Sunita Nadampalli

新一代的CPU因為內置了專門的指令，在機器學習（ML）推理方面提供了顯著的性能提升。結合它們的靈活性、高速開發和低運營成本，這些通用處理器為其他現有硬件解決方案提供了一種替代選擇。

AWS、Arm、Meta等公司幫助優化了基于Arm處理器的PyTorch 2.0推理性能。因此，我們很高興地宣布，AWS Graviton基于實例的PyTorch 2.0推理性能比之前的PyTorch版本提高了3.5倍，Resnet50的速度（請參見下圖），BERT的速度提高了1.4倍，使Graviton基于實例成為AWS上這些模型最快的計算優化實例。

AWS使用基于AWS Graviton3的Amazon Elastic Cloud Compute C7g實例進行PyTorch推理，相對于可比較的EC2實例，跨Torch Hub Resnet50和多個Hugging Face模型，可節省高達50%的成本，如下圖所示。

這是因為AWS Graviton3處理器是最新一代定制的AWS Graviton處理器，可為Amazon Elastic Compute Cloud（Amazon EC2）中的工作負載提供最佳價格性能。它們提供高達2倍的浮點性能、高達2倍的加密性能和高達3倍的ML性能，包括對PyTorch的支持。

此外，推理的延遲也減少了，如下圖所示。

我們在Graviton上的其他工作負載中也看到了類似的價格性能優勢趨勢，例如使用FFmpeg進行視頻編碼（https://aws.amazon.com/blogs/opensource/optimized-video-encoding-with-ffmpeg-on-aws-graviton-processors/）。

優化細節

優化集中在三個關鍵領域：

1. GEMM內核 - PyTorch通過OneDNN后端（以前稱為MKL-DNN）支持Arm Compute Library（ACL）GEMM內核，用于基于Arm的處理器。ACL庫為Neon和SVE優化了GEMM內核，適用于fp32和bfloat16格式。這些內核提高了SIMD硬件利用率并降低了端到端推理延遲。

2. bfloat16支持 - Graviton3中的bfloat16支持允許有效部署使用bfloat16，fp32和AMP（自動混合精度）訓練的模型。標準fp32模型通過OneDNN快速數學模式使用bfloat16內核，無需模型量化，與不帶bfloat16快速數學支持的現有fp32模型推理相比，性能提高了兩倍。

3. 原始緩存 - 我們還為conv、matmul和inner product運算符實現了原始緩存，以避免冗余的GEMM內核初始化和張量分配開銷。

如何利用這些優化

最簡單的方法是使用Amazon Elastic Compute Cloud（Amazon EC2）C7g實例或Amazon SageMaker上的AWS Deep Learning Containers（DLC）。DLC可在Amazon Elastic Container Registry（Amazon ECR）上提供AWS Graviton或x86。有關SageMaker的更多詳細信息，請參閱在基于AWS Graviton的實例上運行機器學習推理工作負載（https://aws.amazon.com/blogs/machine-learning/run-machine-learning-inference-workloads-on-aws-graviton-based-instances-with-amazon-sagemaker/）以及Amazon SageMaker添加了八個基于Graviton的實例以進行模型部署（https://aws.amazon.com/about-aws/whats-new/2022/10/amazon-sagemaker-adds-new-graviton-based-instances-model-deployment/）。

使用AWS DLC

要使用AWS DLC，請使用以下代碼：

udo apt-get update
sudo apt-get -y install awscli docker


# Login to ECR to avoid image download throttling
aws ecr get-login-password --region us-east-1 
| docker login --username AWS 
  --password-stdin 763104351884.dkr.ecr.us-east-1.amazonaws.com


# Pull the AWS DLC for pytorch
# Graviton
docker pull 763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-inference-graviton:2.0.0-cpu-py310-ubuntu20.04-ec2


# x86
docker pull 763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-inference:2.0.0-cpu-py310-ubuntu20.04-ec2

如果您喜歡通過pip安裝PyTorch，請從官方repo安裝PyTorch 2.0 wheel。在這種情況下，在啟動PyTorch之前，您將需要設置兩個環境變量，如下面的代碼所述，以激活Graviton優化。

使用Python wheel

要使用Python wheel，請參考以下代碼：

# Install Python
sudo apt-get update
sudo apt-get install -y python3 python3-pip


# Upgrade pip3 to the latest version
python3 -m pip install --upgrade pip


# Install PyTorch and extensions
python3 -m pip install torch
python3 -m pip install torchvision torchaudio torchtext


# Turn on Graviton3 optimization
export DNNL_DEFAULT_FPMATH_MODE=BF16
export LRU_CACHE_CAPACITY=1024

運行推斷

可以使用PyTorch TorchBench測量CPU推理性能改進，或比較不同的實例類型：

# Pre-requisite: 
# pull and run the AWS DLC
# or 
# pip install PyTorch2.0 wheels and set the previously mentioned environment variables


# Clone PyTorch benchmark repo
git clone https://github.com/pytorch/benchmark.git


# Setup Resnet50 benchmark
cd benchmark
python3 install.py resnet50


# Install the dependent wheels
python3 -m pip install numba


# Run Resnet50 inference in jit mode. On successful completion of the inference runs,
# the script prints the inference latency and accuracy results
python3 run.py resnet50 -d cpu -m jit -t eval --use_cosine_similarity

性能基準測試

您可以使用AmazonSageMaker推理推薦實用程序（https://docs.aws.amazon.com/sagemaker/latest/dg/inference-recommender.html）來自動化不同實例之間的性能基準測試。使用推理推薦程序，您可以找到實時推理端點，該端點可以為給定的ML模型以最低的成本提供最佳性能。我們通過在生產端點上部署模型，使用推理推薦器筆記本收集了前面的數據。有關推理推薦程序的更多詳細信息，請參閱GitHub repo（https://github.com/aws/amazon-sagemaker-examples/blob/main/sagemaker-inference-recommender/huggingface-inference-recommender/huggingface-inference-recommender.ipynb）。我們為這篇文章測試了以下模型：ResNet50圖像分類（https://pytorch.org/hub/pytorch_vision_resnet/）、DistilBERT情緒分析（https://huggingface.co/distilbert-base-uncased-finetuned-sst-2-english）、RoBERTa填充掩碼（https://huggingface.co/roberta-base）和RoBERTa情緒分析（https://huggingface.co/cardiffnlp/twitter-roberta-base-sentiment）。

結論

AWS通過Torch Hub Resnet50中基于AWS Graviton3的亞馬遜彈性云計算C7g實例，以及相對于可比EC2實例的多個擁抱人臉模型，為PyTorch推理節省了高達50%的成本。這些實例可在SageMaker和AmazonEC2上獲得。AWS Graviton技術指南（https://github.com/aws/aws-graviton-getting-started）提供了優化的庫和最佳實踐列表，這些庫和最佳做法將幫助您在不同的工作負載中使用Graviton實例實現成本效益。

如果您發現在AWS Graviton上沒有觀察到類似性能提升的用例，請在AWS Gravaton技術指南（https://github.com/aws/aws-graviton-getting-started）上提交問題，讓我們了解它。我們將繼續添加更多性能改進，使Graviton成為使用PyTorch進行推理的最具成本效益和效率的通用處理器。