À l'aide d'instances P4d Amazon EC2 et de PyTorch, AI21 Labs entraîne un modèle de langage comptant pas moins de 178 milliards de paramètres

Grâce au machine learning, AI21 Labs développe des modèles linguistiques axés sur la compréhension du sens. En 2021, elle s'est fixée pour objectif d'entraîner un modèle de langage autorégressif récemment publié, qui comporte 178 milliards de paramètres : le Jurassic-1 Jumbo. Les développeurs qui s'inscrivent aux tests bêta auront accès à Jurassic-1 Jumbo et pourront immédiatement commencer à personnaliser le modèle en fonction de leur cas d'utilisation. La start-up logicielle souhaitait entraîner le modèle de manière efficace. Elle s'est donc tournée vers Amazon Web Services (AWS) et a créé une solution utilisant Amazon Elastic Compute Cloud (Amazon EC2), un service Web qui fournit une capacité de calcul sécurisée et redimensionnable dans le cloud. Le choix d'Amazon EC2 a permis à l'entreprise de contrôler le processus d'entraînement, y compris l'allocation des nœuds.

Pour bénéficier de puissantes fonctionnalités de calcul et de mise en réseau, l'entreprise a sélectionné les instances P4d Amazon EC2. Ces instances fournissent un réseau à haut débit et à faible latence pour la formation au machine learning et les applications de calcul à haute performance dans le cloud. À l'aide des instances P4d Amazon EC2, AI21 Labs a obtenu les performances et la mémoire requises en distribuant l'entraînement des modèles à des centaines de processeurs graphiques. Le but était de fournir un traitement du langage naturel sous forme de service via son modèle Jurassic-1 Jumbo. Comme l'entreprise entraîne et contrôle désormais son propre modèle à grande échelle, elle peut travailler au développement de nouveaux modèles à la même échelle et innover plus facilement.

Fondée en 2017, AI21 Labs poursuit une mission hybride : mener des recherches sur le traitement du langage naturel, et développer des produits basés sur l'intelligence artificielle pour la lecture et l'écriture. Son produit phare, Wordtune, est un assistant d'écriture et d'édition intelligent qui a été lancé en octobre 2020. Il s'est depuis développé, prenant désormais en charge près d'un million d'utilisateurs. Son autre produit principal, AI21 Studio, offre un accès par API aux modèles de langage Jurassic-1 de l'entreprise et permet le développement de modèles personnalisés. « Nous faisons partie d'une petite cohorte d'entreprises qui proposent des modèles linguistiques sous forme de service, permettant ainsi à tous, des développeurs indépendants aux multinationales, de créer des applications et des services s'appuyant sur une technologie avancée de traitement du langage naturel », explique Yoav Shoham, Cofounder and Co-CEO d'AI21 Labs. « En outre, nous poursuivons des innovations scientifiques et relevons les défis d'ingénierie logicielle posés par les modèles de cette taille et de cette complexité. »

Pour entraîner efficacement son premier mégamodèle de deep learning et répondre aux besoins élevés du modèle en matière d'évolutivité et de performances, AI21 Labs avait besoin d'un calcul puissant, d'une mise en réseau rapide ainsi que d'un accès à une assistance technique et à des conseils. Dans cette optique, début 2021, l'entreprise a commencé à mettre en œuvre une solution sur AWS, en choisissant d'entraîner le modèle à l'aide d'instances P4d Amazon EC2. Les instances EC2 sont déployées dans des clusters hyperscale, appelés Amazon EC2 UltraClusters. Ces clusters fournissent plus de 4 000 GPU NVIDIA A100, une infrastructure réseau non bloquante à l'échelle du pétabit ainsi qu'un stockage à haut débit et à faible latence. 

