Paralelismo de tensores é uma técnica para distribuir o cálculo de uma única camada de rede neural entre vários dispositivos de hardware. Nesta abordagem, as matrizes de pesos de uma camada são divididas em fragmentos, e cada dispositivo detém uma porção dos parâmetros e executa a porção correspondente da multiplicação de matrizes. Os resultados parciais são então combinados por meio de operações de comunicação coletiva, como all-reduce, para produzir a saída da camada. Este método é uma forma de paralelismo de modelo, distinto do paralelismo de dados, onde cada dispositivo detém uma cópia completa do modelo e processa diferentes amostras de dados. O paralelismo de tensores é essencial para treinar e executar inferência em modelos muito grandes, como modelos de linguagem de grande porte, que excedem a capacidade de memória de um único acelerador.
A principal motivação para o paralelismo de tensores são as demandas de memória e computação dos modelos modernos de aprendizado profundo. Por exemplo, um modelo com centenas de bilhões de parâmetros não cabe na memória de uma única GPU ou TPU, que normalmente varia de 16 a 80 gigabytes. Ao fragmentar as matrizes de pesos, o paralelismo de tensores permite que o modelo seja distribuído por muitos dispositivos, possibilitando o uso de modelos maiores e tempos de treinamento mais rápidos. Também reduz a pegada de memória por dispositivo e permite tamanhos de lote maiores. No entanto, o paralelismo de tensores introduz sobrecarga de comunicação, pois os dispositivos devem sincronizar resultados parciais, o que pode se tornar um gargalo, especialmente em interconexões mais lentas.
O paralelismo de tensores é comumente usado em conjunto com outras estratégias de paralelismo, como paralelismo de pipeline e paralelismo de dados, para escalar o treinamento em milhares de dispositivos. É um componente-chave da infraestrutura de treinamento de muitos modelos de última geração, incluindo os desenvolvidos por OpenAI, Anthropic e Google DeepMind. A técnica também é relevante para inferência, pois permite servir modelos grandes com baixa latência, distribuindo o cálculo por vários aceleradores.
Histórico
O conceito de dividir operações de matriz entre vários processadores remonta às primeiras pesquisas em computação paralela. Nas décadas de 1980 e 1990, pesquisadores exploraram implementações paralelas de redes neurais, mas as ferramentas de hardware e software eram limitadas. A forma moderna de paralelismo de tensores surgiu com o avanço do aprendizado profundo baseado em GPU na década de 2010. Um dos primeiros usos notáveis foi no treinamento de grandes redes convolucionais, onde pesquisadores da Universidade de Toronto e outras instituições experimentaram paralelismo de modelo. No entanto, foi o advento dos modelos baseados em transformadores, como a arquitetura Transformer introduzida em 2017, que tornou o paralelismo de tensores uma técnica mainstream. As camadas de atenção e as redes feed-forward do Transformer são compostas por grandes multiplicações de matrizes, que são candidatas naturais para fragmentação.
Em 2019, pesquisadores da NVIDIA e outras organizações publicaram descrições detalhadas de paralelismo de tensores para treinar modelos de linguagem de grande porte. O framework Megatron-LM, desenvolvido pela NVIDIA, demonstrou paralelismo de tensores eficiente para modelos transformadores, alcançando escalonamento quase linear em múltiplas GPUs. Este trabalho lançou as bases para sistemas subsequentes de treinamento em larga escala, como o GShard do Google e o DeepSpeed da Microsoft, que incorporam paralelismo de tensores como um recurso central.
Como Funciona o Paralelismo de Tensores
Em uma camada típica de rede neural, a passagem direta calcula uma multiplicação de matrizes: Y = XW, onde X são as ativações de entrada, W é a matriz de pesos e Y é a saída. No paralelismo de tensores, a matriz de pesos W é particionada ao longo de uma ou mais dimensões. Para uma camada totalmente conectada, uma abordagem comum é dividir W por colunas (ou seja, ao longo da dimensão de saída). Cada dispositivo detém um subconjunto das colunas de W e calcula uma saída parcial. As saídas parciais são então combinadas usando uma operação all-reduce para produzir o Y final. Isso é conhecido como particionamento paralelo por colunas. Alternativamente, W pode ser dividido por linhas (ao longo da dimensão de entrada), exigindo uma operação all-gather na entrada e um all-reduce na saída.
Para modelos transformadores, o paralelismo de tensores é aplicado tanto à atenção multi-cabeça quanto às redes feed-forward. Na atenção multi-cabeça, as matrizes de pesos de consulta, chave e valor são fragmentadas entre os dispositivos, e as cabeças de atenção são distribuídas. A projeção de saída é então combinada. Na rede feed-forward, as duas camadas lineares são tipicamente particionadas: a primeira camada é dividida por colunas, e a segunda camada é dividida por linhas, de modo que as ativações intermediárias sejam fragmentadas e a saída seja reduzida. Este design minimiza a comunicação, exigindo apenas operações all-reduce em pontos específicos.
O padrão de comunicação é crítico. Cada operação all-reduce requer a troca de dados entre todos os dispositivos no grupo de paralelismo de tensores. O volume de comunicação é proporcional ao tamanho das ativações, não dos pesos, que geralmente é menor. No entanto, para modelos muito grandes, a comunicação ainda pode ser significativa. Para mitigar isso, o paralelismo de tensores é tipicamente usado dentro de um único nó com interconexões de alta velocidade, como NVLink da NVIDIA ou Infinity Fabric da AMD, enquanto o paralelismo de dados é usado entre nós.
Comparação com Outras Técnicas de Paralelismo
O paralelismo de tensores é uma das várias estratégias de paralelismo de modelo. O paralelismo de pipeline, por exemplo, divide o modelo por camadas, com cada dispositivo responsável por um conjunto contíguo de camadas. Isso reduz a comunicação em comparação com o paralelismo de tensores, pois apenas as ativações nos limites das camadas são trocadas. No entanto, o paralelismo de pipeline pode sofrer com tempo ocioso devido a bolhas de pipeline. O paralelismo de dados, por outro lado, replica o modelo em cada dispositivo e processa diferentes amostras de dados, exigindo sincronização de gradientes após cada etapa. O paralelismo de tensores é complementar a ambos: pode ser combinado com paralelismo de pipeline para reduzir a memória por dispositivo e com paralelismo de dados para aumentar a taxa de transferência.
Outra técnica relacionada é o paralelismo de sequência, que fragmenta a dimensão de sequência em modelos transformadores. Isso é frequentemente usado em conjunto com paralelismo de tensores para reduzir ainda mais o uso de memória. O paralelismo de especialistas, usado em modelos de mistura de especialistas, fragmenta os módulos especialistas entre dispositivos, o que é uma forma de paralelismo de modelo, mas não estritamente paralelismo de tensores.
Aplicações e Casos de Uso
O paralelismo de tensores é amplamente usado no treinamento e na inferência de modelos de linguagem de grande porte. Por exemplo, o modelo GPT-3, desenvolvido pela OpenAI, foi treinado usando uma combinação de paralelismo de modelo, incluindo paralelismo de tensores, em milhares de GPUs da NVIDIA. Da mesma forma, os modelos Claude da Anthropic e os modelos Gemini do Google DeepMind dependem de paralelismo de tensores para escalar para centenas de bilhões de parâmetros. Na inferência, o paralelismo de tensores permite servir modelos como GPT-4 com baixa latência, distribuindo o cálculo por várias GPUs, como feito pelo Triton Inference Server da NVIDIA e pelo Text Generation Inference da Hugging Face.
Além de modelos de linguagem, o paralelismo de tensores também é usado em outros domínios, como visão computacional e computação científica, onde grandes modelos ou grandes tamanhos de lote exigem computação distribuída. Por exemplo, treinar grandes transformadores de visão (ViTs) em imagens de alta resolução pode se beneficiar do paralelismo de tensores.
Desafios e Limitações
Apesar de seus benefícios, o paralelismo de tensores tem vários desafios. O principal problema é a sobrecarga de comunicação. Operações all-reduce exigem interconexões de alta largura de banda e baixa latência. Em sistemas com redes mais lentas, como Ethernet, o paralelismo de tensores pode se tornar ineficiente. Portanto, é tipicamente usado dentro de um único servidor ou em um cluster firmemente acoplado. Outro desafio é o balanceamento de carga: fragmentação desigual pode levar a alguns dispositivos subutilizados. Além disso, o paralelismo de tensores requer gerenciamento cuidadoso de memória, pois cada dispositivo deve armazenar seu fragmento dos pesos e as ativações intermediárias. Finalmente, implementar paralelismo de tensores manualmente é complexo; requer modificar o código do modelo e lidar com primitivas de comunicação. Felizmente, frameworks como PyTorch e TensorFlow fornecem suporte integrado para paralelismo de tensores, abstraindo grande parte da complexidade.
Suporte de Software e Frameworks
Vários frameworks e bibliotecas de aprendizado profundo suportam paralelismo de tensores. O PyTorch oferece o módulo torch.distributed e a API tensor_parallel para fragmentar tensores. O TensorFlow fornece tf.distribute com suporte para paralelismo de modelo. Bibliotecas especializadas como Megatron-LM, DeepSpeed e os Transformers da Hugging Face implementaram paralelismo de tensores para modelos transformadores. Essas ferramentas permitem que pesquisadores e engenheiros apliquem paralelismo de tensores com mudanças mínimas de código, muitas vezes apenas especificando o número de dispositivos e a estratégia de fragmentação.
Direções Futuras
À medida que os modelos continuam a crescer, o paralelismo de tensores permanecerá uma técnica crítica. Desenvolvimentos futuros podem incluir algoritmos de comunicação mais eficientes, como all-reduce hierárquico, e melhor integração com hardware heterogêneo, incluindo GPUs da AMD, aceleradores da Intel e chips personalizados como AWS Trainium e LPUs da Groq. Além disso, a descoberta automatizada de paralelismo, onde o sistema determina a estratégia de fragmentação ideal, é uma área ativa de pesquisa. Com a crescente demanda por IA em larga escala, o paralelismo de tensores continuará a evoluir para atender aos desafios de escala e eficiência.