Reduzindo o tempo de treinamento do Qwen em 6x no AI Runtime

Treinar ou fazer o fine-tuning de um modelo base com dezenas de bilhões de parâmetros exige um poder de computação significativo, muitas vezes abrangendo de centenas a milhares de horas de treinamento em GPU. No entanto, otimizar o desempenho do treinamento de um LLM é desafiador, e métricas de alto nível, como a utilização da GPU, podem não mostrar sua utilização real. Neste post do blog, mostramos como ir além das camadas de abstração e acelerar o fine-tuning do modelo Qwen 3.8 27B mais recente no AI Runtime em 6x. Por meio de otimização sistemática de desempenho, mesmo com a utilização da GPU já em 98%, passamos de 7 horas e 44 minutos para 1 hora e 17 minutos de fine-tuning em um conjunto de dados fixo.

Apresentaremos uma análise detalhada das técnicas por trás dessas melhorias, desde como medir o que realmente importa até detalhes de implementação e compensações de design envolvidas em cada otimização.

Se quiser começar imediatamente, vá direto para o código de exemplo no AI Runtime e experimente!

Configuração do modelo

Começamos com um loop de treinamento padrão para o Qwen 3.8 27B usando a biblioteca transformers e o conjunto de dados de treinamento SFT HuggingFaceH4/ultrachat_200k.

• Aplicamos o FSDP2, a técnica mais recente de sharding de modelo do torch. Isso é necessário para ajustar os pesos do modelo e o estado do otimizador nas GPUs. Uma abordagem paralela de dados distribuídos pura falharia em ajustar os pesos do modelo, desse tamanho, em uma única GPU.
• Também ativamos o activation checkpointing, uma técnica que troca espaço de memória por computação na etapa retroativa (backwards pass) para permitir um tamanho de batch maior por GPU.
• O dataloader está empacotando batches 2D com um máximo de 8192 tokens por batch, preenchendo (padding) até o tamanho máximo dos exemplos do batch.

Benchmarks iniciais de treinamento

Nosso treinamento de referência do modelo em todo o conjunto de dados leva 7 horas e 44 minutos e, durante o treinamento, a GPU mostra uma utilização média de 98%. O resultado é um treinamento de 9.396 tok/s.

Essa utilização da GPU parece alta, mas, como mencionado anteriormente, esse número é enganoso. A métrica padrão de utilização da GPU conta apenas a porcentagem do tempo em que a GPU estava fazendo qualquer coisa, portanto, isso não nos diz o quão bem o modelo estava usando a GPU enquanto ela estava ativa. Para entender isso, precisamos olhar mais a fundo.

Uma maneira de fazer isso é com o mecanismo de rastreamento do torch; podemos ajustar nosso loop de treinamento da seguinte forma (enquanto também limitamos o número de etapas para obter uma resposta rápida):

Para uma etapa de exemplo, isso nos dá os seguintes dados:

Tabela 1: CPU - 14,4s com 1.714.382 chamadas. A tabela contém as principais operações por tempo total.

Tabela 2: GPU - 10,3s com 269.448 operações. A tabela contém as principais operações por tempo total

Áreas de melhoria

A partir dessas tabelas, notamos alguns problemas de alto nível no modelo atual:

• A CPU leva mais tempo para configurar cada etapa de treinamento do que a GPU leva para executá-la. Isso é inesperado, pois geralmente o lado da CPU na execução da GPU é assíncrono, ou seja, a CPU deve rodar bem à frente da GPU. O principal efeito disso é que a GPU fica ociosa (stalling), por isso temos uma oportunidade de melhorar isso.
• Há operações demais! Nas tabelas acima, temos centenas de milhares de pequenas operações; neste caso, o problema é que a CPU não consegue enfileirar as operações da GPU rápido o suficiente para mantê-la ocupada. A ideia principal é que precisamos reduzir o número de operações que a CPU enfileira para a GPU. Existem algumas técnicas gerais para isso, e uma se encaixa na nossa situação:
  • Os CUDA graphs do torch registram os kernels CUDA vistos e os reproduzem com uma única chamada de CPU. Embora sejam poderosos, existem muitas limitações e entram em conflito com o dinamismo do FSDP2. Portanto, por enquanto, isso está fora de cogitação.
  • O Torch.compile rastreia o código e usa o Triton para fundir kernels automaticamente quando possível. Neste exemplo, usaremos kernels fundidos manualmente e, no caso de LLMs, isso não deixa espaço para melhorias com o torch.compile.
  • Kernels fundidos manualmente usam técnicas de fusão especializadas para operações específicas que o compilador do torch não consegue aplicar por conta própria. Felizmente, existem muitos kernels pré-construídos já consolidados na comunidade. Exemplos disso incluem o Flash Attention, que oferece uma pequena melhoria no desempenho neste caso.