L'approche de l'entreprise a été encore optimisée grâce à la technologie GPUDirectRDMA à faible latence et à haut débit, ainsi qu'à Elastic Fabric Adapter (EFA), une interface réseau pour les instances Amazon EC2. Cette interface permet aux clients d'exécuter des applications nécessitant des niveaux élevés de communications internœuds à grande échelle sur AWS. En raison de la taille du modèle, l'équipe avait besoin d'utiliser le traitement parallèle pour optimiser le temps d'entraînement. Elle s'est donc tournée vers les fonctionnalités réseau d'AWS pour prendre en charge son entraînement distribué et le parallélisme des modèles. « Les instances P4d Amazon EC2 offrent un réseau haute performance de 400 Gbit/s sur EFA », explique Opher Lieber, Jurassic Technical Lead chez AI21 Labs. « La vitesse de mise en réseau GPU à GPU a un impact direct sur la capacité à évoluer efficacement et à rester rentable lorsque l'on passe à des centaines de GPU. »

AI21 Labs a commencé par mettre à jour sa base de code sur les instances P4d Amazon EC2 activées pour EFA. Elle a ensuite testé et vérifié les performances et la mise à l'échelle efficace de son approche d'entraînement multinœuds. Puis, afin de vérifier son fonctionnement et ses performances, l'équipe a lancé un entraînement rapide sur le modèle grandeur nature qui utilise des centaines de GPU. À partir de là, elle a pu commencer à entraîner son modèle Jurassic-1 Jumbo sur AWS. Pour l'orchestration, l'entreprise a choisi une solution interne qui alloue des instances à l'aide d'un kit de développement logiciel AWS : le AWS SDK pour Python (Boto3). Ce kit facilite l'intégration de l'application, de la bibliothèque ou du script Python d'un client à divers services AWS.

Pour le stockage, AI21 Labs a choisi Amazon Simple Storage Service (Amazon S3). Amazon S3 offre une capacité de mise à l'échelle, une disponibilité des données, une sécurité et des performances de premier plan. « Avec l'aide de l'équipe AWS, nous avons pu atteindre de très bonnes performances sur Amazon S3. Le choix s'est donc avéré facile en termes de performances et de prix », explique Lieber. L'équipe utilise les compartiments Amazon S3 pour stocker et charger les points de contrôle de manière efficace et distribuée. Pour enregistrer la progression et les événements de l'entraînement, l'équipe utilise Amazon CloudWatch, un service de surveillance et d'observabilité. 

Lors de la mise en œuvre de sa solution, AI21 Labs a bénéficié du support d'AWS. Son équipe a consulté des spécialistes AWS qui lui ont fourni des conseils sur les questions et préoccupations relatives au niveau de service, à l'architecture et au matériel. En outre, l'entreprise a amélioré les performances de Jurassic-1 Jumbo à l'aide de PyTorch sur AWS. PyTorch est un cadre de deep learning open source qui facilite le développement de modèles de machine learning ainsi que leur déploiement en production. 

AI21 Labs a effectué un entraînement sur plusieurs mois, qui s'est terminé en juin 2021. Le nouveau mégamodèle, un modèle de langage autorégressif, compte 178 milliards de paramètres, ce qui est comparable à l'offre du concurrent de l'entreprise. Il propose également un vocabulaire différencié de 256 000 éléments qui fournit des fonctionnalités de représentation de texte étendues ainsi qu'une prise en charge des entités nommées. La société propose désormais Jurassic-1 Jumbo (ainsi que son homologue Jurassic-1 Large, qui compte 7 milliards de paramètres) en version bêta ouverte via l'offre AI21 Studio de l'entreprise. Grâce à ce service, un large éventail de développeurs peuvent créer des produits sur le modèle Jurassic-1 Jumbo. De plus, AI21 Labs a déjà été adopté dans de nombreux secteurs, dont le marketing, la création de contenu, les jeux vidéo, la recherche médicale, l'automobile, les télécommunications et la finance.

Comme AI21 Labs possède son propre modèle et y a un accès direct, elle peut s'adapter et innover sans dépendre de tiers. Cela signifie qu'elle peut explorer des objectifs d'innovation permanents, ce qui constituent un élément clé de sa mission. AI21 Labs prototype actuellement des modèles supplémentaires, qu'elle prévoit également d'entraîner à grande échelle. « L'entraînement et la possession de nos propres mégamodèles continueront d'être un facteur de différenciation essentiel pour nos offres Wordtune et AI21 Studio », déclare Yoav Shoham.