AI21 Labs entrena un modelo lingüístico de 178.000 millones de parámetros mediante instancias P4d de Amazon EC2 y PyTorch

AI21 Labs utiliza el machine learning para desarrollar modelos de lenguaje centrados en comprender el significado y, en 2021, se fijó el objetivo de entrenar el Jurassic-1 Jumbo, publicado recientemente, un modelo de lenguaje autorregresivo con 178.000 millones de parámetros. Los desarrolladores que se registren para las pruebas beta tendrán acceso a Jurassic-1 Jumbo y podrán empezar inmediatamente a personalizar el modelo para su caso de uso. La empresa emergente de software quería entrenar el modelo de manera eficiente, por lo que recurrió a Amazon Web Services (AWS) y creó una solución con Amazon Elastic Compute Cloud (Amazon EC2), un servicio web que proporciona capacidad informática segura y redimensionable en la nube. La elección de Amazon EC2 permitió a la empresa controlar el proceso de formación, incluida la asignación de nodos.

Para obtener una potente funcionalidad de computación y redes, la empresa seleccionó las instancias P4d de Amazon EC2, que ofrecen redes de alto rendimiento y baja latencia para la formación en machine learning y aplicaciones informáticas de alto rendimiento en la nube. Al utilizar las instancias P4d de Amazon EC2, AI21 Labs obtuvo el rendimiento y la memoria necesarios al distribuir la capacitación sobre modelos en cientos de GPU para ofrecer el procesamiento del lenguaje natural como servicio a través de su modelo Jumbo Jurassic-1. Como la empresa ahora entrena y controla su propio modelo a gran escala, puede trabajar para desarrollar nuevos modelos a la misma escala e innovar con mayor facilidad.

Fundada en 2017, AI21 Labs persigue una misión híbrida: llevar a cabo investigaciones sobre el procesamiento del lenguaje natural y desarrollar productos impulsados por inteligencia artificial para leer y escribir. Su producto estrella, Wordtune, es un asistente inteligente de escritura y edición que se lanzó en octubre de 2020 y que ha crecido hasta dar soporte a casi un millón de usuarios. Su otro producto principal, AI21 Studio, ofrece acceso mediante API a los modelos de lenguaje Jurassic-1 de la empresa, así como el desarrollo de modelos personalizados. «Formamos parte de un pequeño grupo de empresas que ofrecen modelos de lenguaje como servicio, lo que permite a cualquier persona, desde desarrolladores independientes hasta empresas multinacionales, crear aplicaciones y servicios sobre la base de una tecnología avanzada de procesamiento del lenguaje natural», afirma Yoav Shoham, cofundador y codirector ejecutivo de AI21 Labs. «Además, buscamos innovaciones científicas y abordamos los desafíos de ingeniería de software que plantean los modelos de este tamaño y complejidad».

Para entrenar su primer megamodelo de aprendizaje profundo de manera eficiente y satisfacer las necesidades de alto escalamiento y rendimiento del modelo, AI21 Labs necesitaba una computación potente, una velocidad de red eficiente y acceso a soporte técnico y orientación. Por estas razones, a principios de 2021, la empresa comenzó a implementar una solución en AWS y optó por entrenar el modelo mediante instancias P4d de Amazon EC2. Las instancias están implementadas en clústeres a hiperescala, llamados UltraClusters de Amazon EC2 y ofrecen más de 4000 GPU NVIDIA A100, una infraestructura de red sin bloqueo a escala de petabits y almacenamiento de alto rendimiento y baja latencia. 

