Optimisation de l’inférence de modèle PyTorch sur AWS Graviton
Les modèles AI/ML peuvent être une entreprise extrêmement coûteuse. Bon nombre de nos articles se sont concentrés sur une grande variété de conseils, astuces et techniques pour analyser et optimiser les performances d’exécution des charges de travail IA/ML. Notre argument est double :
- L’analyse et l’optimisation des performances doivent faire partie intégrante de chaque projet de développement d’IA/ML et,
- Atteindre des améliorations significatives des performances et une réduction des coûts ne nécessite pas un haut degré de spécialisation. N’importe quel développeur AI/ML peut le faire. Chaque développeur AI/ML devrait le faire.
nous avons relevé le défi de l’optimisation d’une charge de travail d’inférence ML sur un processeur Intel® Xeon®. Nous avons commencé par examiner un certain nombre de scénarios dans lesquels un processeur pourrait être le meilleur choix pour l’inférence IA/ML, même à l’ère de plusieurs puces d’inférence IA dédiées. Nous avons ensuite introduit un modèle PyTorch de classification d’images de jouets et avons procédé à la démonstration d’un grand nombre de techniques permettant d’améliorer ses performances d’exécution sur un ordinateur. Instance Amazon EC2 c7i.xlarge, alimentée par des processeurs Intel Xeon Scalable de 4e génération. Dans cet article, nous étendons notre discussion aux processeurs Graviton basés sur Arm d’AWS. Nous reviendrons sur bon nombre des optimisations dont nous avons discuté dans nos articles précédents – dont certaines nécessiteront une adaptation au processeur Arm – et évaluerons leur impact sur le même modèle de jouet. Compte tenu des profondes différences entre les processeurs Arm et Intel, les chemins vers la configuration la plus performante peuvent emprunter des chemins différents.
AWS Graviton
AWS Graviton est une famille de processeurs basés sur Bras Neoverse Processeurs conçus et construits sur mesure par AWS pour un rapport qualité-prix et une efficacité énergétique optimaux. Leurs moteurs dédiés au traitement vectoriel (NÉON et SVE/SVE2) et multiplication matricielle (MMLA), et leur soutien à Bfloat16 opérations (à partir de Graviton3), en font un candidat incontournable pour exécuter des charges de travail intensives en calcul telles que Inférence IA/ML. Pour faciliter une IA/ML haute performance sur Graviton, l’ensemble de la pile logicielle a été optimisée pour son utilisation :
- Noyaux de calcul de bas niveau de la Bibliothèque de calcul Arm (ACL) sont hautement optimisés pour tirer parti des accélérateurs matériels Graviton (par exemple, SVE et MMLA).
- Bibliothèques de middleware ML tel que unDNN et OuvertBLAS acheminez les opérations d’apprentissage profond et d’algèbre linéaire vers les noyaux ACL spécialisés.
- Cadres IA/ML comme PyTorch et TensorFlow sont compilés et configurés pour utiliser ces backends optimisés.
Dans cet article, nous utiliserons un Amazon EC2 c8g.xlarge instance alimentée par quatre AWSGraviton4 processeurs et un AMI d’apprentissage profond AWS ARM64 PyTorch (DLAMI).
L’intention de cet article est de montrer des conseils pour améliorer les performances sur une instance AWS Graviton. Surtout, notre intention est pas il ne s’agit pas non plus de faire une comparaison entre AWS Graviton et les puces alternatives, ni de préconiser l’utilisation d’une puce plutôt que de l’autre. Le meilleur choix de processeur dépend de tout un tas de considérations qui dépassent le cadre de cet article. L’une des considérations importantes sera les performances d’exécution maximales de votre modèle sur chaque puce. En d’autres termes : combien pouvons-nous en avoir pour notre argent ? Ainsi, prendre une décision éclairée sur le meilleur processeur est l’une des motivations pour optimiser les performances d’exécution de chacun.
Une autre motivation pour optimiser les performances de notre modèle pour plusieurs dispositifs d’inférence est d’augmenter sa portabilité. Le terrain de jeu de l’IA/ML est extrêmement dynamique et la résilience aux circonstances changeantes est cruciale pour réussir. Il n’est pas rare que des instances de calcul de certains types deviennent soudainement indisponibles ou rares. À l’inverse, une augmentation de la capacité des instances AWS Graviton pourrait impliquer leur disponibilité à des prix très réduits, par exemple dans le Instance ponctuelle Amazon EC2 marché, présentant des opportunités d’économies que vous ne voudriez pas manquer.
Avis de non-responsabilité
Les blocs de code que nous partagerons, les étapes d’optimisation dont nous discuterons et les résultats que nous obtiendrons sont destinés à servir d’exemple des avantages que vous pouvez constater grâce à l’optimisation des performances du ML sur une instance AWS Graviton. Ceux-ci peuvent différer considérablement des résultats que vous pourriez obtenir avec votre propre modèle et votre propre environnement d’exécution. Veuillez ne pas vous fier à l’exactitude ou à l’optimalité du contenu de cet article. Veuillez ne pas interpréter la mention d’une bibliothèque, d’un framework ou d’une plateforme comme une approbation de son utilisation.
Optimisation de l’inférence sur AWS Graviton
Comme dans notre article précédent, nous démontrerons les étapes d’optimisation sur un modèle de classification d’images de jouets :
import torch, torchvision
import time
def get_model(channels_last=False, compile=False):
model = torchvision.models.resnet50()
if channels_last:
model= model.to(memory_format=torch.channels_last)
model = model.eval()
if compile:
model = torch.compile(model)
return model
def get_input(batch_size, channels_last=False):
batch = torch.randn(batch_size, 3, 224, 224)
if channels_last:
batch = batch.to(memory_format=torch.channels_last)
return batch
def get_inference_fn(model, enable_amp=False):
def infer_fn(batch):
with torch.inference_mode(), torch.amp.autocast(
'cpu',
dtype=torch.bfloat16,
enabled=enable_amp
):
output = model(batch)
return output
return infer_fn
def benchmark(infer_fn, batch):
# warm-up
for _ in range(20):
_ = infer_fn(batch)
iters = 100
start = time.time()
for _ in range(iters):
_ = infer_fn(batch)
end = time.time()
return (end - start) / iters
batch_size = 1
model = get_model()
batch = get_input(batch_size)
infer_fn = get_inference_fn(model)
avg_time = benchmark(infer_fn, batch)
print(f"\nAverage samples per second: {(batch_size/avg_time):.2f}")Le débit initial est de 12 échantillons par seconde (SPS).
Mise à niveau vers la version la plus récente de PyTorch
Alors que la version de PyTorch dans notre DLAMI est la 2.8, la dernière version de PyTorch, au moment d’écrire ces lignes, est la 2.9. Compte tenu du rythme rapide de développement dans le domaine de l’IA/ML, il est fortement recommandé d’utiliser les packages de bibliothèques les plus récents. Dans un premier temps, nous passons à PyTorch 2.9 lequel inclut des mises à jour clés de son backend Arm.
pip3 install -U torch torchvision --index-url https://download.pytorch.org/whl/cpuDans le cas de notre modèle dans sa configuration initiale, la mise à jour de la version PyTorch n’a aucun effet. Cependant, cette étape est cruciale pour tirer le meilleur parti des techniques d’optimisation que nous évaluerons.
Inférence par lots
Pour réduire les frais généraux de lancement et augmenter l’utilisation des accélérateurs matériels, nous regroupons les échantillons et appliquons l’inférence par lots. Le tableau ci-dessous montre comment le débit du modèle varie en fonction de la taille du lot :
Optimisations de la mémoire
Nous appliquons un certain nombre de techniques de notre article précédent pour optimiser l’allocation et l’utilisation de la mémoire. Ceux-ci incluent le format de mémoire des canaux, précision mixte automatique avec le type de données bfloat16 (pris en charge par Graviton3), le TCMalloc bibliothèque d’allocation et allocation de pages énorme. Veuillez consulter le pour plus de détails. Nous permettons également au mathématiques rapides mode des noyaux ACL GEMM et mise en cache des primitives du noyau — deux optimisations qui apparaissent dans le directives pour exécuter l’inférence PyTorch sur Graviton.
Les instructions de ligne de commande requises pour activer ces optimisations sont présentées ci-dessous :
# install TCMalloc
sudo apt-get install google-perftools
# Program the use of TCMalloc
export LD_PRELOAD=/usr/lib/aarch64-linux-gnu/libtcmalloc.so.4
# Enable huge page memory allocation
export THP_MEM_ALLOC_ENABLE=1
# Enable the fast math mode of the GEMM kernels
export DNNL_DEFAULT_FPMATH_MODE=BF16
# Set LRU Cache capacity to cache the kernel primitives
export LRU_CACHE_CAPACITY=1024Le tableau suivant présente l’impact des optimisations de mémoire, appliquées successivement :
Dans le cas de notre modèle de jouet, le dernières chaînes et bfloat16-précision mixte les optimisations ont eu le plus grand impact. Après avoir appliqué toutes les optimisations de mémoire, le débit moyen est de 53,03 SPS.
Compilation de modèles
La prise en charge de la compilation PyTorch pour AWS Graviton est un domaine d’effort concentré de l’équipe AWS Graviton. Cependant, dans le cas de notre modèle jouet, cela entraîne une légère réduction du débit, de 53,03 SPS à 52,23.
Inférence multi-travailleurs
Bien qu’elle soit généralement appliquée dans des paramètres comportant plus de quatre processeurs virtuels, nous démontrons la mise en œuvre de l’inférence multi-travailleurs en modifiant notre script pour prendre en charge l’épinglage de base :
if __name__ == '__main__':
# pin CPUs according to worker rank
import os, psutil
rank = int(os.environ.get('RANK','0'))
world_size = int(os.environ.get('WORLD_SIZE','1'))
cores = list(range(psutil.cpu_count(logical=True)))
num_cores = len(cores)
cores_per_process = num_cores // world_size
start_index = rank * cores_per_process
end_index = (rank + 1) * cores_per_process
pid = os.getpid()
p = psutil.Process(pid)
p.cpu_affinity(cores[start_index:end_index])
batch_size = 8
model = get_model(channels_last=True)
batch = get_input(batch_size, channels_last=True)
infer_fn = get_inference_fn(model, enable_amp=True)
avg_time = benchmark(infer_fn, batch)
print(f"\nAverage samples per second: {(batch_size/avg_time):.2f}")Nous notons que contrairement aux autres types d’instances de CPU AWS EC2, chaque vCPU Graviton est directement mappé à un seul cœur de CPU physique. Nous utilisons le course aux flambeaux utilitaire pour démarrer quatre travailleurs, chacun fonctionnant sur un seul cœur de processeur :
export OMP_NUM_THREADS=1 #set one OpenMP thread per worker
torchrun --nproc_per_node=4 main.pyCela se traduit par un débit de 55,15 SPS, soit une amélioration de 4 % par rapport à notre meilleur résultat précédent.
Quantification INT8 pour Arm
Un autre domaine de développement actif et d’amélioration continue sur Arm est la quantification INT8. Les outils de quantification INT8 sont généralement fortement liés au type d’instance cible. Dans notre article précédent, nous avons démontré Quantification d’exportation PyTorch 2 avec backend X86 via Inductor en utilisant le TorchAO (0.12.1) bibliothèque. Heureusement, les versions récentes de TorchAO inclure un dédié quantificateur pour Arm. La séquence de quantification mise à jour est présentée ci-dessous. Comme dans notre article précédent, nous nous intéressons uniquement à l’impact potentiel sur les performances. En pratique, la quantification INT8 peut avoir un impact significatif sur la qualité du modèle et nécessiter une stratégie de quantification plus sophistiquée.
from torchao.quantization.pt2e.quantize_pt2e import prepare_pt2e, convert_pt2e
import torchao.quantization.pt2e.quantizer.arm_inductor_quantizer as aiq
def quantize_model(model):
x = torch.randn(4, 3, 224, 224).contiguous(
memory_format=torch.channels_last)
example_inputs = (x,)
batch_dim = torch.export.Dim("batch")
with torch.no_grad():
exported_model = torch.export.export(
model,
example_inputs,
dynamic_shapes=((batch_dim,
torch.export.Dim.STATIC,
torch.export.Dim.STATIC,
torch.export.Dim.STATIC),
)
).module()
quantizer = aiq.ArmInductorQuantizer()
quantizer.set_global(aiq.get_default_arm_inductor_quantization_config())
prepared_model = prepare_pt2e(exported_model, quantizer)
prepared_model(*example_inputs)
converted_model = convert_pt2e(prepared_model)
optimized_model = torch.compile(converted_model)
return optimized_model
batch_size = 8
model = get_model(channels_last=True)
model = quantize_model(model)
batch = get_input(batch_size, channels_last=True)
infer_fn = get_inference_fn(model, enable_amp=True)
avg_time = benchmark(infer_fn, batch)
print(f"\nAverage samples per second: {(batch_size/avg_time):.2f}")Le débit résultant est de 56,77 SPS, soit une amélioration de 7,1 % par rapport à la solution bfloat16.
Compilation AOT utilisant ONNX et OpenVINO
Dans notre article précédent, nous avons exploré les techniques de compilation de modèles anticipées (AOT) en utilisant Échange de réseaux neuronaux ouverts (ONNX) et OuvrirVINO. Les deux bibliothèques incluent une prise en charge dédiée à l’exécution sur AWS Graviton (par exemple, voir ici et ici). Les expériences de cette section nécessitent les installations de bibliothèque suivantes :
pip install onnxruntime onnxscript openvino nncfLe bloc de code suivant illustre la compilation et l’exécution du modèle sur Arm à l’aide d’ONNX :
def export_to_onnx(model, onnx_path="resnet50.onnx"):
dummy_input = torch.randn(4, 3, 224, 224)
batch = torch.export.Dim("batch")
torch.onnx.export(
model,
dummy_input,
onnx_path,
input_names=["input"],
output_names=["output"],
dynamic_shapes=((batch,
torch.export.Dim.STATIC,
torch.export.Dim.STATIC,
torch.export.Dim.STATIC),
),
dynamo=True
)
return onnx_path
def onnx_infer_fn(onnx_path):
import onnxruntime as ort
sess = ort.InferenceSession(
onnx_path,
providers=["CPUExecutionProvider"]
)
sess_options = ort.SessionOptions()
sess_options.add_session_config_entry(
"mlas.enable_gemm_fastmath_arm64_bfloat16", "1")
input_name = sess.get_inputs()[0].name
def infer_fn(batch):
result = sess.run(None, {input_name: batch})
return result
return infer_fn
batch_size = 8
model = get_model()
onnx_path = export_to_onnx(model)
batch = get_input(batch_size).numpy()
infer_fn = onnx_infer_fn(onnx_path)
avg_time = benchmark(infer_fn, batch)
print(f"\nAverage samples per second: {(batch_size/avg_time):.2f}")Il convient de noter que le runtime ONNX prend en charge un Fournisseur d’exécution d’ACL pour fonctionner sur Arm, mais cela nécessite une version ONNX personnalisée (au moment d’écrire ces lignes), ce qui sort du cadre de cet article.
Alternativement, nous pouvons compiler le modèle en utilisant OuvrirVINO. Le bloc de code ci-dessous démontre son utilisation, y compris une option pour la quantification INT8 utilisant NNCF:
import openvino as ov
import nncf
def openvino_infer_fn(compiled_model):
def infer_fn(batch):
result = compiled_model([batch])[0]
return result
return infer_fn
class RandomDataset(torch.utils.data.Dataset):
def __len__(self):
return 10000
def __getitem__(self, idx):
return torch.randn(3, 224, 224)
quantize_model = False
batch_size = 8
model = get_model()
calibration_loader = torch.utils.data.DataLoader(RandomDataset())
calibration_dataset = nncf.Dataset(calibration_loader)
if quantize_model:
# quantize PyTorch model
model = nncf.quantize(model, calibration_dataset)
ovm = ov.convert_model(model, example_input=torch.randn(1, 3, 224, 224))
ovm = ov.compile_model(ovm)
batch = get_input(batch_size).numpy()
infer_fn = openvino_infer_fn(ovm)
avg_time = benchmark(infer_fn, batch)
print(f"\nAverage samples per second: {(batch_size/avg_time):.2f}")Dans le cas de notre modèle de jouet, la compilation OpenVINO entraîne une augmentation supplémentaire du débit à 63,48 SPS, mais la quantification NNCF déçoit, aboutissant à seulement 55,18 SPS.
Résultats
Les résultats de nos expériences sont résumés dans le tableau ci-dessous :
Comme dans notre , nous avons refait nos expériences sur un deuxième modèle : un Vision Transformer (ViT) du Tim bibliothèque – pour démontrer comment l’impact des optimisations d’exécution dont nous avons discuté peut varier en fonction des détails du modèle. Les résultats sont capturés ci-dessous :
Résumé
Dans cet article, nous avons examiné un certain nombre de techniques d’optimisation relativement simples et les avons appliquées à deux modèles de jouets PyTorch. Comme les résultats l’ont démontré, l’impact de chaque étape d’optimisation peut varier considérablement en fonction des détails du modèle, et le cheminement vers des performances optimales peut emprunter de nombreux chemins différents. Les étapes que nous avons présentées dans cet article n’étaient qu’un apéritif ; il existe sans aucun doute de nombreuses autres optimisations qui peuvent débloquer des performances encore plus élevées.
En cours de route, nous avons noté les nombreuses bibliothèques AI/ML qui ont introduit un support approfondi pour l’architecture Graviton, ainsi que l’effort apparemment continu de la communauté en matière d’optimisation continue. Les gains de performances que nous avons obtenus, combinés à ce dévouement apparent, prouvent qu’AWS Graviton est fermement dans la « cour des grands » lorsqu’il s’agit d’exécuter des charges de travail d’IA/ML à forte intensité de calcul.
