Inferencia acelerada de PyTorch con Torch.compile en procesadores Graviton de AWS

Originalmente, PyTorch usaba un modo ansioso en el que cada operación de PyTorch que forma el modelo se ejecuta de forma independiente tan pronto como se alcanza. PyTorch 2.0 introdujo antorcha.compilar para acelerar el código de PyTorch en comparación con el modo ansioso predeterminado. A diferencia del modo ansioso, el archivo Torch.compile compila previamente todo el modelo en un solo gráfico de una manera óptima para ejecutarse en una plataforma de hardware determinada. AWS optimizó la función Torch.compile de PyTorch para Procesadores AWS Graviton3Esta optimización da como resultado un rendimiento hasta dos veces mejor para Cara abrazada Inferencia de modelos (basada en la media geométrica de mejora del rendimiento para 33 modelos) y hasta 1,35 veces mejor rendimiento para Banco de antorcha Inferencia de modelos (media geométrica de mejora del rendimiento para 45 modelos) en comparación con la inferencia de modo ansioso predeterminada en varios modelos de procesamiento de lenguaje natural (NLP), visión artificial (CV) y recomendación en instancias de Amazon EC2 basadas en AWS Graviton3. A partir de PyTorch 2.3.1, las optimizaciones están disponibles en Python de Torch ruedas y AWS Graviton PyTorch contenedor de aprendizaje profundo (DLC).

En esta publicación de blog, mostramos cómo optimizamos el rendimiento de Torch.Compile en instancias EC2 basadas en AWS Graviton3, cómo usar las optimizaciones para mejorar el rendimiento de la inferencia y las aceleraciones resultantes.

¿Por qué Torch.Compile y cuál es el objetivo?

En el modo ansioso, los operadores de un modelo se ejecutan inmediatamente cuando se encuentran. Es más fácil de usar, más adecuado para los investigadores de aprendizaje automático (ML) y, por lo tanto, es el modo predeterminado. Sin embargo, el modo ansioso genera una sobrecarga de tiempo de ejecución debido al lanzamiento redundante del kernel y la sobrecarga de lectura de memoria. Mientras que en el modo de compilación Torch, los operadores primero se sintetizan en un gráfico, donde un operador se fusiona con otro para reducir y localizar las lecturas de memoria y la sobrecarga total de lanzamiento del kernel.

El objetivo del equipo de AWS Graviton era optimizar el backend de Torch.Compile para los procesadores Graviton3. El modo entusiasta de PyTorch ya estaba optimizado para los procesadores Graviton3 con Biblioteca de cómputo Arm (ACL) núcleos que utilizan oneDNN (también conocido como MKLDNN). Entonces, la pregunta era, ¿cómo reutilizar esos núcleos en modo de compilación de antorcha para obtener lo mejor de la compilación de gráficos y el rendimiento optimizado del núcleo juntos?

Resultados

El equipo de AWS Graviton extendió el inductor de antorcha y los primitivos oneDNN que reutilizaron los núcleos ACL y optimizaron el rendimiento del modo de compilación en los procesadores Graviton3. A partir de PyTorch 2.3.1, las optimizaciones están disponibles en las ruedas Python de antorcha y en el DLC de AWS Graviton. Consulte Ejecutar una inferencia Sección que sigue para obtener instrucciones sobre la instalación, la configuración del tiempo de ejecución y cómo ejecutar las pruebas.

Para demostrar las mejoras de rendimiento, utilizamos modelos de PNL, CV y ​​recomendación de Banco de antorcha y los modelos de PNL más descargados de Cara abrazada en tareas de respuesta a preguntas, clasificación de texto, clasificación de tokens, traducción, clasificación de disparo cero, traducción, resumen, extracción de características, generación de texto, generación de texto a texto, máscara de relleno y similitud de oraciones para cubrir una amplia variedad de casos de uso de clientes.

Comenzamos midiendo la latencia de inferencia del modelo TorchBench, en milisegundos (ms), para el modo ansioso, que está marcado como 1.0 con una línea de puntos roja en el siguiente gráfico. Luego, comparamos las mejoras de Torch.compile para la misma inferencia del modelo; los resultados normalizados se representan en el gráfico. Puede ver que para los 45 modelos que evaluamos, hay una mejora de latencia de 1.35x (media geográfica para los 45 modelos).