El enfoque de la empresa se optimizó aún más con GPUDirectRDMA de baja latencia y alto ancho de banda, junto con Elastic Fabric Adapter (EFA), una interfaz de red para instancias de Amazon EC2 que permite a los clientes ejecutar aplicaciones que requieren altos niveles de comunicaciones entre nodos a escala en AWS. Debido al tamaño del modelo, el equipo necesitaba utilizar el procesamiento en paralelo para lograr un tiempo de entrenamiento eficiente, por lo que recurrió a las capacidades de red de AWS para respaldar su capacitación distribuida y su paralelismo de modelos. «Las instancias P4d de Amazon EC2 ofrecen redes de alto rendimiento de 400 Gbps en EFA», afirma Opher Lieber, director técnico de Jurassic en AI21 Labs. «La velocidad de conexión entre GPUs tiene un impacto directo en la capacidad de escalar de manera eficiente y seguir siendo rentable al escalar a cientos de GPU».

AI21 Labs comenzó por mostrar su base de código en las instancias P4d de Amazon EC2 activadas para EFA. Luego probó y verificó el rendimiento y el escalado eficiente de su enfoque de entrenamiento multinodo. A continuación, el equipo inició un entrenamiento rápido sobre el modelo de tamaño completo, que utiliza cientos de GPU, para verificar el funcionamiento y el rendimiento. A partir de ahí, pudo empezar a entrenar su modelo Jurassic-1 Jumbo en AWS. Para la orquestación, la empresa eligió una solución interna que asigna instancias mediante un kit de desarrollo de software de AWS, el AWS SDK para Python (Boto3), que facilita la integración de la aplicación, la biblioteca o el script de Python de un cliente con varios servicios de AWS.

Para el almacenamiento, AI21 Labs escogióAmazon Simple Storage Service (Amazon S3), que ofrece escalabilidad, disponibilidad de datos, seguridad y rendimiento líderes en el sector. «Hemos conseguido un rendimiento muy bueno en Amazon S3 gracias a la ayuda del equipo de AWS, por lo que fue una elección sencilla tanto por el rendimiento como por el precio», afirma Lieber. El equipo utiliza buckets de Amazon S3 para almacenar y cargar los puntos de control de forma eficiente y distribuida. Para registrar el progreso y los eventos del entrenamiento, el equipo utiliza Amazon CloudWatch, un servicio de monitorización y observabilidad. 

Al implementar su solución, AI21 Labs aprovechó el soporte de AWS. Su equipo consultó a especialistas de AWS, quienes proporcionaron orientación sobre preguntas e inquietudes relacionadas con el nivel de servicio, la arquitectura y el hardware. Además, la empresa mejoró el rendimiento de Jurassic-1 Jumbo con PyTorch en AWS, un marco de aprendizaje profundo de código abierto que facilita el desarrollo de modelos de machine learning y su implementación en producción. 

AI21 Labs completó la capacitación a lo largo de varios meses y finalizó en junio de 2021. El nuevo megamodelo, un modelo de lenguaje autorregresivo, tiene 178.000 millones de parámetros, lo que es comparable a la oferta de la empresa de la competencia. También ofrece un vocabulario diferenciado de 256.000 elementos que proporciona capacidades ampliadas de representación de texto, así como soporte para entidades nombradas. La compañía ahora ofrece Jurassic-1 Jumbo (junto con su homólogo, Jurassic-1 Large, que tiene 7.000 millones de parámetros) en beta abierta a través de la oferta AI21 Studio de la compañía. Gracias al servicio, una amplia gama de desarrolladores puede crear productos basados en el modelo Jurassic-1 Jumbo, y AI21 Labs ya lo ha adoptado en muchos sectores, como el marketing, la creación de contenido, los juegos, la investigación médica, la automoción y las telecomunicaciones y las finanzas.

Dado que AI21 Labs posee su modelo y tiene acceso directo a él, puede adaptarse e innovar sin depender de terceros y puede explorar los objetivos de innovación continua, que son una parte clave de su misión. Actualmente, AI21 Labs está creando prototipos de modelos adicionales, que también planea entrenar a gran escala. «Entrenar y tener nuestros propios megamodelos seguirá siendo un factor diferenciador fundamental en nuestras ofertas de Wordtune y AI21 Studio», afirma Shoham.