Frugalidade encontra precisão: treinamento econômico de modelos GPT NeoX e Pythia com AWS Trainium | Amazon Web Services

Grandes modelos de linguagem (ou LLMs) tornaram-se um tópico de conversas diárias. Sua rápida adoção é evidente pela quantidade de tempo necessária para atingir 100 milhões de usuários, que passou de “4.5 anos pelo Facebook” para um mínimo histórico de apenas “2 meses pelo ChatGPT”. Um transformador generativo pré-treinado (GPT) usa atualizações autorregressivas causais para fazer previsões. Uma variedade de tarefas, como reconhecimento de fala, geração de texto e resposta a perguntas, demonstraram ter um desempenho estupendo por essas arquiteturas de modelo. Vários modelos recentes, como Neox, falcão, lhama use a arquitetura GPT como backbone. O treinamento de LLMs requer uma quantidade colossal de tempo de computação, que custa milhões de dólares. Neste post, resumiremos o procedimento de treinamento do GPT Neox on Treinamento AWS, um acelerador de aprendizado de máquina (ML) desenvolvido especificamente para treinamento de aprendizado profundo. Descreveremos como treinamos esses modelos de maneira econômica (3.2 milhões de tokens/$) com o AWS Trainium, sem perder a qualidade do modelo.

Visão geral da solução

Modelos GPT NeoX e Pythia

GPT NeoX e Pythia são os modelos de linguagem causal de código aberto da Eleuther-AI com aproximadamente 20 bilhões de parâmetros em NeoX e 6.9 ​​bilhões em Pythia. Ambos são modelos de decodificadores que seguem um projeto arquitetônico semelhante ao Chat GPT3. Porém, eles também possuem diversas adições, que também são amplamente adotadas nos modelos recentes como o Llama. Particularmente, eles possuem incorporação posicional rotacional (ROPE) com rotação parcial nas dimensões da cabeça. Os modelos originais (NeoX e Pythia 6.9B) são treinados em modelos disponíveis abertamente Empilhar conjunto de dados com desduplicação e usando back-end Megatron e Deepspeed.

Demonstramos o pré-treinamento e o ajuste fino desses modelos em instâncias Trn1 baseadas no AWS Trainium usando Neurônio NeMo biblioteca. Para estabelecer a prova de conceito e a reprodução rápida, usaremos um subconjunto menor do conjunto de dados da Wikipedia tokenizado usando o tokenizer GPT2 de codificação de pares de bytes (BPE).

Passo a passo

Baixe o conjunto de dados da Wikipedia pré-tokenizado conforme mostrado:

export DATA_DIR=~/examples_datasets/gpt2 mkdir -p ${DATA_DIR} && cd ${DATA_DIR} wget https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-vocab.json
wget https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-merges.txt
aws s3 cp s3://neuron-s3/training_datasets/gpt/wikipedia/my-gpt2_text_document.bin . --no-sign-request
aws s3 cp s3://neuron-s3/training_datasets/gpt/wikipedia/my-gpt2_text_document.idx . --no-sign-request
aws s3 cp s3://neuron-s3/training_datasets/gpt/wikipedia/license.txt . --no-sign-request

Tanto o NeoX 20B quanto o Pythia 6.9B utilizam ROPE com rotação parcial, por exemplo, girando 25% das dimensões da cabeça e mantendo o restante sem girar. Para implementar com eficiência a rotação parcial no acelerador AWS Trainium, em vez de concatenar as dimensões rotativas e não rotativas, acrescentamos frequências zero para dimensões não rotativas e, em seguida, giramos o conjunto completo de dimensões do cabeçote. Esse truque simples nos ajudou a melhorar o rendimento (sequências processadas por segundo) no AWS Trainium.

Etapas de treinamento

Para executar o treinamento, usamos Amazon Elastic Compute Cloud de vários nós gerenciados por SLURM (Amazon EC2) Cluster Trn1, com cada nó contendo uma instância trn1.32xl. Cada trn1.32xl possui 16 aceleradores com dois trabalhadores por acelerador. Depois de baixar o mais recente Neurônio NeMo pacote, use o fornecido neox e pitia scripts de pré-treinamento e ajuste fino com hiperparâmetros otimizados e execute o seguinte para um treinamento de quatro nós.