Imagen 1: Mejora del rendimiento de inferencia del modelo PyTorch con Torch.Compile en la instancia c7g basada en Graviton3 de AWS mediante el marco TorchBench. El rendimiento del modo de referencia está marcado como 1.0 (cuanto más alto, mejor).

De manera similar al gráfico de rendimiento de inferencia de TorchBench anterior, comenzamos midiendo la latencia de inferencia del modelo NLP Hugging Face, en ms, para el modo ansioso, que está marcado como 1.0 con una línea de puntos roja en el siguiente gráfico. Luego, comparamos las mejoras de Torch.compile para la misma inferencia de modelo; los resultados normalizados se representan en el gráfico. Puede ver que, para los 33 modelos que evaluamos, hay una mejora de rendimiento de alrededor de 2x (media geográfica para los 33 modelos).

Imagen 2: Mejora del rendimiento de la inferencia del modelo NLP de Hugging Face con Torch.Compile en la instancia c7g basada en Graviton3 de AWS mediante scripts de ejemplo de Hugging Face. El rendimiento del modo ansioso de referencia está marcado como 1.0. (cuanto más alto, mejor)

Ejecutar una inferencia

A partir de PyTorch 2.3.1, las optimizaciones están disponibles en la rueda de Python de Torch y en el DLC PyTorch de AWS Graviton. Esta sección muestra cómo ejecutar la inferencia en los modos anxious y Torch.Compile utilizando ruedas de Python de Torch y scripts de evaluación comparativa de los repositorios Hugging Face y TorchBench.

Para ejecutar correctamente los scripts y reproducir los números de aceleración mencionados en esta publicación, necesita una instancia de hardware de la familia Graviton3 (c7g/r7g/m7g/hpc7g). Para esta publicación, usamos el Instancia c7g.4xl (16 vcpu)La instancia, los detalles de la AMI y las versiones de la biblioteca de antorcha requeridas se mencionan en el siguiente fragmento.

Instance: c7g.4xl instance
Region: us-west-2
AMI: ami-05cc25bfa725a144a (Ubuntu 22.04/Jammy with 6.5.0-1017-aws kernel)

# 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==2.3.1 torchvision==0.18.1 torchaudio==2.3.1

Los ajustes de tiempo de ejecución genéricos implementados para la inferencia en modo ansioso son igualmente aplicables para el modo Torch.Compile, por lo tanto, configuramos las siguientes variables de entorno para mejorar aún más el rendimiento de Torch.Compile en los procesadores AWS Graviton3.

# Enable the fast math GEMM kernels, to accelerate fp32 inference with bfloat16 gemm
export DNNL_DEFAULT_FPMATH_MODE=BF16

# Enable Linux Transparent Huge Page (THP) allocations,
# to reduce the tensor memory allocation latency
export THP_MEM_ALLOC_ENABLE=1

# Set LRU Cache capacity to cache the primitives and avoid redundant
# memory allocations
export LRU_CACHE_CAPACITY=1024

Scripts de evaluación comparativa de TorchBench

TorchBench es una colección de pruebas comparativas de código abierto que se utilizan para evaluar el rendimiento de PyTorch. Evaluamos 45 modelos utilizando los scripts del repositorio de TorchBench. El siguiente código muestra cómo ejecutar los scripts para el modo de ejecución en modo de ejecución en modo de ejecución en modo de ejecución en modo de compilación con el backend Inductor.

# Set OMP_NUM_THREADS to number of vcpus, 16 for c7g.4xl instance
export OMP_NUM_THREADS=16

# Install the dependencies
sudo apt-get install -y libgl1-mesa-glx
sudo apt-get install -y libpangocairo-1.0-0
python3 -m pip install psutil numpy transformers pynvml numba onnx onnxruntime scikit-learn timm effdet gym doctr opencv-python h5py==3.10.0 python-doctr

# Clone pytorch benchmark repo
git clone https://github.com/pytorch/benchmark.git
cd benchmark
# PyTorch benchmark repo doesn't have any release tags. So,
# listing the commit we used for collecting the performance numbers
git checkout 9a5e4137299741e1b6fb7aa7f5a6a853e5dd2295

# Setup the models
python3 install.py

