FFN Fusion da NVIDIA Acelera Eficiência de LLMs

O Equilíbrio Computacional da IA Moderna

Os modelos de linguagem grandes (LLMs) destacam-se como pilares da inteligência artificial contemporânea, demonstrando capacidades notáveis que estão a remodelar indústrias e a descoberta científica. A sua proficiência em gerar texto semelhante ao humano, potenciar agentes conversacionais sofisticados e até auxiliar em tarefas de investigação complexas tornou-os ferramentas indispensáveis. No coração destes modelos poderosos bate a arquitetura transformer, um design caracterizado pelas suas camadas alternadas. Os dados de entrada, divididos em tokens, fluem através de uma sequência de mecanismos de atenção, que ponderam a importância de diferentes tokens, seguidos por redes feed-forward (FFNs), que processam a informação recolhida. Este processamento em camadas e sequencial é fundamental para a forma como os transformers aprendem e geram resultados.

No entanto, esta mesma arquitetura, embora eficaz, apresenta um desafio crescente à medida que os modelos aumentam em tamanho e complexidade. A natureza sequencial significa que cada camada geralmente tem de esperar que a anterior conclua a sua computação antes de poder começar. Este processamento passo a passo cria um gargalo inerente, particularmente durante a fase de inferência – a etapa em que um modelo treinado é efetivamente usado para gerar previsões ou texto. À medida que modelos como os que potenciam assistentes de IA avançados incorporam centenas de biliões, ou mesmo triliões, de parâmetros, os recursos computacionais e o tempo necessários para a inferência aumentam drasticamente. Esta procura crescente traduz-se em latência significativa (atraso na resposta), redução do throughput (número de pedidos tratados ao longo do tempo) e custos operacionais crescentes, dificultando a implementação generalizada e a aplicação em tempo real dos LLMs mais poderosos. Consequentemente, melhorar a eficiência da inferência tornou-se uma preocupação primordial na comunidade de investigação em IA, estimulando uma busca por estratégias inovadoras que possam otimizar a computação sem comprometer o desempenho notável que estes modelos oferecem. O desafio central reside em mitigar as restrições impostas pela execução sequencial, especialmente em ambientes distribuídos onde as computações abrangem múltiplas GPUs, adicionando sobrecarga de comunicação ao tempo de processamento.

Navegando pelo Cenário da Otimização: Ferramentas Existentes e Seus Limites

No esforço contínuo para tornar os LLMs mais eficientes e rápidos, os investigadores desenvolveram um conjunto de técnicas de otimização. Cada uma oferece um caminho para a eficiência, mas muitas vezes vem com o seu próprio conjunto de compromissos, impedindo que qualquer método único seja uma solução universal. Compreender estas trocas é crucial para apreciar a necessidade de abordagens inovadoras como o FFN Fusion.

Uma técnica proeminente é a quantização. Esta envolve a redução da precisão numérica usada para representar os pesos e ativações do modelo. Em vez de usar números de ponto flutuante padrão de 32 bits, os modelos podem usar representações de 16 bits, 8 bits ou até mesmo de bits inferiores. Isto reduz diretamente a pegada de memória do modelo e pode acelerar significativamente os cálculos, uma vez que as operações em números de menor precisão são tipicamente mais rápidas e requerem menos energia. No entanto, a quantização não está isenta de riscos. Reduzir a precisão pode levar a uma perda de informação, potencialmente degradando a precisão do modelo. Este risco torna-se mais pronunciado em larguras de bits muito baixas, exigindo implementação cuidadosa e, por vezes, retreinamento para mitigar quedas de precisão. O desafio reside em encontrar o ponto ideal que maximiza os ganhos de eficiência, mantendo a degradação do desempenho dentro de limites aceitáveis.

Outra estratégia comum é a poda (pruning). Esta técnica opera sob o princípio de que muitos parâmetros dentro de uma grande rede neural podem ser redundantes ou contribuir minimamente para o resultado final. Os algoritmos de poda identificam e removem estas conexões ou neurónios menos importantes, resultando num modelo menor e mais esparso. Tal como a quantização, a poda reduz os requisitos de memória e a carga computacional. No entanto, identificar precisamente quais parâmetros são ‘seguros’ para remover é complexo. A poda agressiva pode inadvertidamente remover componentes cruciais, levando a uma perda substancial de precisão. O ajuste fino (fine-tuning) do modelo após a poda é frequentemente necessário para recuperar o desempenho, adicionando complexidade ao fluxo de trabalho. A calibração cuidadosa é essencial para garantir que o modelo podado permaneça eficaz.

