Une approche concrète pour accélérer votre inférence BERT avec ONNX-Torchscript

Les progrès des LLMs s'accompagnent d'un coût de calcul élevé, la plupart des modèles basés sur les transformateurs sont massifs et le nombre de paramètres ainsi que les données utilisées pour l'apprentissage augmentent constamment. Alors que le modèle BERT original comptait déjà 110 millions de paramètres, le dernier GPT-3 en compte 175 milliards, soit une augmentation stupéfiante de ~1700x en deux ans de recherche.

Ces modèles massifs nécessitent généralement des centaines de GPU et plusieurs jours d'entraînement pour être efficaces. Heureusement, grâce à l'apprentissage par transfert (transfer learning), nous pouvons télécharger des modèles pré-entraînés et les affiner rapidement sur nos propres ensembles de données beaucoup plus petits pour un coût modique.

Ceci étant dit, une fois l'entraînement terminé, vous avez toujours un modèle massif sur les bras que vous pouvez vouloir déployer en production. L'inférence prend un temps relativement long par rapport à des modèles plus modestes et peut être trop lente pour atteindre le débit dont vous avez besoin.

Bien que vous puissiez investir dans du matériel plus rapide ou utiliser plus de serveurs pour faire le travail, il existe différentes façons de réduire le temps d'inférence de votre modèle :

• L'élagage du modèle : Réduire le nombre de couches, la dimension des embeddings ou le nombre d'unités dans les couches cachées.
• Quantification : Au lieu d'utiliser des flottants 32 bits (FP32) pour les poids, utilisez la demi-précision (FP16) ou même des entiers 8 bits.
• Exportation d'un modèle de Pytorch/Tensorflow natif vers un format ou un moteur d'inférence approprié (Torchscript/ONNX/TensorRT...).
• Batching : Prédire sur un lot d'échantillons plutôt que sur des échantillons individuels.

La première et la deuxième approche impliquent généralement un réentraînement de votre modèle, tandis que les deux dernières approches sont effectuées après l'entraînement et sont essentiellement indépendantes de votre tâche particulière.

Si la vitesse d'inférence est extrêmement importante pour votre cas d'utilisation, vous devrez très probablement expérimenter toutes ces méthodes pour produire un modèle fiable et extrêmement rapide. Dans la plupart des cas cependant, l'exportation de votre modèle vers un format/framework approprié et la prédiction sur des lots vous donnera des résultats beaucoup plus rapides pour une quantité de travail minimale. Nous nous concentrerons ici sur cette approche pour voir l'impact qu'elle peut avoir sur le débit de notre modèle.

Nous allons explorer les effets de la modification du format du modèle et du batching avec quelques expériences :

• Baseline avec vanilla Pytorch CPU/GPU
• Exportation du modèle Pytorch vers Torchscript CPU/GPU
• Modèle Pytorch vers ONNX CPU/GPU
• Toutes les expériences sont effectuées sur des lots de 1/2/4/8/16/32/64 échantillons.

A ce jour, il n'est pas encore possible d'exporter directement un modèle transformer de Pytorch vers TensorRT en raison du manque de support de int64 utilisé par les embeddings de Pytorch, nous allons donc l'ignorer pour le moment.

Nous allons effectuer la classification de phrases sur camemBERT (~100M paramètres), une variante française de Roberta. Comme la grande majorité des calculs sont effectués dans le modèle de transformation, vous devriez obtenir des résultats similaires quelle que soit votre tâche.

Tout d'abord, nous allons jeter un coup d'oeil rapide sur la façon d'exporter un modèle Pytorch vers le format/framework approprié, si vous ne voulez pas lire de code, vous pouvez passer à la section des résultats plus bas.

Comment exporter votre modèle

Vanilla Pytorch

Sauvegarder et charger un modèle dans Pytorch est assez simple, bien qu'il y ait différentes façons de procéder. Pour l'inférence, la documentation officielle recommande de sauvegarder le 'state_dict' de votre modèle, qui est un dictionnaire python contenant les paramètres apprenables de votre modèle. Cette méthode est plus légère et plus robuste que le pickling de votre modèle entier.

#Savingmodel = SequenceClassifier()train_model(model)torch.save(model.state_dict(), 'pytorch_model.pt')#Loadingmodel = SequenceClassifier()model.load_state_dict(torch.load(PATH))model.eval() #Set dropout and batch normalization layers to evaluation modeavec torch.go_grad() :   logits = modèle(**batch_x)