1. Compilar: Pré-compile o modelo com três iterações de treinamento para gerar e salvar os gráficos:

   sbatch --nodes 4 compile.slurm ./neoX_20B_slurm.sh

2. Executar: Execute o treinamento carregando os gráficos em cache das primeiras etapas

   sbatch --nodes 4 run.slurm ./neoX_20B_slurm.sh

3. Monitore os resultados

   tensorboard --logdir=nemo_experiments/megatron_neox

As mesmas etapas precisam ser seguidas para executar o modelo Pythia 6.9B com substituição neox_20B_slurm.sh by pythia_6.9B_slurm.sh.

Experimentos de pré-treinamento e ajuste fino

Demonstramos o pré-treinamento dos modelos GPT-NeoX e Pythia no AWS Trainium usando Neurônio NeMo biblioteca para iterações de 10k e também mostra o ajuste fino desses modelos para etapas de 1k. Para pré-treinamento, usamos o tokenizer GPT2 BPE dentro do NeMo e seguimos o mesmo configuração como usado no modelo original. O ajuste fino no AWS Trainium requer alteração de alguns parâmetros (como fator de divisão do tamanho do vocabulário), que são fornecidos nos scripts de ajuste fino para acomodar diferenças entre Megatron e NeMo e alterações entre GPU e AWS Trainium. A taxa de transferência de treinamento distribuído de vários nós com número variável de nós é mostrada na Tabela 1.

Tabela 1. Comparação do rendimento médio dos modelos GPT NeoX e Pythia para treinamento de até 500 etapas com alteração do número de nós. O preço de trn1.32xl baseia-se na tarifa efetiva reservada por hora por 3 anos.

A seguir, também avaliamos a trajetória de perda do treinamento do modelo no AWS Trainium e a comparamos com a execução correspondente em um cluster P4d (núcleos de GPU Nvidia A100). Junto com a perda de treinamento, também comparamos indicadores úteis, como norma de gradiente, que é a norma 2 dos gradientes do modelo calculados em cada iteração de treinamento para monitorar o progresso do treinamento. Os resultados do treinamento são mostrados na Figura 1, 2 e o ajuste fino do NeoX 20B na Figura 3.

Figura 1. Perda de treinamento calculada em média para todos os trabalhadores (esquerda) e norma de gradiente (direita) no treinamento em cada etapa. NeoX 20B é treinado em 4 nós com pequeno conjunto de dados wiki em GPU e Trainium com os mesmos hiperparâmetros de treinamento (tamanho global do lote = 256). A GPU está usando BF16 e precisão mista padrão, enquanto o AWS Trainium está usando BF16 completo com arredondamento estocástico. As trajetórias de norma de perda e gradiente correspondem à GPU e ao AWS Trainium.

Figura 2. Perda de treinamento calculada em média para todos os trabalhadores (esquerda) e norma de gradiente (direita) no treinamento em cada etapa. Semelhante ao GPT NeoX na Figura 1, Pythia 6.9B é treinado em 4 nós com pequeno conjunto de dados wiki em GPU e Trainium com os mesmos hiperparâmetros de treinamento (tamanho global do lote = 256). As trajetórias de norma de perda e gradiente correspondem a GPU e Trainium.

Figura 3. Modelo GPT NeoX 20B de ajuste fino em GPU e AWS Trainium com média de perda de treinamento em todos os trabalhadores (esquerda) e norma de gradiente (direita). Um pequeno conjunto de dados wiki é usado para demonstração de ajuste fino. As trajetórias de norma de perda e gradiente correspondem à GPU e ao AWS Trainium.