Uma abordagem arquitetonicamente distinta é o modelo Mixture-of-Experts (MoE). Em vez de processar cada entrada através de toda a rede, os modelos MoE consistem em múltiplas sub-redes ‘especialistas’ (tipicamente FFNs). Para cada token de entrada, um mecanismo de gating seleciona dinamicamente um pequeno subconjunto destes especialistas para realizar a computação. Esta computação condicional significa que apenas uma fração dos parâmetros totais do modelo é ativada para qualquer entrada dada, levando a economias computacionais significativas, especialmente durante o treino e a inferência em modelos muito grandes. Os modelos MoE podem escalar para triliões de parâmetros, mantendo custos computacionais razoáveis. No entanto, a sua eficiência depende muito da carga de trabalho. Eles destacam-se no tratamento de tamanhos de lote muito grandes, onde o padrão de ativação seletiva leva a uma boa utilização do hardware. Em tamanhos de lote menores ou intermédios, os modelos MoE podem sofrer de subutilização dos recursos computacionais, uma vez que o hardware paralelo pode não ser mantido consistentemente ocupado pelos especialistas ativados esparsamente. Além disso, implementar e balancear a carga de modelos MoE pode ser mais complexo do que implementar arquiteturas ‘densas’ padrão.

Embora a quantização, a poda e os modelos MoE representem avanços valiosos na otimização de LLMs, as suas limitações inerentes destacam a necessidade de estratégias alternativas ou complementares. A busca continua por métodos que possam proporcionar melhorias amplas de eficiência em vários cenários, idealmente com menos compromissos em termos de precisão ou complexidade de implementação, particularmente para as arquiteturas de modelos densos que permanecem populares devido à sua relativa simplicidade no treino e implementação.

FFN Fusion: Repensando o Paralelismo em Transformers

No meio deste cenário de técnicas de otimização, investigadores da NVIDIA introduziram uma nova abordagem convincente denominada FFN Fusion. Esta técnica confronta diretamente o gargalo sequencial inerente à arquitetura transformer, não alterando parâmetros ou ativando seletivamente partes, mas repensando fundamentalmente como sequências de computações podem ser paralelizadas. A inovação decorre de uma observação crucial sobre o comportamento das camadas FFN dentro de modelos transformer profundos.

Usando uma ferramenta de diagnóstico chamada Puzzle, os investigadores analisaram o funcionamento interno de modelos grandes. Quando removeram experimentalmente as camadas de atenção, notaram que os modelos frequentemente mantinham sequências surpreendentemente longas de camadas FFN consecutivas. Mais importante, a análise revelou que as computações realizadas por estas FFNs adjacentes frequentemente exibiam interdependência mínima. Essencialmente, a saída de uma FFN na sequência muitas vezes não alterava drasticamente o caminho direcional ou a informação central necessária pela FFN imediatamente seguinte. Isto sugeriu que estas FFNs, tradicionalmente executadas uma após a outra, poderiam possuir o potencial para execução simultânea e paralela sem perturbar significativamente a função geral do modelo.

Esta perceção formou a base do FFN Fusion. A ideia central é elegantemente simples, mas poderosa: identificar sequências de camadas FFN consecutivas com baixa dependência computacional e fundi-las numa única camada FFN mais larga que realiza a computação equivalente em paralelo. Em vez de uma cadeia como Entrada -> FFN1 -> FFN2 -> FFN3 -> Saída, a estrutura fundida torna-se Entrada -> FFN_Fundida (Equivalente a FFN1+FFN2+FFN3 em paralelo) -> Saída. Esta transformação arquitetónica encurta efetivamente a profundidade sequencial da rede, substituindo múltiplos passos por um único passo computacional mais amplo. Ao visar estas sequências FFN de baixa dependência, o FFN Fusion visa reduzir a latência e o custo computacional, preservando ao mesmo tempo o poder representacional e a precisão do modelo. O desenvolvimento do Ultra-253B-Base a partir do Llama-3.1-405B-Instruct serviu como uma demonstração primordial do potencial desta técnica.

A Alquimia Arquitetônica: Como Funciona o FFN Fusion

A magia por trás do FFN Fusion reside na sua manipulação inteligente da estrutura matemática subjacente das redes feed-forward. Não se trata apenas de executar camadas existentes lado a lado; envolve a criação de uma nova camada unificada que replica o comportamento coletivo da sequência original, mas fá-lo concorrentemente.

