A AI21 Labs treina modelo de linguagem com 178 bilhões de parâmetros usando instâncias P4d do Amazon EC2, PyTorch

A AI21 Labs usa machine learning para desenvolver modelos de linguagem focados na compreensão do significado e, em 2021, estabeleceu a meta de treinar o recém-lançado Jurassic-1 Jumbo, um modelo de linguagem autorregressivo com 178 bilhões de parâmetros. Os desenvolvedores que se inscreverem para o teste beta terão acesso ao Jurassic-1 Jumbo e poderão começar imediatamente a personalizar o modelo para seu caso de uso. A startup de software queria treinar o modelo com eficiência, então procurou a Amazon Web Services (AWS) e criou uma solução usando o Amazon Elastic Compute Cloud (Amazon EC2), um serviço web que fornece capacidade de computação segura e redimensionável na nuvem. A escolha do Amazon EC2 deu à empresa o controle sobre o processo de treinamento, incluindo a alocação de nós.

Para uma poderosa funcionalidade de computação e de rede, a empresa selecionou as instâncias P4d do Amazon EC2, que oferecem rede de alto throughput e baixa latência para treinamento de machine learning e aplicações de computação de alto desempenho na nuvem. Usando instâncias P4d do Amazon EC2, a AI21 Labs obteve o desempenho e a memória necessários ao distribuir o treinamento do modelo em centenas de GPUs para fornecer processamento de linguagem natural como um serviço por meio do modelo Jurassic-1 Jumbo. Como a empresa agora treina e controla seu próprio modelo em grande escala, ela pode trabalhar para desenvolver novos modelos na mesma escala e inovar com maior facilidade.

Fundada em 2017, a AI21 Labs busca uma missão híbrida: conduzir pesquisas sobre processamento de linguagem natural e desenvolver produtos baseados em inteligência artificial para leitura e escrita. Seu principal produto, o Wordtune, é um assistente inteligente de redação e edição lançado em outubro de 2020 e cresceu para oferecer suporte a quase um milhão de usuários. Seu outro produto principal, o AI21 Studio, oferece acesso por API aos modelos de linguagem Jurassic-1 da empresa, bem como desenvolvimento de modelos personalizados. “Somos parte de um pequeno grupo de empresas que estão oferecendo modelos de linguagem como serviço, capacitando desde de desenvolvedores independentes a empresas multinacionais a criar aplicações e serviços com base na avançada tecnologia de processamento de linguagem natural”, diz Yoav Shoham, cofundador e coCEO da AI21 Labs. “Além disso, estamos buscando inovações científicas e enfrentando os desafios de engenharia de software impostos por modelos desse tamanho e complexidade.”

Para treinar seu primeiro megamodelo de aprendizado profundo com eficiência e atender às necessidades de alta escala e desempenho do modelo, a AI21 Labs precisava de uma computação poderosa, velocidade de rede eficiente e acesso a suporte técnico e orientação. Por esses motivos, no início de 2021, a empresa começou a implementar uma solução na AWS, optando por treinar o modelo usando instâncias P4d do Amazon EC2. Essas instâncias são implantadas em clusters de hiperescala chamados Amazon EC2 UltraClusters e oferecem mais de 4 mil GPUs do NVIDIA A100, infraestrutura de rede sem bloqueio em escala de petabits e alto throughput, além de armazenamento de baixa latência. 

A abordagem da empresa foi ainda mais otimizada em relação ao GPUDirectRDMA de baixa latência e alta largura de banda, juntamente com o Elastic Fabric Adapter (EFA), uma interface de rede para instâncias do Amazon EC2 que permite que os clientes executem aplicações que exigem altos níveis de comunicação entre nós em grande escala na AWS. Devido ao tamanho do modelo, a equipe precisava usar o processamento paralelo para obter um tempo de treinamento eficiente. Por isso, utilizou os recursos de rede na AWS para apoiar seu treinamento distribuído e o paralelismo de modelos. “As instâncias P4d do Amazon EC2 oferecem rede de alto desempenho de 400 Gbps no EFA”, diz Opher Lieber, líder técnico do Jurassic na AI21 Labs. “A velocidade da rede de GPU a GPU afeta diretamente a capacidade de escalar com eficiência e permanecer econômica ao escalar para centenas de GPUs.”

A AI21 Labs começou criando sua base de código nas instâncias P4d do Amazon EC2 ativadas para o EFA. Em seguida, testou e verificou o desempenho e a escalabilidade eficiente de sua abordagem de treinamento multinó. Em seguida, a equipe lançou um treinamento rápido do modelo em tamanho real, que usa centenas de GPUs, para verificar a função e o desempenho. A partir daí, ela pôde começar a treinar seu modelo Jurassic-1 Jumbo na AWS. Para orquestração, a empresa escolheu uma solução interna que aloca instâncias usando um kit de desenvolvimento de software da AWS — o AWS SDK para Python (Boto3), que facilita a integração da aplicação, biblioteca ou script Python de um cliente com vários serviços da AWS.

Para armazenamento, a AI21 Labs escolheu o Amazon Simple Storage Service (Amazon S3), que oferece escalabilidade, disponibilidade de dados, segurança e desempenho líderes do setor. “Conseguimos alcançar um desempenho muito bom no Amazon S3 usando a ajuda da equipe da AWS. Por isso, foi uma escolha fácil por causa do desempenho e do preço”, diz Lieber. A equipe usa buckets do Amazon S3 para armazenar e carregar pontos de verificação de forma eficiente e distribuída. Para registrar o progresso e os eventos do treinamento, a equipe usa o Amazon CloudWatch, um serviço de monitoramento e observabilidade. 

Ao implementar sua solução, a AI21 Labs aproveitou o suporte da AWS. Sua equipe consultou especialistas da AWS que forneceram orientação sobre questões e preocupações relacionadas a nível de serviço, arquitetura e hardware. Além disso, a empresa melhorou o desempenho do Jurassic-1 Jumbo usando o PyTorch na AWS, uma estrutura de aprendizado profundo de código aberto que facilita o desenvolvimento de modelos de machine learning e a implantação deles na produção. 

A AI21 Labs concluiu o treinamento ao longo de vários meses, finalizando em junho de 2021. O novo megamodelo, um modelo de linguagem autorregressiva, possui 178 bilhões de parâmetros, o que é comparável à oferta da concorrente da empresa. Ele também oferece um vocabulário diferenciado de 256 mil itens que fornece recursos expandidos de representação de texto, bem como suporte para entidades nomeadas. A empresa agora oferece o Jurassic-1 Jumbo (junto com seu homólogo, o Jurassic-1 Large, que possui 7 bilhões de parâmetros) em versão beta aberta por meio da oferta do AI21 Studio da empresa. Usando o serviço, uma grande variedade de desenvolvedores pode criar produtos no modelo Jurassic-1 Jumbo, e a AI21 Labs já foi adotado em muitos setores, incluindo marketing, criação de conteúdo, jogos, pesquisa médica, automotivo, telecomunicações e finanças.

Como a AI21 Labs possui e tem acesso direto ao seu modelo, ela pode se adaptar e inovar sem depender de terceiros e pode explorar metas de inovação contínuas, que são uma parte fundamental de sua missão. Atualmente, a AI21 Labs está prototipando modelos adicionais, que também planeja treinar em grande escala. “Treinar e possuir nossos próprios megamodelos continuará sendo um fator crucial de diferenciação em nossas ofertas do Wordtune e do AI21 Studio”, diz Shoham.