Folding Tensor and Sequence Parallelism for Memory-Efficient Transformer Training & Inference

1. Declaração do Problema

O treinamento e a inferência de grandes modelos de linguagem (LLMs) enfrentam dois principais gargalos de memória:

1. Memória de Parâmetros: O armazenamento necessário para pesos do modelo, gradientes e estados do otimizador.
2. Memória de Ativação: O armazenamento necessário para ativações intermediárias, que escala quadraticamente com o comprimento da sequência ($S^2$) em mecanismos de atenção.

Estratégias de paralelismo existentes abordam esses problemas de forma isolada ou com compensações:

• Paralelismo de Tensor (TP): Fragmenta os pesos do modelo entre dispositivos, reduzindo a memória de parâmetros, mas deixando a memória de ativação totalmente replicada. Torna-se ineficiente em comprimentos de sequência longos devido à alta memória de ativação e volume de comunicação.
• Paralelismo de Sequência (SP): Fragmenta a sequência de entrada entre dispositivos, reduzindo a memória de ativação, mas deixando os pesos do modelo totalmente replicados. É ineficaz quando a memória de parâmetros é o gargalo.
• TP+SP Convencional: Combina ambos atribuindo-os a eixos ortogonais em uma malha de dispositivos (por exemplo, uma grade $T \times \Sigma$). Embora eficaz, isso consome um grande número de dispositivos para paralelismo de modelo ($T \times \Sigma$), deixando menos dispositivos para paralelismo de dados. Além disso, frequentemente força grupos paralelos a atravessar links inter-nós mais lentos, subutilizando interconexões intra-nós de alta largura de banda (como NVLink ou AMD Infinity Fabric).

O Desafio Central: Como reduzir simultaneamente a memória de parâmetros e de ativação em um único eixo de dispositivo para maximizar a eficiência de memória e liberar dispositivos para paralelismo de dados, sem incorrer em custos de comunicação proibitivos.

2. Metodologia: Paralelismo de Tensor e Sequência (TSP)

Os autores propõem o TSP, uma estratégia que "dobra" o paralelismo de tensor e de sequência em um único eixo de dispositivo. Em vez de atribuir TP e SP a dimensões ortogonais, o TSP atribui a cada rank (GPU) tanto um fragmento dos pesos do modelo quanto um fragmento da sequência.

Componentes Principais de Design:

• Fragmentação Unificada: Cada dispositivo mantém $1/D$ dos parâmetros do modelo e $1/D$ dos tokens da sequência. Isso reduz tanto a memória proporcional aos parâmetros quanto a memória de ativação por um fator de $D$ em um único eixo.
• Programações de Runtime: Para lidar com o fato de que pesos e ativações são fragmentados de forma diferente, o TSP emprega programações de comunicação específicas para diferentes camadas:
  • Camadas de Atenção: Utiliza uma programação de broadcast.
    • Os ranks iteram sobre fragmentos de pesos (empacotados $W_Q, W_K, W_V, W_O$).
    • Em cada iteração, um rank transmite seu fragmento de peso para todos os outros.
    • Todos os ranks calculam projeções locais em seus fragmentos de sequência.
    • Uma operação All-Gather reconstrói o contexto completo de Chave/Valor ao longo da dimensão da sequência para realizar o FlashAttention.
    • Um esquema de particionamento em zigue-zague é utilizado para equilibrar a carga de trabalho da atenção causal (garantindo que tokens posteriores não incorram em trabalho desproporcionalmente maior).
  • Camadas MLP com Portão: Utiliza uma programação de comunicação em anel.
    • Os fragmentos de peso ($W_1, W_2, W_3$) circulam ponto a ponto ao redor do anel.
    • Cada rank acumula saídas parciais localmente.
    • Isso elimina a necessidade de uma operação All-Reduce (padrão no TP), pois a sequência permanece local e apenas os pesos se movem.
• Consciência de Hardware: O design é otimizado para topologias como a MI300X da AMD (Infinity Fabric), onde a largura de banda coletiva escala com o número de GPUs participantes. O TSP garante que os primitivos de comunicação (broadcast e anel) envolvam o grupo completo ou sejam sobrepostos ao cálculo para maximizar a eficiência.

