Modelos de raciocínio de grandes linguagens (LLMs) são projetados para resolver problemas complexos, dividindo-os em uma série de etapas menores. Esses modelos poderosos são particularmente bons em tarefas desafiadoras, como programação avançada e planejamento em várias etapas.
Mas o desenvolvimento de modelos de raciocínio exige uma enorme quantidade de computação e energia devido a ineficiências no processo de treinamento. Enquanto alguns dos processadores de alta potência trabalham continuamente em consultas complicadas, outros do grupo ficam ociosos.
Pesquisadores do MIT e de outros lugares encontraram uma maneira de usar esse tempo de inatividade computacional para acelerar com eficiência o treinamento do modelo de raciocínio.
Seu novo método treina automaticamente um modelo menor e mais rápido para prever os resultados do LLM de raciocínio maior, que o modelo maior verifica. Isso reduz a quantidade de trabalho que o modelo de raciocínio deve realizar, acelerando o processo de treinamento.
A chave para este sistema é a sua capacidade de treinar e implantar o modelo menor de forma adaptativa, de modo que ele entre em ação somente quando alguns processadores estiverem ociosos. Ao aproveitar recursos computacionais que de outra forma seriam desperdiçados, acelera o treinamento sem incorrer em sobrecarga adicional.
Quando testado em vários LLMs de raciocínio, o método dobrou a velocidade de treinamento, preservando a precisão. Isto poderia reduzir o custo e aumentar a eficiência energética do desenvolvimento de LLMs avançados para aplicações como a previsão de tendências financeiras ou a detecção de riscos nas redes eléctricas.
“As pessoas querem modelos que possam lidar com tarefas mais complexas. Mas se esse é o objetivo do desenvolvimento de modelos, então precisamos priorizar a eficiência. Encontramos uma solução sem perdas para esse problema e, em seguida, desenvolvemos um sistema full-stack que pode fornecer acelerações bastante drásticas na prática”, diz Qinghao Hu, pós-doutorado no MIT e co-autor principal de um livro. artigo sobre esta técnica.
Ele é acompanhado no artigo pelo co-autor principal Shang Yang, um estudante de graduação em engenharia elétrica e ciência da computação (EECS); Junxian Guo, estudante de pós-graduação do EECS; autor sênior Music Han, professor associado do EECS, membro do Laboratório de Pesquisa de Eletrônica e ilustre cientista da NVIDIA; bem como outros na NVIDIA, ETH Zurich, MIT-IBM Watson AI Lab e na Universidade de Massachusetts em Amherst. A pesquisa será apresentada na Conferência Internacional ACM sobre Suporte Arquitetural para Linguagens de Programação e Sistemas Operacionais.
Gargalo de treinamento
Os desenvolvedores desejam LLMs de raciocínio para identificar e corrigir erros em seu processo de pensamento crítico. Esse recurso permite que eles respondam a consultas complicadas que atrapalhariam um LLM padrão.
Para ensinar-lhes essa habilidade, os desenvolvedores treinam LLMs de raciocínio usando uma técnica chamada aprendizagem por reforço (RL). O modelo gera múltiplas respostas potenciais para uma consulta, recebe uma recompensa para o melhor candidato e é atualizado com base na resposta principal. Essas etapas são repetidas milhares de vezes à medida que o modelo aprende.
Mas os pesquisadores descobriram que o processo de geração de múltiplas respostas, denominado rollout, pode consumir até 85% do tempo de execução necessário para o treinamento de RL.
“Atualizar o modelo – que é a verdadeira parte do ‘treinamento’ – consome muito pouco tempo em comparação”, diz Hu.
Esse gargalo ocorre em algoritmos RL padrão porque todos os processadores no grupo de treinamento devem terminar suas respostas antes de passarem para a próxima etapa. Como alguns processadores podem estar trabalhando em respostas muito longas, outros que geraram respostas mais curtas aguardam que elas terminem.
“Nosso objetivo period transformar esse tempo ocioso em aceleração sem nenhum desperdício de custos”, acrescenta Hu.
Eles procuraram usar uma técnica existente, chamada decodificação especulativa, para acelerar as coisas. A decodificação especulativa envolve treinar um modelo menor chamado redator para adivinhar rapidamente os resultados futuros do modelo maior.
O modelo maior verifica as suposições do redator e as respostas que ele aceita são usadas para treinamento.
Como o modelo maior pode verificar todas as suposições do redator de uma só vez, em vez de gerar cada resultado sequencialmente, ele acelera o processo.
Uma solução adaptativa
Mas na decodificação especulativa, o modelo do redator normalmente é treinado apenas uma vez e permanece estático. Isso inviabiliza a técnica de aprendizagem por reforço, uma vez que o modelo de raciocínio é atualizado milhares de vezes durante o treinamento.
Um redator estático rapidamente se tornaria obsoleto e inútil após algumas etapas.
Para superar esse problema, os pesquisadores criaram um sistema flexível conhecido como “Taming the Lengthy Tail” ou TLT.
A primeira parte do TLT é um treinador de rascunho adaptativo, que usa o tempo livre em processadores ociosos para treinar o modelo de rascunho em tempo actual, mantendo-o bem alinhado com o modelo alvo sem usar recursos computacionais extras.
O segundo componente, um mecanismo de implementação adaptativo, gerencia a decodificação especulativa para selecionar automaticamente a estratégia perfect para cada novo lote de insumos. Esse mecanismo altera a configuração de decodificação especulativa com base nos recursos da carga de trabalho de treinamento, como o número de entradas processadas pelo modelo preliminar e o número de entradas aceitas pelo modelo de destino durante a verificação.
Além disso, os pesquisadores projetaram o modelo preliminar para ser leve, para que pudesse ser treinado rapidamente. O TLT reutiliza alguns componentes do processo de treinamento do modelo de raciocínio para treinar o redator, levando a ganhos extras em aceleração.
“Assim que alguns processadores terminam suas consultas curtas e ficam ociosos, nós imediatamente os alternamos para fazer o treinamento do modelo preliminar usando os mesmos dados que estão usando para o processo de implementação. O mecanismo principal é a nossa decodificação especulativa adaptativa – esses ganhos não seriam possíveis sem ela”, diz Hu.
Eles testaram o TLT em vários LLMs de raciocínio que foram treinados usando conjuntos de dados do mundo actual. O sistema acelerou o treinamento entre 70 e 210 por cento, preservando a precisão de cada modelo.
Como um bônus adicional, o modelo de pequeno rascunho poderia ser facilmente utilizado para implantação eficiente como um subproduto gratuito.
No futuro, os pesquisadores desejam integrar o TLT em mais tipos de estruturas de treinamento e inferência e encontrar novas aplicações de aprendizagem por reforço que possam ser aceleradas usando esta abordagem.
“À medida que o raciocínio continua a se tornar a principal carga de trabalho que impulsiona a demanda por inferência, o TLT de Qinghao é um excelente trabalho para lidar com o gargalo computacional do treinamento desses modelos de raciocínio. Acho que esse método será muito útil no contexto da computação de IA eficiente”, diz Han.
Este trabalho é financiado pelo MIT-IBM Watson AI Lab, pelo MIT AI {Hardware} Program, pelo MIT Amazon Science Hub, pela Hyundai Motor Firm e pela Nationwide Science Basis.