Considere uma sequência de k camadas FFN consecutivas. Num transformer padrão, a entrada x passa pela FFN1, a sua saída torna-se a entrada para a FFN2, e assim por diante, até FFNk. Cada passo depende explicitamente da conclusão do anterior. O FFN Fusion quebra esta cadeia de dependência. Matematicamente, uma FFN envolve tipicamente duas transformações lineares com uma função de ativação não linear (como GeLU ou SwiGLU) entre elas: FFN(x) = W_out * Ativação(W_in * x). O FFN Fusion aproveita o facto de que as transformações lineares podem muitas vezes ser combinadas.

O processo de fusão funciona através da concatenação dos pesos das camadas FFN individuais. Especificamente, as matrizes de pesos de entrada (W_in) das FFNs consecutivas são combinadas (por exemplo, em bloco diagonal) numa única matriz de pesos de entrada maior para a camada fundida. Da mesma forma, as matrizes de pesos de saída (W_out) são concatenadas para formar uma única matriz de pesos de saída mais larga. A função de ativação é aplicada elemento a elemento dentro desta estrutura maior. Esta construção garante que a FFN fundida opera sobre a entrada original x simultaneamente através de caminhos paralelos correspondentes às FFNs originais. As saídas destes caminhos paralelos são então implicitamente agregadas pela estrutura dos pesos de saída concatenados.

A base teórica confirma que esta estrutura fundida pode manter a mesma capacidade representacional que a sequência original de FFNs, desde que as dependências entre as camadas originais fossem de facto baixas. A chave é identificar quais sequências são adequadas para fusão. Para fazer isso sistematicamente, os investigadores da NVIDIA empregaram uma técnica de análise de dependência. Eles mediram a distância cosseno entre os estados ocultos de saída de camadas FFN consecutivas para um conjunto representativo de tokens de entrada. Uma pequena distância cosseno indica que o vetor de saída de uma FFN aponta numa direção muito semelhante ao vetor de saída da próxima FFN na sequência. Esta similaridade sugere baixa dependência funcional – a segunda FFN não está a mudar drasticamente a representação da informação estabelecida pela primeira. Sequências de FFNs exibindo distâncias cosseno consistentemente baixas entre camadas foram identificadas como candidatas principais para fusão, pois fundi-las era menos provável de perturbar as representações aprendidas do modelo e o desempenho geral. Esta abordagem orientada por dados permite a aplicação direcionada do FFN Fusion às partes do modelo onde será mais eficaz e menos disruptivo.

De Behemoth a Sprinter: A Transformação Ultra-253B-Base

O poder prático do FFN Fusion foi vividamente demonstrado através da sua aplicação a um dos maiores modelos publicamente conhecidos na altura, o Llama-3.1-405B-Instruct. Este modelo, ostentando 405 biliões de parâmetros, representava um empreendimento computacional significativo para a inferência. Os investigadores embarcaram num processo de refinamento arquitetónico, combinando FFN Fusion com poda estratégica, para criar um novo modelo mais eficiente apelidado de Ultra-253B-Base.

O processo de transformação envolveu vários passos:

1. Análise: Usando as suas ferramentas de análise de dependência (medindo distâncias cosseno), os investigadores identificaram sequências de camadas FFN consecutivas dentro da arquitetura Llama-405B que exibiam baixa dependência inter-camadas.
2. Fusão: Estas sequências FFN identificadas foram então fundidas em camadas FFN únicas e mais largas, como descrito anteriormente (concatenando pesos). Isto reduziu diretamente o número de passos sequenciais na rede.
3. Poda (Pruning): Concorrentemente ou subsequentemente, parâmetros considerados menos críticos (potencialmente identificados através de técnicas de poda padrão ou informados pelo processo de fusão) foram removidos do modelo.

Esta abordagem combinada resultou no Ultra-253B-Base, um modelo com 253 biliões de parâmetros. Isto representa uma redução substancial – mais de 37% menos parâmetros do que o modelo original de 405B. As alterações arquitetónicas alcançadas através da fusão foram chave para permitir uma redução de tamanho tão significativa, visando ao mesmo tempo reter o desempenho. O objetivo não era apenas um modelo menor, mas um fundamentalmente mais rápido e computacionalmente mais frugal, graças ao aumento do paralelismo desbloqueado pelo FFN Fusion. Este estudo de caso serviu como uma prova de conceito crucial, mostrando que modelos de grande escala poderiam ser substancialmente reestruturados para eficiência.