3. Contribuições Principais

1. Estratégia de Paralelismo Inovadora: Introdução do TSP, que colapsa dois eixos de paralelismo ortogonais em um, permitindo a redução simultânea da memória de parâmetros e de ativação.
2. Análise Teórica: Um modelo matemático rigoroso comparando o TSP com DP, TP, SP e TP+SP em relação a:
   • Memória: O TSP é o único esquema que reduz tanto a memória proporcional aos pesos quanto a memória de ativação por $1/D$ sem exigir $D^2$ dispositivos.
   • Comunicação: Embora o TSP introduza comunicação de movimento de pesos (que o TP+SP evita), os autores provam que, para sequências longas ou tamanhos de lote moderados, a comunicação proporcional à ativação domina, tornando o TSP competitivo.
   • FLOPs: O TSP mantém a mesma eficiência computacional por dispositivo que o TP e o SP.
3. Implementação e Sobreposição: Implementação detalhada dos blocos de atenção e MLP do TSP com estratégias explícitas de sobreposição comunicação/cálculo (pipelining de transferências de pesos atrás de cálculos GEMM).
4. Otimização de Topologia: Demonstração de que o TSP permite que toda a réplica do modelo caiba dentro de um único nó de alta largura de banda (intra-nó), evitando as penalidades de desempenho de atravessar links inter-nós mais lentos frequentemente exigidos por malhas 2D TP+SP.

4. Resultados

Os autores avaliaram o TSP em um cluster de até 1.024 GPUs AMD MI300X (128 nós) com um modelo transformer denso de 7B de parâmetros.

• Eficiência de Memória:
  • O TSP consistentemente alcançou a memória de pico mais baixa em todos os comprimentos de sequência testados (16K a 128K tokens).
  • Em contextos curtos, o TSP igualou o TP (dominado pela memória de parâmetros).
  • Em contextos longos, o TSP superou significativamente o TP e o TP+SP (dominados pela memória de ativação), prevenindo erros de Memória Esgotada (OOM) onde outras estratégias falharam.
• Vazão (Throughput):
  • Apesar de um volume de comunicação teoricamente maior (devido ao movimento de pesos), o TSP manteve vazão competitiva com o TP+SP.
  • A comunicação extra foi ocultada com sucesso atrás do cálculo via pipelining.
  • As vantagens de vazão cresceram com o aumento dos graus de paralelismo ($D$) e dos comprimentos de sequência.
• Escalabilidade:
  • O TSP permitiu micro-lotes maiores devido à redução da pegada de memória, melhorando ainda mais a vazão.
  • Permitiu encaixar réplicas de modelo inteiramente dentro de um único nó (8 GPUs), liberando ranks para Paralelismo de Dados (DP) em clusters maiores.

5. Significado

• Escalabilidade Consciente de Hardware: O TSP fornece uma solução prática para clusters modernos de HPC onde a largura de banda intra-nó é alta, mas a largura de banda inter-nó é um gargalo. Maximiza a utilidade de interconexões intra-nós caras.
• Viabilidade de Longo Contexto: À medida que os modelos avançam para janelas de contexto mais longas (100k+ tokens), a memória de ativação torna-se o gargalo primário. O TSP está posicionado de forma única para lidar com isso sem a sobrecarga massiva de dispositivos do paralelismo 2D.
• Flexibilidade: O TSP pode ser composto ortogonalmente com outros esquemas de paralelismo (Paralelismo de Pipeline, Paralelismo de Especialistas, Paralelismo de Dados), oferecendo um "eixo" flexível para projetistas de sistema otimizarem para topologias de hardware específicas.
• Eficiência de Custos: Ao reduzir o número de dispositivos necessários para paralelismo de modelo, o TSP permite que mais dispositivos sejam usados para paralelismo de dados, potencialmente reduzindo o tamanho total do cluster necessário para uma determinada tarefa de treinamento ou aumentando o tamanho de lote efetivo.

Em conclusão, o TSP representa uma evolução significativa nas estratégias de treinamento distribuído, trocando um aumento gerenciável no volume de comunicação por ganhos substanciais na eficiência de memória e utilização de hardware, tornando-o uma ferramenta crítica para o treinamento de LLMs de próxima geração com longo contexto.