Aceleração da Inferência em LLM por Meio de Fusão Eficiente de Operações

O rápido desenvolvimento dos Modelos de Linguagem de Grande Porte (LLMs) baseados em Transformer nos últimos anos está intimamente ligado ao seu tamanho cada vez maior e já enorme. Muitos LLMs contêm centenas de bilhões de parâmetros e requerem recursos de hardware dedicados para treinamento e inferência. Um dos principais desafios inerentes à arquitetura Transformer é a necessidade de suportar inúmeras transformações não lineares que envolvem normalização. Por exemplo, cada bloco do decodificador normalmente contém pelo menos uma operação Softmax e duas Layernorm. O cálculo dos fatores de escala correspondentes à normalização torna-se um grande gargalo, pois exige operações coletivas espaciais. Em outras palavras, no cálculo dos denominadores para Softmax e Layernorm, todos os elementos do vetor devem ser agregados em um único local, exigindo comunicação significativa. Essas operações coletivas desaceleram a inferência em Transformers em aproximadamente 20%, prejudicando o objetivo principal do cálculo distribuído em memória. Neste trabalho, propomos uma técnica extremamente eficiente que pode ocultar completamente a sobrecarga causada por tais operações coletivas. Note que cada operação Softmax e Layernorm é tipicamente seguida por uma camada linear. Como operações não lineares e lineares são realizadas em motores de hardware diferentes, elas podem ser facilmente paralelizadas uma vez que a álgebra permita tal comutação. Aproveitando as propriedades inerentes das operações lineares, podemos adiar a normalização do Softmax e Layernorm precedentes até depois do cálculo da camada linear. Agora podemos computar os fatores coletivos de escala concomitantemente com a multiplicação matricial e ocultar completamente a latência da primeira atrás da segunda. Tal paralelização preserva a precisão numérica enquanto melhora significativamente a utilização do hardware e reduz a latência geral.