Nos últimos anos, os grandes modelos de linguagem (LLM) baseados na arquitetura Transformer têm atraído ampla atenção devido ao seu desempenho excepcional. LLMs industriais precisam processar sequências longas para fornecer serviços de alta qualidade. No entanto, o consumo de memória cresce de forma quadrática com o comprimento da sequência, limitando a capacidade de escalabilidade do treinamento de sequências longas. Métodos paralelos existentes geram tensores redundantes durante a execução, deixando espaço para otimização de memória; ao mesmo tempo, o paralelismo de tensor (TP) não consegue realizar uma sobreposição eficaz entre cálculo e comunicação. Para abordar esses problemas, este artigo propõe um método paralelo universal — paralelismo de tensor com alta eficiência de memória (METP), projetado especialmente para as unidades centrais de cálculo no treinamento do Transformer (ou seja, duas multiplicações matriciais consecutivas e possíveis operações funcionais entre elas O=f(AB)C). METP distribui as subtarefas do cálculo de O em vários dispositivos, usando comunicação ponto a ponto (send/recv) em vez de comunicação coletiva para trocar submatrizes e concluir o cálculo, evitando a geração de tensores redundantes. Por meio da técnica de double buffering, alcança-se uma sobreposição profunda entre cálculo e comunicação, e são propostas condições teóricas de sobreposição completa para orientar o treinamento do Transformer em sequências longas. Análises teóricas mostram que, com paralelismo de grau p, o custo de memória de METP sem uso de FlashAttention para cálculo da atenção é O(1/p3); ao usar FlashAttention para cálculo de atenção multi-cabeça, METP pode economizar pelo menos 41,7% de memória em comparação com TP. Experimentos demonstram que, com uma configuração de 8 GPUs A100, METP pode aumentar o comprimento da sequência de 2,38 a 2,99 vezes em relação a outros métodos.
Keywords
aprendizado distribuído;grandes modelos de linguagem;sequências longas;sistemas de aprendizado de máquina;eficiência de memória;paralelismo de tensor