Medindo os Ganhos: Desempenho, Velocidade e Economia de Recursos

O verdadeiro teste de qualquer técnica de otimização reside no seu impacto mensurável. Para o Ultra-253B-Base, os resultados derivados da aplicação do FFN Fusion e da poda à base Llama-405B foram convincentes, demonstrando melhorias significativas em múltiplas dimensões sem compromissos substanciais na capacidade.

Velocidade e Custo de Inferência: Os ganhos mais notáveis foram observados na eficiência da inferência. Comparado com o modelo original de 405B parâmetros, o Ultra-253B-Base alcançou:

• Uma melhoria de 1.71x na latência de inferência. Isto significa que o modelo podia gerar respostas significativamente mais rápido, crucial para aplicações em tempo real.
• Uma redução de 35x no custo computacional por token quando medido com um tamanho de lote de 32. Esta diminuição dramática nas operações computacionais (FLOPs) por token traduz-se diretamente em menor consumo de energia e requisitos de hardware reduzidos para servir o modelo.

Benchmarks de Desempenho do Modelo: Criticamente, estas melhorias de eficiência não vieram à custa da inteligência ou capacidades do modelo. O Ultra-253B-Base foi rigorosamente avaliado num conjunto de benchmarks padrão de LLM, alcançando pontuações que foram altamente competitivas com, e em alguns casos excederam, o modelo original, muito maior:

• MMLU (Massive Multitask Language Understanding): 85.17%
• MMLU-Pro (Uma versão mais desafiadora): 72.25%
• Arena Hard (Avaliação de preferência humana em prompts difíceis): 84.92%
• HumanEval (Capacidade de geração de código): 86.58%
• MT-Bench (Qualidade de conversação multi-turno): 9.19

Estas pontuações indicam que o modelo fundido e podado reteve um nível muito alto de compreensão, raciocínio, capacidade de codificação e qualidade conversacional, comparável ao seu progenitor de 405B parâmetros, apesar de ter apenas 253 biliões de parâmetros.

Eficiência de Memória: Além da velocidade e custo computacional, o FFN Fusion também contribuiu para a economia de memória. As alterações arquitetónicas, potencialmente combinadas com outras otimizações possibilitadas pela fusão, levaram a uma redução de 2x no tamanho do cache chave-valor (KV cache) necessário durante a inferência. O KV cache armazena ativações intermédias (chaves e valores de atenção) e pode consumir memória GPU substancial, especialmente para sequências de entrada longas. Reduzir este requisito para metade torna viável executar o modelo em hardware com menos memória intensiva ou processar contextos mais longos dentro das mesmas restrições de memória.

Estes resultados quantificáveis sublinham a eficácia do FFN Fusion. Permitiu a criação de um modelo que não era apenas menor, mas fundamentalmente mais eficiente em termos de velocidade, operações computacionais e uso de memória, tudo isto mantendo um desempenho de topo em benchmarks desafiadores.

Preservando o Conhecimento: O Papel Crucial do Treinamento e Fine-Tuning

Modificar arquitetonicamente um modelo de linguagem pré-treinado massivo como o Llama-405B através de técnicas como FFN Fusion e poda inevitavelmente perturba o delicado equilíbrio dos seus parâmetros aprendidos. Embora a equivalência matemática vise preservar a função localmente, o comportamento global da rede pode mudar. Para garantir que o modelo Ultra-253B-Base resultante não só se tornasse mais eficiente, mas também retivesse o seu alto nível de desempenho, um processo de treino pós-modificação cuidadosamente orquestrado foi essencial.

Este processo envolveu duas fases principais:

1. Destilação de Conhecimento: O primeiro passo foi transferir o conhecimento do modelo original, maior (ou um modelo professor adequado) de volta para a arquitetura modificada. Isto foi alcançado através da destilação, onde o modelo Ultra-253B-Base foi treinado para imitar as saídas ou representações internas do modelo professor. Esta fase utilizou um conjunto de dados substancial, especificamente 54 biliões de tokens, processados com uma janela de contexto de 8k. A destilação ajuda o modelo fundido e podado a recapturar nuances e capacidades que podem ter sido ligeiramente perturbadas durante as alterações arquitetónicas.