# Colect eager mode performance using the following command. The results will be
# stored at .userbenchmark/cpu/metric-<timestamp>.json.
python3 run_benchmark.py cpu --model BERT_pytorch,hf_Bert,hf_Bert_large,hf_GPT2,hf_Albert,hf_Bart,hf_BigBird,hf_DistilBert,hf_GPT2_large,dlrm,hf_T5,mnasnet1_0,mobilenet_v2,mobilenet_v3_large,squeezenet1_1,timm_efficientnet,shufflenet_v2_x1_0,timm_regnet,resnet50,soft_actor_critic,phlippe_densenet,resnet152,resnet18,resnext50_32x4d,densenet121,phlippe_resnet,doctr_det_predictor,timm_vovnet,alexnet,doctr_reco_predictor,vgg16,dcgan,yolov3,pytorch_stargan,hf_Longformer,timm_nfnet,timm_vision_transformer,timm_vision_transformer_large,nvidia_deeprecommender,demucs,tts_angular,hf_Reformer,pytorch_CycleGAN_and_pix2pix,functorch_dp_cifar10,pytorch_unet --test eval --metrics="latencies,cpu_peak_mem"

# Collect torch.compile mode performance with inductor backend
# and weights pre-packing enabled. The results will be stored at
# .userbenchmark/cpu/metric-<timestamp>.json
python3 run_benchmark.py cpu --model BERT_pytorch,hf_Bert,hf_Bert_large,hf_GPT2,hf_Albert,hf_Bart,hf_BigBird,hf_DistilBert,hf_GPT2_large,dlrm,hf_T5,mnasnet1_0,mobilenet_v2,mobilenet_v3_large,squeezenet1_1,timm_efficientnet,shufflenet_v2_x1_0,timm_regnet,resnet50,soft_actor_critic,phlippe_densenet,resnet152,resnet18,resnext50_32x4d,densenet121,phlippe_resnet,doctr_det_predictor,timm_vovnet,alexnet,doctr_reco_predictor,vgg16,dcgan,yolov3,pytorch_stargan,hf_Longformer,timm_nfnet,timm_vision_transformer,timm_vision_transformer_large,nvidia_deeprecommender,demucs,tts_angular,hf_Reformer,pytorch_CycleGAN_and_pix2pix,functorch_dp_cifar10,pytorch_unet --test eval --torchdynamo inductor --freeze_prepack_weights --metrics="latencies,cpu_peak_mem"

Una vez finalizadas las ejecuciones de inferencia, el script almacena los resultados en formato JSON. A continuación, se muestra el resultado de muestra:

{
"name": "cpu"
"environ": {
"pytorch_git_version": "d44533f9d073df13895333e70b66f81c513c1889"
},

"metrics": {
"BERT_pytorch-eval_latency": 56.3769865,
"BERT_pytorch-eval_cmem": 0.4169921875
}
}

Scripts de evaluación comparativa de Hugging Face

El modelo de traducción de texto pequeño T5 de Google es uno de los 30 modelos Hugging Face que evaluamos como referencia. Lo estamos usando como modelo de muestra para demostrar cómo ejecutar la inferencia en los modos de compilación y de preparación. Las configuraciones y API adicionales necesarias para ejecutarlo en el modo de compilación se destacan en ATREVIDOGuarde el siguiente script como google_t5_small_text_translation.py .

import argparse
from transformers import T5Tokenizer, T5Model
import torch
from torch.profiler import profile, record_function, ProfilerActivity
import torch._inductor.config as config config.cpp.weight_prepack=True config.freezing=True

def test_inference(mode, num_iter):
tokenizer = T5Tokenizer.from_pretrained("t5-small")
model = T5Model.from_pretrained("t5-small")

input_ids = tokenizer(
"Studies have been shown that owning a dog is good for you", return_tensors="pt"
).input_ids  # Batch size 1
decoder_input_ids = tokenizer("Studies show that", return_tensors="pt").input_ids  # Batch size 1

    if (mode == 'compile'):         model = torch.compile(model)

with torch.no_grad():
for _ in range(50):
outputs = model(input_ids=input_ids, decoder_input_ids=decoder_input_ids)

with profile(activities=[ProfilerActivity.CPU]) as prof:
with record_function("model_inference"):
for _ in range(num_iter):
outputs = model(input_ids=input_ids, decoder_input_ids=decoder_input_ids)

print(prof.key_averages().table(sort_by="self_cpu_time_total"))

