Este blog é coautor de Sarah Jane Hong CSO, Darryl Barnhart CTO e Ian Thompson CEO da Latent Space e Prem Ranga da AWS.
O espaço latente é uma representação oculta de ideias abstratas que os modelos de aprendizado de máquina (ML) aprendem. Por exemplo, “cachorro”, “flor” ou “porta” são conceitos ou localizações no espaço latente. No Espaço Latente, estamos trabalhando em um mecanismo que permite manipular e explorar esse espaço com linguagem e prompts visuais. A equipe do Espaço Latente vem de duas áreas que há muito têm pouca sobreposição: gráficos e processamento de linguagem natural (PNL). Tradicionalmente, as modalidades de imagens e texto têm sido tratadas separadamente, cada um com sua própria história de engenharia de recursos complexos, caros e frágeis. Tarefas de PNL, como compreensão de documentos ou resposta a perguntas, geralmente têm pouco em comum com tarefas de visão, como compreensão ou renderização de cenas, e geralmente usamos abordagens e modelos muito diferentes para cada tarefa. Mas isso está mudando rapidamente.
Essa fusão de modalidades em um único espaço latente compartilhado abre uma nova geração de aplicativos criativos e comerciais, de jogos a compreensão de documentos. Mas desbloquear esses novos aplicativos em um único modelo abre novos desafios de dimensionamento, conforme destacado em “The Bitter Lesson” de Richard Sutton, e o trabalho empolgante dos últimos anos sobre leis de dimensionamento. Para tornar isso possível, o Espaço Latente está trabalhando em pesquisas de ponta para fundir essas modalidades em um único modelo, mas também para dimensionar e fazê-lo de forma eficiente. É aqui que entra o paralelismo do modelo.
Amazon Sage MakerO particionamento de modelo automatizado exclusivo e a abordagem de pipelining eficiente tornaram nossa adoção do paralelismo de modelo possível com pouco esforço de engenharia, e escalamos nosso treinamento de modelos além de 1 bilhão de parâmetros (usamos o p4d.24xlarge A100 instâncias), que é um requisito importante para nós. Além disso, observamos que ao treinar com uma configuração de treinamento de 16 nós e oito GPU com a biblioteca de paralelismo de modelo SageMaker, registramos uma melhoria de 38% na eficiência em comparação com nossas execuções de treinamento anteriores.
Desafios com o treinamento de transformadores de grande escala
No Latent Space, estamos fundindo linguagem e visão em modelos de transformador com bilhões de parâmetros para dar suporte a casos de uso “fora da distribuição” da imaginação de um usuário ou que ocorram no mundo real, mas não em nossos dados de treinamento. Estamos lidando com os desafios inerentes ao dimensionamento para bilhões de parâmetros e além de duas maneiras diferentes:
As técnicas de recuperação de informações têm sido um componente-chave dos mecanismos de pesquisa e das tarefas de controle de qualidade. Recentemente, um progresso emocionante foi feito combinando técnicas clássicas de IR com transformadores modernos, especificamente para tarefas de resposta a perguntas em que um modelo é treinado em conjunto com um recuperador neural que aprende a recuperar documentos relevantes para ajudar a responder a perguntas. Para uma visão geral, veja o trabalho recente da FAIR em Geração Aumentada de Recuperação: Simplificando a criação de modelos de processamento de linguagem natural inteligente e Fusão no Decodificador, Do Google Brain REINOe da Nvidia Recuperador Neural para responder a perguntas.
Embora as técnicas de recuperação aumentada ajudem com custos e eficiência, ainda não conseguimos ajustar o modelo em uma única GPU para nosso modelo maior. Isso significa que precisamos usar o paralelismo de modelo para treiná-lo. No entanto, devido à natureza da nossa arquitetura de recuperação, projetar a divisão do nosso modelo foi um desafio devido às interdependências entre os contextos recuperados nas entradas de treinamento. Além disso, mesmo se determinarmos como dividimos nosso modelo, introduzir o paralelismo de modelo foi uma tarefa de engenharia significativa para fazermos manualmente em todo o nosso ciclo de vida de pesquisa e desenvolvimento.
A biblioteca de paralelismo de modelo SageMaker
Paralelismo de modelo é o processo de dividir um modelo entre vários dispositivos ou nós (como instâncias equipadas com GPU) e criar um pipeline eficiente para treinar o modelo nesses dispositivos para maximizar a utilização da GPU. O biblioteca de paralelismo de modelos no SageMaker torna o paralelismo de modelo mais acessível, fornecendo divisão de modelo automatizada, também conhecida como particionamento de modelo automatizado e programação sofisticada de execução de pipeline. Os algoritmos de divisão do modelo podem otimizar a velocidade ou o consumo de memória. A biblioteca usa um algoritmo de particionamento que equilibra a memória, minimiza a comunicação entre os dispositivos e otimiza o desempenho.
Particionamento de modelo automatizado
Para o nosso caso de uso PyTorch, a biblioteca paralela do modelo executa internamente uma etapa de rastreamento (na primeira etapa de treinamento) que constrói o gráfico do modelo e determina as formas do tensor e dos parâmetros. Em seguida, ele constrói uma árvore, que consiste em nn.Module objetos no modelo, bem como dados adicionais coletados do rastreamento, como a quantidade de nn.Parameters, e tempo de execução para cada nn.Module.
A biblioteca então atravessa essa árvore a partir da raiz e executa um algoritmo de particionamento que equilibra a carga computacional e o uso de memória e minimiza a comunicação entre as instâncias. Se vários nn.Modules compartilharem o mesmo nn.Parameter, esses módulos serão colocados no mesmo dispositivo para evitar a manutenção de várias versões do mesmo parâmetro. Após a decisão de particionamento ser tomada, os módulos atribuídos e pesos são carregados em seus dispositivos.
Agendamento de execução de pipeline
Outro recurso central da biblioteca paralela de modelo distribuído SageMaker é execuções em pipeline, que determinam a ordem em que os cálculos são feitos e os dados são processados nos dispositivos durante o treinamento do modelo. O pipelining é baseado na divisão de um minilote em microbatches, que são alimentados no pipeline de treinamento um por um e seguem um cronograma de execução definido pelo tempo de execução da biblioteca.
O pipeline de microbatch garante que todas as GPUs sejam totalmente utilizadas, o que é algo que teríamos que construir nós mesmos, mas com a biblioteca de paralelismo de modelo isso acontece perfeitamente nos bastidores. Por fim, podemos usar Amazon FSx, que é importante para garantir que nossas velocidades de leitura sejam rápidas, dado o número de arquivos lidos durante o treinamento de um modelo multimodal com recuperação.
Arquitetura de treinamento
O diagrama a seguir representa como configuramos nossa arquitetura de treinamento. Nossos objetivos principais eram melhorar a velocidade do treinamento e reduzir custos. Os transformadores de imagem e linguagem que estamos treinando são altamente complexos, com um número significativamente grande de camadas e pesos internos, atingindo bilhões de parâmetros, o que os torna incapazes de caber na memória de um único nó. Cada nó carrega um subconjunto do modelo, por meio do qual os fluxos de dados e as transformações são compartilhados e compilados. Nós configuramos 16 p4d.24xgrande instâncias, cada uma com oito GPUs, usando a seguinte representação de arquitetura:
À medida que aumentamos nossos modelos, uma tendência comum é ter tudo armazenado nos pesos da rede. No entanto, para fins práticos, queremos aumentar nossos modelos para aprender como procurar contextos relevantes para ajudar na tarefa de renderização. Isso nos permite manter nossos custos de serviço baixos, sem comprometer a qualidade da imagem. Usamos um grande modelo de PNL baseado em transformador e, conforme mencionado antes, observamos um aumento de 38% na eficiência de treinamento com a biblioteca de paralelismo de modelo SageMaker, conforme mostrado a seguir:
- Precisamos de um allreduce para cada cálculo no caso de paralelismo de nível de tensor. Isso leva O (log_2 n) etapas paralelas. Ou seja, n máquinas executando O (n) etapas, para O (n log_2 n) operações totais.
- Para o paralelismo do pipeline, exigimos O (1) etapas paralelas para passar dados pelo pipeline
- Dadas 16 máquinas com oito GPUs, temos custo O (1) para paralelo de pipeline e custo O (log_2 (8)) = O (3) para paralelo de modelo em profundidade.
- Neste caso, vemos que o custo da rede é reduzido para 1/3 mudando para o paralelo de pipeline que usamos com o paralelismo do modelo SageMaker, e o custo geral de treinamento reduz para 1/2 + 1/2 * 1 / log_2 (16 ) = 0.625 do custo original, levando a uma melhoria de eficiência correspondente.
Em geral, quando a necessidade justifica o treinamento distribuído (problemas com o tamanho do modelo de dimensionamento ou dados de treinamento), podemos seguir um conjunto de práticas recomendadas para determinar qual abordagem funciona melhor.
Melhores práticas para treinamento distribuído
Com base em nossa experiência, sugerimos começar com uma abordagem paralela de dados distribuídos. Paralelismo de dados distribuídos, como o Biblioteca paralela de dados distribuídos SageMaker resolve a maioria dos problemas de rede com réplicas de modelo, portanto, você deve ajustar os modelos no menor número de nós e, em seguida, replicar para dimensionar o tamanho do lote conforme necessário.
Se ficar sem memória durante o treinamento, como fizemos neste cenário, você pode mudar para uma abordagem de modelo paralelo. No entanto, considere essas alternativas antes de tentar o treinamento paralelo de modelo:
- Em hardware equipado com NVIDIA Tensor Core, use treinamento de precisão mista para aumentar a velocidade e reduzir o consumo de memória.
- Reduza o tamanho do lote (ou reduza a resolução da imagem ou o comprimento da sequência NLP, se possível).
Além disso, preferimos projetos de modelo que não tenham normalização em lote, conforme descrito em Reconhecimento de imagem em grande escala de alto desempenho sem normalização. Se isso não puder ser evitado, certifique-se de que a normalização em lote seja sincronizada entre os dispositivos. Quando você usa o treinamento distribuído, seu lote é dividido em GPUs, portanto, estatísticas de lote precisas exigem sincronização em todos os dispositivos. Sem isso, a normalização terá erro aumentado e, portanto, prejudicará a convergência.
Comece com o treinamento paralelo do modelo quando tiver as seguintes restrições:
- Seu modelo não cabe em um único dispositivo
- Devido ao tamanho do seu modelo, você está enfrentando limitações na escolha de tamanhos de lote maiores, como se os pesos do seu modelo ocupassem a maior parte da memória da GPU e você fosse forçado a escolher um tamanho de lote menor e abaixo do ideal
Ao otimizar para desempenho, faça o seguinte:
- Use pipelining para comunicações entre nós para minimizar a latência e aumentar o rendimento
- Mantenha as tubulações o mais curtas possível para minimizar quaisquer bolhas. O número de microbatches deve ser ajustado para equilibrar a eficiência computacional com o tamanho da bolha e ter pelo menos o comprimento do pipeline. Se necessário, você pode formar microbatches no nível do token, conforme descrito em TeraPipe: paralelismo de pipeline de nível de token para treinar modelos de linguagem em grande escala
Ao otimizar o custo, use Instâncias Spot gerenciadas pelo SageMaker para treinamento. Isso pode otimizar o custo dos modelos de treinamento em até 90% em relação às instâncias On-Demand. SageMaker gerencia as interrupções Spot em seu nome.
Outros fatores a serem considerados:
- Em um nó, quando há uma interconexão rápida, há mais nuances. Se houver ampla capacidade de rede intra-nó, reorganizar os dados para uma computação mais otimizada pode mostrar um benefício.
- Se as ativações forem muito maiores do que os tensores de peso, um otimizador fragmentado também pode ajudar. Por favor, consulte Zero para mais detalhes.
A tabela a seguir fornece alguns cenários de scaleup de treinamento comuns e como você pode configurá-los na AWS.
| Cenário | Quando isso se aplica? | Solução |
| Escalonamento de uma única GPU para várias GPUs | Quando a quantidade de dados de treinamento ou o tamanho do modelo é muito grande | Mude para uma instância multi-GPU, como p3.16xlarge, que tem oito GPUs, com os dados e o processamento divididos entre as oito GPUs e produzindo um aumento de velocidade quase linear no tempo que leva para treinar seu modelo. |
| Escalonamento de uma única instância para várias instâncias | Quando o dimensionamento precisa ir além da alteração do tamanho da instância | Escale o número de instâncias com a função de estimador do SageMaker Python SDK, definindo seu instance_type para p3.16xlarge e instance_count para 2. Em vez de oito GPUs em um único p3.16xlarge, você tem 16 GPUs em duas instâncias idênticas. Considere usar o Biblioteca paralela de dados distribuídos SageMaker. |
| Seleção de um modelo de abordagem paralela para treinamento | Ao encontrar erros de falta de memória durante o treinamento | Mude para uma abordagem de modelo paralelo usando o Biblioteca paralela de modelo distribuído SageMaker. |
| Desempenho de rede para comunicações entre nós | Para treinamento distribuído com várias instâncias (por exemplo, comunicação entre os nós no cluster ao fazer uma operação AllReduce) | Suas instâncias precisam estar na mesma região e mesma zona de disponibilidade. Quando você usa o SageMaker Python SDK, isso é feito para você. Seus dados de treinamento também devem estar na mesma Zona de disponibilidade. Considere usar o Biblioteca paralela de dados distribuídos SageMaker. |
| GPU, rede e armazenamento otimizados | Para necessidades de treinamento distribuído em larga escala | O tipo de instância p4d.24xlarge foi projetado para armazenamento local rápido e um backplane de rede rápido com até 400 gigabits, e é altamente recomendável como a opção de melhor desempenho para treinamento distribuído. |
Conclusão
Com a biblioteca paralela de modelos no SageMaker, obtemos muitos benefícios prontos para uso, como particionamento de modelo automatizado e pipelining eficiente. Nesta postagem, compartilhamos nossos desafios com nosso caso de uso de ML, nossas considerações sobre diferentes abordagens de treinamento e como usamos a biblioteca de paralelismo de modelos Amazon SageMaker para acelerar nosso treinamento. O melhor de tudo é que agora pode levar apenas algumas horas para adotar as práticas recomendadas para paralelismo de modelo e melhorias de desempenho descritas aqui. Se esta postagem ajudar ou inspirar você a resolver um problema, adoraríamos saber! Por favor, compartilhe seus comentários e feedback.
Referências
Para obter mais informações, consulte o seguinte:
Sobre os autores
Prem Ranga é arquiteto de soluções corporativas baseado em Atlanta, GA. Ele faz parte da comunidade técnica de aprendizado de máquina e adora trabalhar com clientes em sua jornada de ML e IA. Prem é apaixonado por robótica, é pesquisador de veículos autônomos e também construiu o Beer Pours, controlado pelo Alexa, em Houston e outros locais.
Sarah Jane Hong é o co-fundador e Diretor de Ciência da Latent Space. Sua experiência está na interseção da interação humano-computador e aprendizado de máquina. Anteriormente, ela liderou a pesquisa de PNL na Sonar (adquirida pela Marchex), que atende empresas no espaço de IA de conversação. Ela também é uma conceituada desenvolvedora de AR / VR, tendo recebido prêmios e bolsas de estudo da Oculus, Mozilla Mixed Reality e Microsoft Hololens.
Darryl Barnhart é cofundador e diretor de tecnologia da Latent Space. Ele é um desenvolvedor experiente com experiência em aceleração de GPU, computação gráfica, dados em grande escala e aprendizado de máquina. Outras paixões incluem matemática, desenvolvimento de jogos e o estudo da informação.
Ian Thompson é o fundador e CEO da Latent Space. Ian é um engenheiro e pesquisador inspirado pelo “possível adjacente” - tecnologias prestes a ter um grande impacto em nossas vidas. Atualmente focado em simplificar e dimensionar o aprendizado de representação multimodal para ajudar a construir IA segura e criativa. Anteriormente, ele ajudou a construir empresas em gráficos / realidade virtual (AltspaceVR, adquirida pela Microsoft) e educação / PNL (HSE).