Além disso, temos um problema com excesso de preenchimento (padding) no batch: o tensor de entrada deve ser retangular e cada linha é preenchida até o número máximo de tokens para esse batch. Mostraremos um truque para LLMs que pode ser usado para minimizá-lo.

Nas seções a seguir, nos aprofundaremos em cada tópico técnico, cobrindo o que é, as melhores práticas recomendadas e as principais compensações a serem consideradas ao aplicá-lo.

Kernels fundidos

Você pode notar que alguns kernels ocorrem muito! Em particular, para o Qwen, o kernel de multiplicação de matrizes em batch executa mais de 30.000 vezes em uma única etapa.

Neste caso, exploraremos primeiro os kernels fundidos. Um kernel fundido pega vários kernels de GPU menores e os combina em uma única unidade. Isso pode ajudar no desempenho, evitando transferências dispendiosas entre a memória de alta largura de banda da GPU e sua memória on-die de largura de banda ainda maior. Além disso, também pode permitir um agendamento mais eficiente dos núcleos da GPU.

Com tantos kernels extras, embora o torch.compile possa funcionar, ele levará muito tempo para ser executado. Portanto, reduzir o número bruto de ops com kernels fundidos manualmente é um bom primeiro passo. Além disso, fusões conhecidas como Flash Attention não são algo que o torch.compile encontraria por conta própria.

Por padrão, a biblioteca transformers usa a SDPA do torch para as camadas de atenção padrão. Mas acontece que muitas dessas pequenas operações que vemos vêm das camadas não padrão de gated delta network. Temos a sorte de que, para essas camadas, existem dois kernels, causal-conv1d e flash-linear-attention, que foram escritos para fundir muitas dessas operações por camada. Como descobrimos isso? A implementação da biblioteca transformers do Qwen já nos mostra a quais kernels ela dá suporte:

• flash-linear-attention
• causal-conv1d

Existem dois métodos de alto nível para obter esses kernels:

• Instalar os kernels você mesmo. Isso pode variar entre:
  • Um simples pip install (por exemplo, para o flash-linear-attention) a partir do PyPi
  • Um git clone e compilação do repositório relevante
  • Encontrar outra fonte para um pacote pré-compilado, caso a compilação demore muito
• Usar o Kernel Hub da Hugging Face para buscar kernels pré-compilados.

Geralmente recomendamos a primeira opção, especialmente combinada com nosso suporte a contêineres Docker. Crie seu contêiner com todos os seus kernels instalados no momento do build para reduzir o tempo necessário para o início do seu job.

Embora o kernelhub seja conveniente, como os kernels geralmente dependem explicitamente das versões instaladas do torch e do cuda, pode não haver um kernel disponível para sua configuração específica, e o modelo de transformers usará silenciosamente um fallback para um kernel não otimizado com baixo desempenho.

Depois de instalar os pacotes Python flash-linear-attention e causal-conv1d, melhoramos nosso tempo de ponta a ponta em um fator de 3,7x, caindo para 2 horas e 5 minutos. Nas tabelas abaixo, você pode ver os novos kernels (em destaque) ocupando o lugar das funções elementwise nas tabelas anteriores. Além disso, o número total de operações executadas foi bastante reduzido, por um fator de 15x:

Tabela 3: CPU - 3,9s com 118.316 chamadas. A tabela contém as principais operações por tempo total.

Tabela 4: GPU - 3,9s com 18.060 chamadas. A tabela contém as principais operações por tempo total.

Aviso: durante o treinamento, você precisará de um kernel tanto para o passo forward quanto backward, e alguns kernels são apenas forward, pois foram implementados apenas para acelerar a inferência.

Principais conclusões

• Considere usar kernels altamente otimizados e bem estabelecidos que estejam prontamente disponíveis para modelos comuns. Aproveitar esses kernels pode proporcionar melhorias significativas de desempenho com relativamente pouco esforço. Neste caso específico, vimos uma melhoria de 3,7x na velocidade de treinamento, reduzindo nosso tempo de treinamento de 7 horas e 44 minutos para 2 horas e 5 minutos.

