SWattention: desenhando atenção rápida e eficiente em memória para um novo Supercomputador Sunway

Resumo Nos últimos anos, modelos de linguagem grandes (LLM) baseados em Transformer tornaram-se a tecnologia dominante em uma série de aplicações. Para ampliar o comprimento da sequência do Transformer, o FlashAttention foi proposto para calcular a atenção exata com requisitos de memória reduzidos e execução mais rápida. No entanto, implementar o algoritmo FlashAttention na nova geração do Supercomputador Sunway enfrenta muitas restrições, como a arquitetura heterogênea única e a largura de banda de memória limitada. Este trabalho propõe SWattention, um método altamente eficiente para calcular a atenção exata no processador SW26010pro. Para utilizar totalmente os 6 grupos de núcleos (CG) e 64 núcleos por CG no processador, projetamos uma estratégia de particionamento de tarefas paralelas em dois níveis. O acesso assíncrono à memória é empregado para garantir que o acesso à memória coincida com o cálculo. Além disso, uma estratégia de tiling é introduzida para determinar tamanhos ótimos de blocos SRAM. Comparado à atenção padrão, o SWattention alcança cerca de 2,0x de aceleração para treinamento FP32 e 2,5x para treinamento de precisão mista. Os comprimentos das sequências variam de 1k a 8k e escalam até 16k sem falta de memória. Quanto ao desempenho de ponta a ponta, o SWattention alcança até 1,26x de aceleração para treinamento de modelos estilo GPT, o que demonstra que o SWattention possibilita sequências mais longas para o treinamento de LLM.