2. Fine-Tuning Faseado: Após a destilação, o modelo passou por uma série de fases de ajuste fino (fine-tuning) especificamente projetadas para adaptá-lo ao manuseamento de comprimentos de contexto progressivamente mais longos. Isto é crucial para os LLMs modernos, que são frequentemente esperados para processar e gerar texto com base em entradas extensas. O fine-tuning prosseguiu em fases:

   • Fine-tuning numa janela de contexto de 16k.
   • Fine-tuning adicional numa janela de contexto de 32k.
   • Fase final de fine-tuning numa janela de contexto de 128k.

Esta abordagem faseada permite que o modelo adapte gradualmente os seus parâmetros, incluindo as camadas FFN recém-formadas e os mecanismos otimizados de KV cache, para gerir eficazmente dependências e fluxo de informação ao longo de sequências muito longas. Cada fase baseia-se na anterior, garantindo estabilidade e desempenho robusto em diferentes tamanhos de contexto.

Este regime de treino meticuloso, combinando destilação em larga escala com fine-tuning faseado de longo contexto, foi instrumental para colmatar a lacuna entre a eficiência arquitetónica e o desempenho de alta fidelidade. Garantiu que os benefícios de velocidade, custo e memória proporcionados pelo FFN Fusion não comprometessem a precisão e as capacidades do modelo em benchmarks exigentes.

Horizontes Mais Amplos: Generalização e Direções Futuras

A transformação bem-sucedida do Llama-405B no Ultra-253B-Base fornece fortes evidências do potencial do FFN Fusion, mas o seu verdadeiro valor reside na sua aplicabilidade mais ampla e nas perspetivas que oferece para o design futuro de LLMs. A investigação demonstrou que isto não foi apenas um truque único aplicável apenas a modelos enormes.

Validação em Diferentes Escalas: Os investigadores da NVIDIA testaram explicitamente a metodologia FFN Fusion em modelos de vários tamanhos. Aplicaram com sucesso a técnica a modelos de 70B parâmetros, alcançando ganhos de eficiência semelhantes em relação aos seus homólogos originais. Também relataram validação numa escala de 49B, reforçando ainda mais a ideia de que a independência das FFNs e o potencial para fusão não são características exclusivas dos maiores modelos, mas podem ser uma propriedade mais geral da arquitetura transformer, potencialmente tornando-se mais pronunciada em escalas maiores onde sequências FFN mais profundas ocorrem naturalmente. Isto sugere que o FFN Fusion pode tornar-se uma ferramenta padrão no arsenal de otimização de LLMs, aplicável a uma gama de tamanhos de modelo.

FFN vs. Fusão de Bloco Completo: A investigação também lançou luz sobre o papel específico das camadas FFN em comparação com as camadas de atenção dentro do bloco transformer. Enquanto camadas FFN consecutivas frequentemente mostravam baixa dependência, tornando-as ideais para fusão, tentativas de paralelizar blocos transformer inteiros (incluindo tanto as camadas de atenção como as FFN) provaram ser mais desafiadoras. A análise indicou interdependências mais fortes envolvendo os mecanismos de atenção. Fundir blocos inteiros simultaneamente resultou numa degradação de desempenho mais significativa, sugerindo que as camadas de atenção desempenham um papel mais crítico e sequencialmente dependente na integração da informação através dos tokens. Esta descoberta ajuda a delinear os limites da paralelização eficaz – sequências FFN são terreno fértil, enquanto os mecanismos de atenção podem exigir estratégias de otimização diferentes.

Implicações para a Arquitetura de LLMs: O FFN Fusion oferece mais do que apenas uma técnica de otimização post-hoc; fornece informações valiosas para projetar futuros LLMs. A descoberta de que sequências de FFNs podem muitas vezes ser tratadas como unidades paraleláveis desafia a suposição estritamente sequencial que muitas vezes sustenta o design do transformer. Isto poderia inspirar novas arquiteturas que são inerentemente mais amigáveis ao paralelismo desde o início. Modelos futuros podem ser projetados com estruturas FFN explicitamente destinadas à fusão ou execução paralela, potencialmente levando a co-design hardware-software onde as arquiteturas GPU são ainda mais otimizadas para explorar este tipo de paralelismo. O método sistemático usando distância cosseno para quantificar a dependência inter-camadas também fornece uma ferramenta analítica valiosa para entender e redesenhar estruturas de redes neurais. Ao demonstrar que ganhos significativos de eficiência são possíveis através de um redesenho arquitetónico ponderado focado na paralelização de componentes existentes, o FFN Fusion abre caminho para o desenvolvimento de LLMs que são simultaneamente poderosos e mais computacionalmente sustentáveis. Destaca um caminho para mitigar as crescentes exigências de recursos da IA de ponta.