Batching, tokenização e padding

Nosso modelo Qwen está bem rápido agora, mas podemos deixá-lo ainda mais rápido. Por padrão, nossa configuração de batching usa tensores retangulares 2D de 8.192 tokens, pois esse é o maior tamanho que conseguimos encaixar sem ficar sem memória. A desvantagem dessa abordagem é que tendemos a desperdiçar uma quantidade considerável de espaço no tensor com padding.

Se o Qwen usasse apenas mecanismos de atenção padrão, o NestedTensor do torch agora seria bem suportado e ajudaria a reduzir esse custo removendo completamente o padding. Mas o flash-linear-attention suporta apenas uma abordagem de tensor compactado 1D mais tradicional:

Mas não podemos fazer isso de forma arbitrária. Apenas passar o tensor 1D para o modelo pode ser aceitável para o pré-treinamento, mas, para o ajuste fino (finetuning), precisamos garantir que os tokens de um exemplo atendam apenas aos tokens do mesmo exemplo. Para corrigir isso, temos que passar os position_ids de exemplo corretos para o modelo. Por sua vez, o modelo passará esses IDs para os kernels de atenção fundida, e eles aplicarão a máscara de atenção correta. Se você olhar o código de exemplo aqui, verá como calculamos as entradas necessárias para fazer isso funcionar.

Com essa mudança aparentemente simples, nosso tempo de treinamento cai de 2 horas e 5 minutos para 1 hora e 17 minutos.

Há dois motivos para isso:

1. O motivo principal é que esse empacotamento (packing) aumenta o número de sequências por lote e reduz o número total de lotes em aproximadamente 36%. Observe que isso depende da distribuição e do comprimento dos exemplos no seu conjunto de dados.
2. Um motivo secundário é que os kernels de flash linear attention precisam ser compilados just-in-time para cada formato de entrada (shape), e o formato anterior era dinâmico, levando a alguns minutos extras de compilação.

Compromisso (trade-off): Esse tipo de mudança torna o pipeline de entrada mais complexo, e há uma grande chance de introduzir erros que impactem a qualidade do modelo. Mas, neste caso, a economia de 36% vale a pena. Existem outras técnicas para reduzir o padding, como manter o mecanismo de batching original, mas ordenar todas as entradas por comprimento para que o padding seja minimizado. No entanto, essa ordem pode introduzir correlações inesperadas em amostras adjacentes, gerando viés no modelo.

Principais conclusões: Ao treinar LLMs em um conjunto de dados com entradas menores do que o máximo de tokens treináveis (por exemplo, 8192 neste caso), empacotar os dados em lotes com eficiência pode economizar um tempo significativo. Para este conjunto de dados, vimos uma melhoria de 1,6x na velocidade de treinamento, reduzindo nosso tempo de treinamento de 2 horas e 5 minutos para 1 hora e 17 minutos.

Conclusões

Neste blog, mostramos como aumentar a eficiência do seu treinamento em 6x. Resumiu-se a (1) primeiro medir o que importa, não apenas a utilização da GPU, mas também onde seu modelo está gastando tempo no treinamento, (2) otimizar e fundir kernels, usando bibliotecas conhecidas como Flash Linear Attention, e (3) ajustar como os lotes de dados são fornecidos ao modelo.

Para fazer isso em escala, você precisa de uma plataforma que torne o caminho de alto desempenho o padrão. No AIR, muitas dessas melhorias de desempenho já vêm integradas. Incluímos as versões mais recentes do CUDA e do flash attention em nossos ambientes base, monitoramos automaticamente métricas essenciais como utilização da GPU, temperatura e integridade para você, e garantimos que as versões do torch (e os profilers) estejam atualizadas. Nosso agente interno, o Genie Code, também é ajustado para procurar oportunidades como essas para otimizar ainda mais o seu treinamento. Além disso, a leitura de dados do Databricks é otimizada por meio do nosso UCVolumeDataset.

Experimente o AIR na sua próxima carga de trabalho de treinamento! Pronto para testar? Confira o cookbook do AI Runtime, com detalhes sobre como executar esse modelo você mesmo.

(Esta publicação no blog foi traduzida utilizando ferramentas baseadas em inteligência artificial) Publicação original