Como o Amazon Search M5 economizou 30% no custo de treinamento LLM usando o AWS Trainium | Amazon Web Services

Durante décadas, a Amazon foi pioneira e inovou no aprendizado de máquina (ML), trazendo experiências encantadoras aos seus clientes. Desde os primeiros dias, a Amazon usou ML para vários casos de uso, como recomendações de livros, pesquisa e detecção de fraudes. Semelhante ao resto da indústria, os avanços do hardware acelerado permitiram que as equipes da Amazon buscassem arquiteturas de modelo usando redes neurais e aprendizagem profunda (DL).

O programa M5 do Amazon Search possui a estratégia de aprendizagem de descoberta para a Amazon e constrói modelos em grande escala multilíngues, multilocais, multientidades, multitarefas e multimodais, como texto, imagem e vídeo. O programa M5 tem fornecido incorporações universais e modelos básicos em grande escala para centenas de equipes de ML em toda a Amazon, ao mesmo tempo que mantém controles rígidos sobre a otimização de custos. Para conseguir isso, a equipe M5 avalia regularmente novas técnicas para reduzir custos.

Como muitas organizações de ML, os aceleradores são amplamente usados ​​para acelerar o treinamento e a inferência de DL. Quando a AWS lançou aceleradores específicos com a primeira versão do Inferência da AWS em 2020, a equipe M5 rapidamente começou a utilizá-los para implantar cargas de trabalho de produção com mais eficiência, economizando custos e reduzindo a latência. No ano passado, a AWS lançou seu Treinamento AWS aceleradores, que otimizam o desempenho por custo para desenvolver e construir modelos DL de próxima geração. Neste post, discutimos como a M5 conseguiu reduzir em 30% o custo de treinamento de seus modelos e compartilhamos algumas das melhores práticas que aprendemos ao longo do caminho.

Instâncias de treinamento

Com os avanços em aceleradores específicos, a Amazon também fornece aceleradores atraentes na forma de AWS Inferentia e Trainium. Como o nome indica, esses chips são otimizados para exceder as necessidades de inferência e cargas de trabalho de treinamento, respectivamente. Para treinamento em larga escala de modelos de fundação que atingem bilhões de parâmetros de tamanho, o Trainium Instâncias Trn1 e Trn1n são escolhas ideais devido às suas características. As instâncias Trn1 são alimentadas pelo estado da arte NeuronCore-v2e tem uma grande quantidade de computação e memória do acelerador. As instâncias Trn1n também podem ser escolhidas para uma quantidade maior de largura de banda de rede (1,600 Gbs), portanto, são ideais para treinamento de desempenho com otimização de custos em mente.

Para usar aceleradores, você precisa de uma camada de software para suportá-los. Com os chips Trn e Inf, o SDK do AWS Neuron desbloqueia aceleradores específicos da Amazon com a ajuda do PyTorch XLA. PyTorch XLA converte o modo ansioso do PyTorch em implementação baseada em gráfico no modo preguiçoso. Esses gráficos são então usados ​​e compilados para serem usados ​​com o acelerador. PyTorch Neuron (parte do Neuron SDK) permite que os usuários do PyTorch treinem seus modelos no Trainium NeuronCores com algumas linhas de código.

Modelo e carga de trabalho

A equipe M5 treina e implanta modelos fundamentais e representações universais para ajudar diversas equipes em toda a Amazon a levar satisfação aos clientes. Amazon.com clientes. Um desses modelos é um modelo de codificador de texto seguido por um perceptron multicamadas (MLP) com interações de recursos explícitas ou implícitas definidas pela arquitetura da rede neural com centenas de milhões de parâmetros treináveis. Este modelo é treinado em bilhões de tokens e usado para gerar milhões de embeddings em uma configuração de inferência em lote offline. Essas incorporações são entradas para um serviço Amazon de nível 1 voltado para o cliente.

A infraestrutura para o pipeline de produção utiliza Lote da AWS de estratégias de filas de compartilhamento justo, usando um cluster trn1.32xlarge de vários nós habilitado para EFA como cálculo para treinamento de modelo. Funcionalmente, o pipeline de produção realiza treinamento incremental do modelo, avaliação do modelo treinado e inferência em lote offline no modelo treinado, tudo usando PyTorch como biblioteca DL subjacente.

Objetivos

Encantar nossos clientes é um princípio fundamental. Dada a natureza do pipeline voltado para o cliente, é fundamental que todos os acordos de nível de serviço (SLAs) sejam cumpridos sem regressões. Identificamos dois critérios críticos de aceitação para adaptar nosso pipeline de produção de GPU existente e fazer a transição para o Trainium:

• Qualidade do modelo – A qualidade dos nossos modelos impacta diretamente a experiência do cliente. Exigimos que haja menos de 0.1% de diferença na qualidade do modelo entre GPU e Trainium.
• Taxa de transferência de treinamento – Treinamos iterativamente nossos modelos periodicamente para fornecer a experiência mais recente aos nossos clientes. Exigimos que a convergência do modelo seja alcançada dentro de um período de tempo predefinido (como 1 semana) para cumprir nossos SLAs de produção.

Nas seções a seguir, compartilhamos nossa jornada de retrocesso a partir desses critérios e nossos aprendizados para dar suporte a cargas de trabalho de produção em escala amazônica.

Roteiro de treinamento

Antes de iniciar o treinamento do modelo, precisamos fazer alterações no script de treinamento para torná-lo compatível com XLA. Dado o tamanho do modelo, usamos dados paralelos distribuídos (DDP) para treinar o modelo. O DDP nos permite aumentar o rendimento do treinamento do modelo, aumentando o número de máquinas usadas para executar o treinamento do modelo, sem nenhuma alteração no código. Seguimos as instruções fornecidas no Tutorial de treinamento Neuron PyTorch MLP para adicionar construções específicas do XLA em nossos scripts de treinamento. Essas alterações de código são simples de implementar. A seguir estão alguns aprendizados técnicos significativos do exercício que melhoraram muito o rendimento do nosso modelo:

• Colocação de xm.mark_step() - xm.mark_step() compila e executa os gráficos de computação coletados preguiçosamente. Invocando mark_step muitas vezes levará a um número maior de gráficos pequenos, enquanto invocá-lo poucas vezes levará a poucos, mas grandes gráficos. Dependendo da sua aplicação, o rendimento e a implementação do treinamento do seu modelo variarão com base no posicionamento do seu modelo. xm.mark_step(). Nossa implementação coloca um xm.mark_step() após uma passagem para frente e para trás e uma após a etapa do otimizador.
• Envolvimento do carregador de dados com carregador de dispositivo de multiprocessamento XLA – Esta é uma etapa crítica que pode ser facilmente perdida. O carregador de dispositivo de multiprocessamento torch_xla.distributed.parallel_loader.MpDeviceLoader carrega dados de treinamento em cada dispositivo XLA com opções para pré-carregar e sobrepor o carregamento de dados com execuções do dispositivo para melhorar o rendimento. O carregador de dispositivo também invoca xm.mark_step() e, portanto, é capaz de construir gráficos para carregamento de dados do host para o dispositivo.

Compilação para Trainium

Tradicionalmente, o ciclo de desenvolvimento de modelo com GPUs envolve fazer alterações no modelo ou script de treinamento e executá-lo diretamente no dispositivo GPU. Aceleradores como o Trainium que usam XLA exigem uma etapa adicional antes que o treinamento do modelo possa ser executado no acelerador. Os gráficos de computação XLA só podem ser executados após serem compilados. Geralmente, há duas maneiras de realizar esta compilação: Ahead of Time (AOT), onde você rastreia e compila todos os gráficos primeiro e depois os executa, ou Just In Time (JIT), onde os gráficos são rastreados, compilados e executados conforme eles aparecem. são encontrados. O Neuron SDK fornece ambos prontos para uso. Normalmente, a compilação AOT é executada primeiro. Os gráficos são então executados após esta compilação. Se novos gráficos forem encontrados, o tempo de execução do Neuron invoca uma compilação JIT antes de executá-los. Para realizar a compilação AOT, o Neuron SDK fornece neuron_parallel_compile, um utilitário de compilação que extrai gráficos de uma execução de teste do script de treinamento e executa compilação AOT paralela.

Um aspecto importante da compilação AOT é garantir que nenhum novo gráfico de computação seja criado durante o treinamento. Uma fonte de novos gráficos de computação (e, portanto, recompilações) são as formas dinâmicas dos lotes de treinamento durante o treinamento do modelo. Descobrimos que o uso de formas estáticas e lotes de tamanho fixo elimina compilações de tempo de treinamento e melhora muito o rendimento do treinamento sem qualquer efeito na precisão do modelo. Ao impor tais restrições ao treinamento, observamos que apenas 4 a 5 etapas de treinamento do modelo, uma etapa de validação do modelo e verificação do modelo uma vez são necessárias para rastrear todos os gráficos durante a compilação AOT. É importante observar que o Neuron SDK está em constante evolução e, no futuro, também oferecerá suporte a formas dinâmicas.

Além disso, os gráficos compilados são armazenados no Cache Persistente de Neurônios em disco ou em um Serviço de armazenamento simples da Amazon (Amazon S3) balde. Isso é especialmente útil para cargas de trabalho de produção em que a arquitetura do modelo e a configuração do treinamento não mudam. Portanto, a sobrecarga de compilação ocorre apenas uma vez. Usar o cache é tão simples quanto definir um sinalizador de ambiente:

export NEURON_COMPILE_CACHE_URL="s3://BUCKET/KEY"

O compilador Neuron também fornece três opções de otimização em nível de compilador (O1, O2, O3) para equilibrar o tempo de compilação e o rendimento da execução do modelo. O1 permite otimizações básicas no gráfico de computação e minimiza o tempo de compilação, O3 fornece melhor rendimento de execução do modelo ao custo de maior tempo de compilação e O2 (opção padrão) é um equilíbrio entre os dois. Para nosso caso de uso, usamos a otimização O1 e observamos uma redução de 86% no tempo de compilação sem nenhuma alteração nas métricas de precisão do modelo, enquanto observamos uma redução de aproximadamente 5–7% no rendimento em comparação com a otimização padrão (O2). Dependendo do caso de uso, você pode escolher diferentes níveis de otimização.

Para resumir, usamos os seguintes sinalizadores para compilação:

NEURON_CC_FLAGS="--target trn1 --auto-cast all --auto-cast-type bf16 --model-type transformer --optlevel O1"

Compatibilidade de pontos de verificação

Quando a compilação for concluída com sucesso, podemos prosseguir com o treinamento de nossos modelos no Trainium. Conforme mencionado anteriormente, treinamos nossos modelos de forma incremental, o que significa que carregamos um ponto de verificação do modelo previamente treinado e continuamos o treinamento com novos dados. PyTorch e PyTorch XLA permitem uma transição perfeita entre aceleradores por meio da interoperabilidade de pontos de verificação. Ter a flexibilidade de alternar entre GPU e Trainium nos permitiu carregar perfeitamente o modelo de GPU anterior e treinar em máquinas Trainium. Isso foi fundamental para garantir que pudéssemos inicializar nosso modelo com o melhor modelo previamente treinado, sem qualquer tempo de inatividade da produção ou perda na precisão do modelo.

Como o modelo de GPU foi salvo usando utilitários de salvamento de modelo PyTorch padrão, pudemos usar o utilitário de carregamento de ponto de verificação PyTorch para carregar o modelo de GPU em dispositivos Trainium.

Por exemplo, em GPU/CPU, você pode salvar o modelo com o seguinte código:

torch.save(model.state_dict(), PATH)

Então você carrega o modelo de volta no Trainium:

import torch_xla.core.xla_model as xm
xla_device = xm.xla_device()
model = MyModel(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.to(xla_device)

Da mesma forma, você pode salvar o modelo no Trainium com o seguinte código:

import torch_xla.core.xla_model as xm
# automatically moves the data to CPU for the master device
xm.save(model.state_dict(), PATH) 

E carregue o modelo de volta na GPU/CPU:

model = MyModel(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.to(device) # can be any device

Na verdade, como usamos DDP para treinamento de modelo, o carregamento do modelo é independente do número de máquinas usadas para treinar o ponto de verificação anterior. Isso nos permite dimensionar horizontalmente a frota Trn1 sem alterações de código ou efeitos adversos no treinamento do modelo. Esses pontos de verificação baseados em PyTorch podem ser usados ​​diretamente ou até mesmo com scripts de tocha para casos de uso de inferência no AWS Inferentia2 ou outros aceleradores.

Estabilidade operacional

Nunca é demais enfatizar que a execução de cargas de trabalho em produção exige o cumprimento de vários SLAs. Para nosso caso de uso, além dos SLAs de qualidade do modelo e rendimento de treinamento, é imperativo que o pipeline de produção seja operacionalmente estável, o que significa tempo de inatividade mínimo e interrupções durante o treinamento, avaliação e inferência do modelo.

Tal como acontece com o pipeline existente baseado em GPU, adicionamos vários mecanismos para tornar o pipeline operacionalmente estável. Antes de iniciar o treinamento do modelo, executamos vários testes de integridade para avaliar a integridade das máquinas. Esses testes geralmente incluem operações simples de tensor para verificar a integridade dos dispositivos aceleradores. Observamos que para treinamentos distribuídos é importante realizar testes para verificar também a comunicação coletiva entre instâncias. Nós usamos o Conjunto de testes NCCOM do Neuron SDK para conseguir isso, executando uma variedade de operações, como coleta total, redução total e dispersão reduzida.

Mesmo depois de seguir as sugestões que mencionamos, observamos que problemas transitórios são inevitáveis ​​em qualquer pipeline, independentemente do acelerador subjacente. Para criar resiliência em qualquer pipeline de treinamento, recomendamos a criação de mecanismos de repetição para resolver esses possíveis problemas. Nós usamos Novas tentativas automatizadas do AWS Batch para tentar novamente trabalhos que encontrem uma falha transitória durante o treinamento do modelo. Essas reinicializações podem custar caro se for encontrada uma falha no final do treinamento. Para combater esse problema, adaptamos nossos scripts de treinamento para carregar um ponto de verificação do modelo previamente treinado e continuar o treinamento a partir desse ponto. Com essa funcionalidade, podemos reiniciar agressivamente trabalhos de treinamento com falha com sobrecarga mínima.

Com esses mecanismos de resiliência implementados, conseguimos atingir taxas de sucesso de 98.5% para nossas cargas de trabalho no Trn1, comparáveis ​​às taxas de sucesso de nosso pipeline de GPU existente.

Resultados

Para validar a precisão de nossos modelos, inicializamos dois modelos do mesmo ponto de verificação de GPU e treinamos um no Trainium e o outro em uma GPU comparável. Ambos os modelos foram treinados com os mesmos hiperparâmetros de treinamento. O conjunto de dados usado para cálculo de métricas é um conjunto de dados de validação e avaliamos a precisão do modelo neste conjunto de dados a cada N etapas globais. O eixo X é a etapa global e o eixo Y é a precisão do modelo. Observamos menos de 0.1% de diferença na precisão do modelo em cada ponto do gráfico a seguir.

Além disso, para avaliar a relação custo-benefício do treinamento do modelo, preferimos comparar o tempo necessário para atingir a convergência do modelo. Acreditamos que isso fornece uma visão mais prática da economia de custos em comparação com medidas como custo por token, FLOPS/dólar alcançado e outros fatores. Considerando o tempo de treinamento de trn1.32xl e comparável Amazon Elastic Compute Nuvem (Amazon EC2), observamos que o Trainium oferece um custo até 30% mais barato para modelar a convergência.

Conclusão

Há muitos fatores a serem considerados ao avaliar diferentes aceleradores para suas cargas de trabalho de DL. Alguns dos mais importantes são qualidade do modelo, rendimento, custo e disponibilidade. É fundamental garantir que a qualidade e o rendimento do seu modelo não sejam sacrificados com base no acelerador que você escolher.

Graças à nossa parceria e colaboração com a equipe do Annapurna Neuron, a equipe do Amazon Search M5 conseguiu economizar até 30% em custos ao mudar para o Trainium. A equipe é capaz de usar o Trainium e alcançar qualidade de modelo e paridade de rendimento com aceleradores comparáveis ​​no mercado. A interoperabilidade dos pontos de verificação e as alterações mínimas de código com suporte para XLA permitiram que o M5 escolhesse entre vários aceleradores para suas cargas de trabalho. Isso permitiu que a equipe M5 aproveitasse o grande poder computacional do Trainium e construísse soluções independentes de aceleradores para encantar os clientes da Amazon.com. Do ponto de vista operacional, o Trainium provou ser capaz de oferecer suporte a serviços de nível 1 na escala Amazon. A equipe M5 continua transferindo mais cargas de trabalho para o Trainium para fornecer os melhores modelos para a Amazon com os custos mais baixos.

Em resumo, a equipe M5 conseguiu realizar treinamento de ML de nível de produção com boa relação custo-benefício, adicionando o Trainium à frota de aceleradores. Incentivamos você a dar uma olhada no Trainium e em outros dispositivos Neuron, como o AWS Inferentia, para colher os benefícios do silício da Amazon desenvolvido especificamente para cargas de trabalho de ML. Comece facilmente com um dos muitos tutoriais com diferentes modelos, como Lhama 2, disponível no Trainium.

Sobre os autores

Abhinandan Patni é engenheiro de software sênior na Amazon Search. Ele se concentra na construção de sistemas e ferramentas para treinamento de aprendizado profundo distribuído e escalável e inferência em tempo real.

James Park é arquiteto de soluções na Amazon Web Services. Ele trabalha com a Amazon.com para projetar, construir e implantar soluções tecnológicas na AWS e tem interesse particular em IA e aprendizado de máquina. Nas horas vagas ele gosta de buscar novas culturas, novas experiências e de se manter atualizado com as últimas tendências tecnológicas. Você pode encontrá-lo em LinkedIn.

Jerry Mannil é engenheiro de software na Amazon Search. Ele trabalha na melhoria da eficiência, robustez e escalabilidade da infraestrutura de treinamento distribuída.

Ken Su é engenheiro de software na Amazon Search. Ele trabalha para melhorar a eficiência do treinamento e o fluxo de trabalho de treinamento distribuído escalonável. Fora do trabalho, gosta de caminhadas e tênis.

RJ é um engenheiro da Amazon. Ele cria e otimiza sistemas para sistemas distribuídos para treinamento e trabalha na otimização da adoção de sistemas para reduzir a latência para inferência de ML. Fora do trabalho, ele está explorando o uso de IA generativa para criar receitas de alimentos.

• 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/how-amazon-search-m5-saved-30-for-llm-training-cost-by-using-aws-trainium/