O carregador de dados da IBM Research melhora o treinamento de modelos de IA para a comunidade de código aberto

Como superar gargalos quando está treinando modelos de IA com grandes quantidades de dados? Nesta conferência PyTorch do ano, a IBM pesquisa apresentou um carregador de dados inovador para treinamento de LLM em grande escala LLM. A ferramenta, agora disponível para usuários do PyTorch, visa simplificar o treinamento em grande escala para uma audiência o mais ampla possível.

A ideia do carregador de dados de alta taxa de transferência surgiu de questões práticas que os cientistas de pesquisa observaram durante o treinamento de modelos, pois seu trabalho exigia uma ferramenta que pudesse processar grandes quantidades de dados em vários dispositivos, tudo isso enquanto acompanhava GPUs progressivamente eficientes. Como IBM Research observa em seu blog sobre o lançamento, “Tudo isso é graças a uma equipe de pesquisadores que estavam simplesmente construindo as ferramentas de que precisavam para realizar um trabalho”.

Davis Wertheimer, da IBM Research, explica alguns dos desafios que podem surgir durante o treinamento em grande escala: "Quando se trata de treinamento em grande escala, existe uma espécie de regra 80/20. Oitenta por cento de toda a literatura publicada analisa os compromissos algorítmicos entre a memória da GPU e a comunicação e a computação. Mas quando você realmente tenta construir algo, 80% do tempo, você pode contar com uma fila muito longa de todas essas outras questões práticas, porque o pipeline funciona na velocidade do gargalo mais estreito."

À medida que a equipe da IBM desenvolvia sua plataforma de treinamento, eles continuaram encontrando gargalos. “À medida que melhoramos cada vez mais o uso de nossas GPUs, cada vez mais o gargalo é o carregador de dados”, observa Wertheimer.

Essa percepção levou a um processo de desenvolvimento duplo. "Houve uma jornada paralela de, por um lado, evoluir nossa plataforma de treinamento e, por outro lado, evoluir constantemente nosso carregador de dados para acompanhar as demandas de velocidade de nossa plataforma de treinamento para evitar gargalos", explica.

Linsong Chu, da IBM Research, descreve os recursos essenciais do carregador de dados:

Com estado e ponto de verificação: "Sempre que você salva um modelo, o estado do carregador de dados também é salvo e sempre que você se recupera de um ponto de verificação, o estado do modelo e os estados do carregador de dados precisam ser recuperados ao mesmo tempo", diz Chu.

Redimensionamento automático de pontos de verificação: o carregador de dados se ajusta automaticamente às mudanças na carga de trabalho durante sessões de treinamento estendidas. "O treinamento pode facilmente levar semanas ou meses, e há muitos motivos pelos quais você pode ter que redimensionar sua carga de trabalho no meio", observa Chu.

Fluxo de dados eficiente: o sistema oferece suporte ao fluxo de dados com sobrecarga de compilação zero para embaralhamento.

Operação distribuída assíncrona: "Queremos que o carregador de dados não bloqueie", explica Chu. "Ao salvar o estado do carregador de dados, queremos que a economia seja distribuída de uma forma que não haja comunicação zero."

Combinação dinâmica de dados: o carregador de dados pode se adaptar a diferentes proporções de combinação de dados, o que é útil para as necessidades de treinamento em evolução.

Embaralhamento global eficiente: a ferramenta aborda gargalos de memória ao lidar com grandes conjuntos de dados, tornando o embaralhamento eficiente mesmo à medida que os dados crescem.

PyTorch nativo, modular e extensivo: projetado para adaptabilidade e escalabilidade, o carregador de dados está preparado para o crescimento futuro. “E se no próximo ano tivermos que lidar com 30 trilhões, 50 trilhões ou 100 trilhões de tokens?” pergunta Chu. "O mundo está mudando rapidamente, então precisamos construir o carregador de dados para que ele possa não apenas sobreviver hoje, mas também sobreviver para amanhã."

A equipe da IBM Research testou rigorosamente seu carregador de dados ao longo de vários meses, executando centenas de trabalhos pequenos e grandes. Eles observaram números de código estáveis e suaves. Além disso, todo o carregador de dados opera de forma assíncrona e sem bloqueio.

"Aproveitamos muitos recursos integrados do PyTorch para que tudo isso acontecesse", diz Wertheimer. “É por isso que estamos contribuindo, contribuindo de volta.”