Suivez ce lien pour des informations supplémentaires sur la sauvegarde/le chargement sur Pytorch.

Torchscript JIT

TorchScript est un moyen de créer des modèles sérialisables et optimisables à partir de votre code Pytorch. Une fois exporté vers Torchscript, votre modèle sera exécutable depuis Python et C++.

• Trace : Une entrée est envoyée à travers le modèle et toutes les opérations sont enregistrées dans un graphe qui définira votre modèle torchscript.
• Script : Si votre modèle est plus complexe et possède un flux de contrôle tel que des instructions conditionnelles, le scripting inspectera le code source du modèle et le compilera en tant que code TorchScript.

Notez que puisque votre modèle sera sérialisé, vous ne serez pas en mesure de le modifier après qu'il ait été sauvegardé, donc vous devriez le mettre en mode évaluation et l'exporter sur le dispositif approprié avant de le sauvegarder.

Si vous voulez faire de l'inférence à la fois sur le CPU et le GPU, vous devez sauvegarder 2 modèles différents.

#savingjit_sample = (batch_x['input_ids'].int().to(device), batch_x['attention_mask'].int().to(device))

model.eval()model.to(device)module = torch.jit.trace(model, jit_sample)torch.jit.save('model_jit.pt')#loadingmodel = torch.jit.load('model_jit.pt', map_location=torch.device(device))logits = modèle(**batch_x)

