AI21 Labs acelera a adoção de modelos de IA generativa usando o Amazon SageMaker

Fundada em 2017, a AI21 oferece às empresas acesso a seus modelos de linguagem proprietários com o AI21 Studio, que mais de 30.000 desenvolvedores usam para criar suas próprias aplicações de IA generativa. A empresa também oferece seu assistente de escrita e leitura baseado em IA, o Wordtune, que ajuda dezenas de milhões de usuários em todo o mundo a se envolverem com a linguagem escrita.

Em agosto de 2021, a AI21 lançou seu modelo de linguagem Jurassic-1 em dois tamanhos: o modelo Large é rápido e econômico com 7,5 bilhões de parâmetros, e o modelo Jumbo oferece saída de texto de alta qualidade a um custo maior com 178 bilhões de parâmetros. Embora modelos maiores ofereçam a mais alta qualidade, eles podem ser caros de operar em grande escala e são menos ágeis de operar. Para ajudar seus clientes a otimizar a compensação entre custo e qualidade ao operar em grande escala, a AI21 pré-treinou e lançou seu terceiro modelo, o Grande, com 17 bilhões de parâmetros usando o Amazon SageMaker em dezembro de 2022.

A AI21 concluiu o projeto rapidamente, em menos de dois meses desde o início, levando menos de 20 dias para pré-treinar o modelo. Como os LLMs são grandes redes neurais com bilhões de parâmetros, o treinamento é um projeto desafiador e demorado, que exige grandes recursos computacionais. Usando o Amazon SageMaker, a AI21 experimentou um processo de treinamento de modelos mais simples e eficiente, e a empresa pôde escalar os trabalhos de treinamento distribuídos em quantas GPUs fossem necessárias. “Os arquitetos de soluções da AWS foram responsivos e interativos, e conseguimos solucionar problemas e concluir o projeto a tempo”, diz Dan Padnos, vice-presidente de plataforma da AI21.

A empresa já tinha experiência no uso da AWS e escolheu o Amazon SageMaker porque ele é econômico, simples de usar e totalmente gerenciado. A AI21 também poderia continuar usando sua pilha de software de treinamento existente e começar a funcionar rapidamente, o que era importante enquanto a empresa estava desenvolvendo seus negócios. Para pré-treinar o modelo Grande em menos de 20 dias, a AI21 precisou usar 256 GPUs A100 distribuídas em 32 instâncias. O treinamento em grande escala exigia uma ferramenta que pudesse orquestrar a alocação dos nós, disponibilizar o registro em um local central e reduzir a supervisão manual. “Quando você executa um trabalho de treinamento distribuído dessa escala, todos os tipos de desafios técnicos que podem parecer triviais ou mundanos podem se tornar uma dor de cabeça”, diz Padnos. “O Amazon SageMaker tem recursos que você pode usar para gerenciar essa complexidade e reduzir a quantidade de esforço que sua equipe precisa investir nos detalhes.” Por exemplo, o Amazon SageMaker tem recursos como verificações de saúde e registro central que as empresas podem usar para aumentar a eficiência.

Usando o Amazon SageMaker, a AI21 lançou o novo modelo rapidamente. A empresa estima ter economizado várias semanas em comparação com seus métodos de treinamento anteriores. “Como o Amazon SageMaker lida com falhas de nós, reinicia com elegância e orquestra grandes execuções distribuídas, a equipe que trabalha no pré-treinamento do modelo pode se concentrar nas tarefas principais”, diz Padnos. “Em vez de enfrentar os desafios técnicos, eles podem avaliar a performance do modelo e o progresso do treinamento”.

O cronograma acelerado foi importante porque os recursos do modelo Grande atendem melhor às necessidades da maioria dos clientes da AI21. Clientes com casos de uso para consumidores, como redação automática de e-mails, valorizaram a migração do modelo Jumbo para o modelo Grande porque sua grande escala exige eficiência de custos. Apenas alguns meses após a introdução do modelo Grande, ele representou cerca de dois terços do tráfego da empresa. “Vimos uma adoção rápida e estamos muito satisfeitos com o resultado”, diz Padnos. “Nossa experiência com o Amazon SageMaker foi muito positiva. Alcançamos o resultado que esperávamos dentro do prazo, do orçamento e sem desafios inesperados”.

Uma consideração importante para aplicações de IA generativa é a baixa latência de inferência, pois a experiência do usuário precisa ser agradável. Quando os usuários redigem conteúdo usando uma ferramenta como o Wordtune, eles querem que a IA sirva como uma referência rápida sem retardar o processo de pensamento. Usando o Amazon SageMaker, a AI21 alcançou baixa latência de inferência com o modelo Grande para melhor atender às necessidades do cliente, reduzindo a latência em quatro vezes para um de seus grandes clientes. Como resultado, os clientes da AI21 podem atender a milhões de usuários diariamente, quase em tempo real, sem prejudicar a experiência do usuário. “Um de nossos clientes de grande escala viu uma melhoria significativa nas métricas de satisfação do usuário, atribuída à enorme redução na latência ao migrar do modelo Jumbo para o modelo Grande”, diz Padnos.

O lançamento do modelo Grande também contribuiu para o crescimento da AI21 e de seus clientes. “Depois de lançar o modelo Grande, que foi treinado usando o Amazon SageMaker, observamos um crescimento em nosso tráfego geral”, diz Padnos. “Clientes individuais que migraram para o modelo Grande também aumentaram seu tráfego”.

O modelo Grande (agora chamado de Mid) está disponível no Amazon SageMaker JumpStart, um hub de ML com algoritmos integrados, modelos básicos e soluções de ML pré-criadas que os usuários do Amazon SageMaker podem implantar com apenas alguns cliques. O ciclo de vida dos dados está contido no ambiente do usuário para manter a privacidade, e uma organização pode aplicar o modelo de linguagem aos dados sem escrever código ou precisar de um playground de código. A série de modelos básicos de próxima geração da AI21, Jurassic-2, bem como modelos específicos para tarefas, também estão disponíveis no Amazon SageMaker JumpStart.

A AI21 está entusiasmada com a crescente adoção da IA generativa em todo o mundo nos próximos meses e anos. Usando os serviços da AWS, a empresa está trabalhando ativamente em LLMs que serão mais rápidos, precisos, confiáveis e econômicos. “Temos um relacionamento muito bom com a equipe da AWS”, diz Padnos. “Os membros da equipe se aprofundaram nos detalhes técnicos conosco e colaboraram em tarefas desafiadoras. Durante todo o processo, a equipe da AWS foi criativa e conheceu nossos desafios e metas”.

Para saber mais, acesse https://aws.amazon.com/sagemaker.