def main() -> None:
global m, args
parser = argparse.ArgumentParser(__doc__)
parser.add_argument(
"-m",
"--mode",
choices=["eager", "compile"],
default="eager",
help="Which test to run.",
)
parser.add_argument(
"-n",
"--number",
type=int,
default=100,
help="how many iterations to run.",
)
args = parser.parse_args()
test_inference(args.mode, args.number)

if __name__ == "__main__":
main()

Ejecute el script con los siguientes pasos.

# Set OMP_NUM_THREADS to number of vcpus to 4 because
# the scripts are running inference in sequence, and
# they don't need large number of vcpus
export OMP_NUM_THREADS=4

# Install the dependencies
python3 -m pip install transformers

# Run the inference script in Eager mode
# using number of iterations as 1 just to show the torch profiler output
# but for the benchmarking, we used 1000 iterations.
python3 google_t5_small_text_translation.py -n 1 -m eager

# Run the inference script in torch compile mode
python3 google_t5_small_text_translation.py -n 1 -m compile

Una vez finalizadas las ejecuciones de inferencia, el script imprime la salida del generador de perfiles de antorcha con el desglose de latencia de los operadores de antorcha. A continuación, se muestra la salida de muestra del generador de perfiles de antorcha:

# Torch profiler output for the eager mode run on c7g.xl (4vcpu)
---------------    ------------  -----------  ------------  -----------  ------------  ------------
Name                 Self CPU %   Self CPU     CPU total %   CPU total   CPU time avg    # of Calls
---------------    ------------  -----------  ------------  -----------  ------------  ------------
aten::mm            40.71%         12.502ms       40.71%      12.502ms     130.229us            96
model_inference     26.44%         8.118ms       100.00%      30.708ms      30.708ms             1
aten::bmm            6.85%         2.102ms         9.47%       2.908ms      80.778us            36
aten::matmul         3.73%         1.146ms        57.26%      17.583ms     133.205us           132
aten::select         1.88%       576.000us         1.90%     583.000us       0.998us           584
aten::transpose      1.51%       464.000us         1.83%     563.000us       3.027us           186
---------------    ------------  -----------  ------------  -----------  ------------  -------------
Self CPU time total: 30.708ms

# Torch profiler output for the compile mode run for the same model on the same instance
------------------------- ----------  -----------  ------------  ------------  ------------  ------------
Name                      Self CPU %    Self CPU    CPU total %    CPU total   CPU time avg   # of Calls
------------------------- ----------  -----------  ------------  ------------  ------------  ------------
mkldnn::_linear_pointwise   37.98%       5.461ms        45.91%       6.602ms      68.771us            96
Torch-Compiled Region       29.56%       4.251ms        98.53%      14.168ms      14.168ms             1
aten::bmm                   14.90%       2.143ms        21.73%       3.124ms      86.778us            36
aten::select                 4.51%     648.000us         4.62%     665.000us       1.155us           576
aten::view                   3.29%     473.000us         3.29%     473.000us       1.642us           288
aten::empty                  2.53%     364.000us         2.53%     364.000us       3.165us           115
-------------------------  ---------  -----------  ------------  ------------  ------------ -------------
Self CPU time total: 14.379ms

Que sigue

A continuación, ampliaremos el soporte del backend de la CPU del inductor de antorcha para compilar el modelo Llama y agregaremos soporte para kernels GEMM fusionados para permitir la optimización de la fusión del operador del inductor de antorcha en los procesadores AWS Graviton3.

Conclusión

En este tutorial, explicamos cómo optimizamos el rendimiento de Torch.Compile en instancias EC2 basadas en AWS Graviton3, cómo usar las optimizaciones para mejorar el rendimiento de la inferencia del modelo de PyTorch y demostramos las aceleraciones resultantes. ¡Esperamos que lo pruebes! Si necesitas ayuda con el software de aprendizaje automático en Graviton, abre un problema en la Guía técnica de AWS Graviton GitHub.

Sobre el Autor

Sunita Nadampalli es gerente de desarrollo de software y experta en IA/ML en AWS. Dirige las optimizaciones de rendimiento del software AWS Graviton para cargas de trabajo de IA/ML y HPC. Le apasiona el desarrollo de software de código abierto y la entrega de soluciones de software sustentables y de alto rendimiento para SoC basados ​​en Arm ISA.