Pour une introduction plus complète, vous pouvez suivre le [tutoriel] officiel (https://pytorch.org/tutorials/beginner/Intro_to_TorchScript_tutorial.html).

ONNX

ONNX fournit un format open source pour les modèles d'IA, la plupart des frameworks peuvent exporter leur modèle au format ONNX. En plus de l'interopérabilité entre les frameworks, ONNX est livré avec quelques optimisations qui devraient accélérer l'inférence.

L'exportation vers ONNX est légèrement plus compliquée mais Pytorch fournit une fonction d'exportation directe, vous devez seulement fournir quelques informations clés.

• opset_version, pour chaque version il y a un ensemble d'opérateurs qui sont supportés, certains modèles avec des architectures plus exotiques peuvent ne pas être exportables pour le moment.
• input_names et output_names sont les noms à attribuer aux noeuds d'entrée et de sortie du graphe.
• L'argument axes_dynamiques est un dictionnaire qui indique quelle dimension de vos variables d'entrée et de sortie peut changer, par exemple le batch_size ou la longueur de la séquence.

.

#savinginput_x = jit_sample ## taking sample from previous exampletorch.onnx.export(model, input_x, 'model_onnx.pt',export_params=True, opset_version=11, do_constant_folding=True, input_names = ['input_ids', 'attention_mask'], output_names = ['output'],dynamic_axes= {'input_ids' : {0 : 'batch_size', 1:'length'}, 'attention_mask' : {0 : 'batch_size', 1:'length'},'output' : {0 : 'batch_size'}})#chargementmodel = onnxruntime.InferenceSession(model_onnx)batch_x = {'input_ids':sample['input_ids'].cpu().numpy(),"attention_mask":sample['attention_mask'].cpu().numpy()}logits = model.run(None, batch_x)

Le runtime ONNX peut être utilisé avec un GPU bien qu'il nécessite des versions spécifiques de CUDA, cuDNN et du système d'exploitation, ce qui rend le processus d'installation difficile au début.

Pour un tutoriel plus complet, vous pouvez suivre la [documentation] officielle (https://pytorch.org/tutorials/advanced/super_resolution_with_onnxruntime.html).

Résultats expérimentaux

Chaque configuration a été exécutée 5 fois sur un jeu de données de 1 000 phrases de différentes longueurs. Nous avons testé 2 différents GPU populaires : T4 et V100 avec torch 1.7.1 et ONNX 1.6.0. Gardez à l'esprit que les résultats varieront en fonction de votre matériel spécifique, des versions des paquets et du jeu de données.

Le temps d'inférence varie d'environ 50 ms par échantillon en moyenne à 0,6 ms sur notre jeu de données, selon la configuration matérielle.

Sur le CPU, le format ONNX est clairement gagnant pour une taille de lot <32, à partir de laquelle le format semble ne plus avoir d'importance. Si nous prédisons échantillon par échantillon, nous voyons que ONNX parvient à être aussi rapide que l'inférence sur notre base de référence sur GPU pour une fraction du coût.

Comme prévu, l'inférence est beaucoup plus rapide sur GPU, surtout avec une taille de lot plus élevée. Nous pouvons également voir que la taille idéale des lots dépend du GPU utilisé :

• Pour le T4, la meilleure configuration est d'exécuter ONNX avec des lots de 8 échantillons, ce qui donne un gain de vitesse de ~12x par rapport à la taille de lot 1 sur pytorch.
• Pour le V100, avec des lots de 32 ou 64 échantillons, nous pouvons atteindre une accélération de ~28x par rapport à la ligne de base pour le GPU et de ~90x pour la ligne de base sur le CPU.

Dans l'ensemble, nous constatons que le choix d'un format approprié a un impact significatif pour les lots de petite taille, mais que cet impact se réduit au fur et à mesure que les lots augmentent, pour des lots de 64 échantillons, les 3 configurations sont à ~10% les unes des autres.

Impact de la longueur de la séquence et de la stratégie de mise en lots

Un autre élément à prendre en compte est la longueur de la séquence. Les transformateurs sont généralement limités à des séquences de 512 tokens, mais il y a une différence massive dans la vitesse et la mémoire requise pour différentes longueurs de séquences dans cette gamme.

Le temps d'inférence augmente de façon linéaire avec la longueur de la séquence pour les lots plus importants, mais pas pour les échantillons individuels. Cela signifie que si vos données sont constituées de longues séquences de texte (articles de presse par exemple), vous n'obtiendrez pas une accélération aussi importante avec les lots. Comme toujours, cela dépend de votre matériel, un V100 est plus rapide qu'un T4 et ne souffrira pas autant lors de la prédiction de longues séquences alors que d'un autre côté, notre CPU est complètement débordé :

Pour vérifier rapidement, regardons ce qui se passe lorsque nous trions notre ensemble de données avant d'exécuter l'inférence :

Comme nous nous y attendions, pour les lots de grande taille, il y a une incitation significative à regrouper les échantillons de longueur similaire. Pour les données non triées, plus les lots sont grands, plus il y a de chances de se retrouver avec des échantillons plus longs qui augmenteront de manière significative le temps d'inférence de l'ensemble du lot. Nous pouvons voir que passer de 16 à 64 tailles de lot ralentit l'inférence de 20% alors qu'elle devient 10% plus rapide avec des données triées.

Cette stratégie peut également être utilisée pour réduire de manière significative votre temps d'apprentissage, cependant cela doit être fait avec prudence car cela peut avoir un impact négatif sur la performance de votre modèle, surtout s'il y a une certaine corrélation entre vos étiquettes et la longueur de vos échantillons.

Prochaines étapes

Bien que ces expériences aient été exécutées directement en Python, les modèles Torchscript et ONNX peuvent être chargés directement en C++, ce qui pourrait fournir une augmentation supplémentaire de la vitesse d'inférence.

Si votre modèle est encore trop lent pour votre cas d'utilisation, Pytorch fournit différentes options pour la quantification. La "quantification dynamique" peut être effectuée après l'apprentissage mais aura probablement un impact sur la précision de votre modèle, tandis que la "quantification consciente" nécessite un ré-entraînement mais devrait avoir moins d'impact sur la performance de votre modèle.

Conclusion

Comme nous l'avons vu, il n'y a pas de réponse simple pour optimiser votre temps d'inférence, car cela dépend principalement de votre matériel spécifique et du problème que vous essayez de résoudre. Par conséquent, vous devez réaliser vos expériences avec votre propre matériel et vos propres données pour obtenir des résultats fiables.

Néanmoins, il y a quelques lignes directrices qui devraient être vraies et qui sont faciles à mettre en œuvre :

• La prédiction sur des lots peut fournir une accélération significative jusqu'à une certaine taille (en fonction de votre matériel spécifique), en particulier si vous pouvez regrouper des échantillons de longueur similaire ensemble.
• L'utilisation de Torchscript ou ONNX fournit une accélération significative pour des lots de taille et de longueur de séquence inférieures, l'effet est particulièrement fort lors de l'exécution de l'inférence sur des échantillons individuels.
• ONNX semble être la plus performante des trois configurations que nous avons testées, bien qu'elle soit aussi la plus difficile à installer pour l'inférence sur GPU.
• Torchscript fournit une accélération fiable pour les lots de petite taille et est très facile à installer.