Os pesquisadores do Laboratório Nacional de Oak Ridge treinaram um modelo grande de linguagem (LLM) do tamanho do ChatGPT no supercomputador Frontier, usando apenas 3.072 de suas 37.888 GPUs. A equipe publicou um artigo de pesquisa que detalha como realizaram o feito e os desafios enfrentados ao longo do caminho.
O supercomputador Frontier possui 9.472 CPUs Epyc 7A53 e 37.888 GPUs Radeon Instinct. No entanto, a equipe usou apenas 3.072 GPUs para treinar um LLM com um trilhão de parâmetros e 1.024 GPUs para treinar outro LLM com 175 bilhões de parâmetros.
O artigo observa que o principal desafio no treinamento de um LLM tão grande é a quantidade de memória necessária, que era de no mínimo 14 terabytes. Isso significava que várias GPUs MI250X com 64 GB de VRAM cada precisavam ser usadas, mas isso introduziu um novo problema: paralelismo. Colocar mais GPUs em um LLM requer uma comunicação cada vez melhor para realmente usar mais recursos de forma eficaz. Caso contrário, a maior parte ou toda a potência extra da GPU seria desperdiçada.
O artigo de pesquisa se aprofunda nos detalhes de como exatamente esses engenheiros de computação fizeram isso, mas a versão resumida é que eles iteraram em estruturas como Megatron-DeepSpeed e FSDP, mudando as coisas para que o programa de treinamento funcionasse de maneira mais otimizada no Frontier. No final, os resultados foram bastante impressionantes – a eficiência de escalonamento fraca ficou em 100%, o que basicamente significa que mais GPUs foram usadas da maneira mais eficiente possível com um tamanho de carga de trabalho crescente.
Enquanto isso, a forte eficiência de escalonamento foi ligeiramente inferior, 89% para o parâmetro LLM de 175 bilhões e 87% para o LLM de um trilhão de parâmetros. O escalonamento forte refere-se ao aumento da contagem de processadores sem alterar o tamanho da carga de trabalho, e é aí que contagens de núcleos mais altas se tornam menos úteis, de acordo com a lei de Amdahl. Mesmo 87% é um resultado decente, considerando quantas GPUs eles usaram.
No entanto, a equipe notou alguns problemas para alcançar essa eficiência no Frontier, afirmando que “é preciso haver mais trabalho para explorar o desempenho de treinamento eficiente em GPUs AMD, e a plataforma ROCm é escassa”. Como diz o artigo, a maior parte do aprendizado de máquina nessa escala é feita no ecossistema de hardware-software CUDA da Nvidia, o que deixa as soluções da AMD e da Intel subdesenvolvidas em comparação. Naturalmente, esforços como estes irão promover o desenvolvimento destes ecossistemas.
Mesmo assim, o supercomputador mais rápido do mundo continua a ser o Frontier, com seu hardware totalmente AMD. Em segundo lugar está o Aurora com seu hardware puramente Intel, incluindo GPUs, embora no momento apenas metade dele tenha sido usado para envios de benchmark. As GPUs Nvidia alimentam o terceiro supercomputador mais rápido, Eagle. Se a AMD e a Intel quiserem manter a classificação desta forma, as duas empresas precisarão acompanhar as soluções de software da Nvidia.