Nesta postagem, mostramos o treinamento econômico de LLMs em hardware de aprendizado profundo da AWS. Treinamos os modelos GPT NeoX 20B e Pythia 6.9B no AWS Trn1 com a biblioteca Neuron NeMo. A taxa de transferência normalizada de custo para 20 bilhões de modelos com AWS Trainium é de aproximadamente 3.2 milhões de tokens/US$ gastos. Junto com o treinamento econômico no AWS Trainium, obtemos uma precisão de modelo semelhante, o que é evidente na perda de etapas de treinamento e na trajetória da norma de gradiente. Também ajustamos os pontos de verificação disponíveis para o modelo NeoX 20B no AWS Trainium. Para obter informações adicionais sobre o treinamento distribuído com NeMo Megatron no AWS Trainium, consulte Referência do AWS Neuron para NeMo Megatron. Um bom recurso para iniciar o ajuste fino do modelo Llama pode ser encontrado aqui, Ajuste fino do Llama2. Para começar a usar o AWS Trainium gerenciado em Amazon Sage Maker, Ver Treine seus modelos de ML com AWS Trainium e Amazon SageMaker.

Sobre os autores

Gaurav Gupta atualmente é Cientista Aplicado nos laboratórios de IA da Amazon Web Services (AWS). Dr. Gupta completou seu doutorado pela USC Viterbi. Seus interesses de pesquisa abrangem o domínio da modelagem sequencial de dados, aprendizado de equações diferenciais parciais, teoria da informação para aprendizado de máquina, modelos dinâmicos fracionários e redes complexas. Atualmente trabalha em problemas aplicados e matemáticos sobre comportamento de treinamento de LLMs, modelos de visão com PDEs, modelos multimodais de teoria da informação. Gupta tem publicações em revistas/conferências importantes, como Neurips, ICLR, ICML, Nature, IEEE Control Society, ACM cyber-physical society.

Ben Snyder é um cientista aplicado do AWS Deep Learning. Seus interesses de pesquisa incluem modelos fundamentais, aprendizagem por reforço e otimização assíncrona. Fora do trabalho, ele gosta de andar de bicicleta e acampar no interior.

Amith (R) Mamidala é engenheiro sênior de aplicativos de aprendizado de máquina do AWS Annapurna Labs. Dr. Mamidala completou seu doutorado na Ohio State University em computação e comunicação de alto desempenho. Durante sua gestão na pesquisa da IBM, o Dr. Mamidala contribuiu para a classe de computadores BlueGene, que frequentemente liderava o ranking Top500 dos supercomputadores mais poderosos e eficientes em termos de energia. O projeto foi premiado com a Medalha Nacional de Tecnologia e Inovação 2009. Após uma breve passagem como engenheiro de IA em um fundo de hedge financeiro, o Dr. Mamidala ingressou nos laboratórios Annapurna com foco no treinamento de modelos de linguagem grande.

Jun (Lucas) Huan é cientista principal do AWS AI Labs. Dr. Huan trabalha com IA e ciência de dados. Ele publicou mais de 180 artigos revisados ​​por pares nas principais conferências e periódicos. Ele recebeu o Prêmio NSF de Desenvolvimento de Carreira em Início de Carreira em 2009. Antes de ingressar na AWS, ele trabalhou na pesquisa do Baidu como um cientista ilustre e chefe do Laboratório de Big Data do Baidu. Ele fundou a StylingAI Inc., uma start-up de IA, e trabalhou como CEO e cientista-chefe em 2019-2021. Antes de ingressar na indústria, foi professor Charles E. e Mary Jane Spahr no Departamento de EECS da Universidade do Kansas.

Shruti Koparkar é gerente sênior de marketing de produtos da AWS. Ela ajuda os clientes a explorar, avaliar e adotar a infraestrutura de computação acelerada Amazon EC2 para suas necessidades de aprendizado de máquina.

• Conteúdo com tecnologia de SEO e distribuição de relações públicas. Seja amplificado hoje.
• PlatoData.Network Gerativa Vertical Ai. Capacite-se. Acesse aqui.
• PlatoAiStream. Inteligência Web3. Conhecimento Amplificado. Acesse aqui.
• PlatãoESG. Carbono Tecnologia Limpa, Energia, Ambiente, Solar, Gestão de resíduos. Acesse aqui.
• PlatoHealth. Inteligência em Biotecnologia e Ensaios Clínicos. Acesse aqui.
• Fonte: https://aws.amazon.com/blogs/machine-learning/frugality-meets-accuracy-cost-efficient-training-of-gpt-neox-and-pythia-models-with-aws-trainium/