Learned Subspace Compression for Communication-Efficient Pipeline Parallelism

Resumo Técnico: Compressão de Subespaço Aprendida para Paralelismo de Pipeline com Eficiência de Comunicação

Definição do Problema
O treinamento de grandes modelos de linguagem (LLMs) através de hardware heterogêneo e geograficamente distribuído enfrenta gargalos de comunicação significativos, particularmente no paralelismo de pipeline. Embora o paralelismo de pipeline particione os parâmetros do modelo entre dispositivos para superar limites de memória, ele exige a troca de ativações intermediárias entre estágios durante as passagens direta (forward) e inversa (backward). Em redes de baixa largura de banda, essa comunicação entre estágios torna-se o gargalo dominante.

Abordagens existentes, como as Redes de Subespaço (SSNs), tentam mitigar isso projetando ativações em um subespaço de baixo posto compartilhado e fixo definido por uma matriz ortogonal. No entanto, este método requer modificações arquiteturais intrusivas, incluindo otimização restrita para manter os pesos dentro de um subespaço compartilhado e o uso de otimizadores modificados. Empiricamente, esses subespaços globais fixos levam a uma degradação significativa de desempenho, especialmente conforme as taxas de compressão aumentam, porque falham em se adaptar à geometria evolutiva das ativações durante o treinamento. Além disso, o aprendizado ingênuo de matrizes de projeção sem restrições faz com que elas escapem da variedade de Stiefel (o conjunto de matrizes ortogonais), destruindo as propriedades isométricas necessárias para uma compressão eficaz e levando a quedas severas de desempenho.

Metodologia: Aprendizado de Projeção Consciente de Variedade (MAPL)
Os autores propõem o Manifold Aware Projection Learning (MAPL), um método que trata a compressão entre estágios como uma projeção geométrica aprendível em vez de uma restrição arquitetural fixa. Os componentes principais do MAPL são:

1. Projetores Ortogonais Aprendíveis por Estágio: Diferente das SSNs que impõem um único subespaço global, o MAPL equipa cada fronteira de estágio de pipeline com seu próprio projetor ortogonal aprendível $A_p \in \mathbb{R}^{d \times r}$. Esses projetores são restritos à variedade de Stiefel ($St(d, r)$), garantindo que $A_p^\top A_p = I_r$. Isso permite que cada estágio descubra e adapte continuamente seu próprio subespaço de compressão ideal para a tarefa.
2. Otimização Restrita à Variedade (SPEL): Para evitar que os projetores escapem da variedade durante as atualizações de gradiente padrão, os autores utilizam o Descida de Gradiente Espectral na Variedade de Stiefel (SPEL). Este algoritmo atualiza os projetores usando projeções tangentes e operações de retração (via o método Polar Express) em cada etapa de otimização, garantindo que a restrição de ortogonalidade seja estritamente mantida.
3. Embeddings de Âncora Fatorizados: Para lidar com os offsets de alto posto e orientados a tokens nos fluxos residuais sem esgotar a capacidade do projetor, o MAPL introduz embeddings de âncora fatorizados. O offset é decomposto em uma matriz de baixo posto treinável $E_{small}^p$ e uma matriz ortogonal congelada $P_p$. Apenas IDs de tokens inteiros são transmitidos para reconstruir esses offsets no estágio receptor, adicionando um overhead de comunicação desprezível.
4. Quantização Vetorial (VQ): Para compressão adicional, o método suporta quantização vetorial residual (RVQ) nas representações projetadas de baixo posto. Um protocolo de sincronização de dicionário em streaming é usado para amortizar a comunicação do dicionário, permitindo altas taxas de compressão (até 16×) com custo de largura de banda mínimo.

Principais Contribuições

• Aprendizado de Subespaço por Estágio: O artigo demonstra que impor um subespaço global compartilhado degrada o aprendizado em configurações de correspondência de tokens. Em vez disso, o MAPL permite que cada estágio aprenda um subespaço de compressão distinto e restrito pela variedade de Stiefel, o que melhor se alinha com a estrutura intrínseca de baixo posto das ativações em diferentes profundidades.
• Otimização Restrita à Variedade: Os autores identificam o "escape da variedade" como o principal modo de falha do aprendizado de projetores não restritos. Eles introduzem o SPEL para manter a ortogonalidade, mostrando que o aprendizado não restrito é estritamente pior do que o uso de baselines fixos.
• Reconstrução Eficiente: Ao combinar projetores aprendíveis com embeddings de âncora fatorizados, o MAPL alcança a reconstrução de ativações de posto total com overhead de comunicação desprezível, eliminando a necessidade das decomposições arquiteturais intrusivas exigidas pelas SSNs.
• Melhoria da Fronteira de Pareto: O método traça a fronteira de Pareto de entropia cruzada de validação versus custo de comunicação, alcançando alta compressão com degradação de desempenho negligenciável em comparação com métodos anteriores.

Resultados Experimentais
Os autores avaliaram o MAPL em modelos LLaMA apenas decodificadores (decoder-only) variando de 150M a 1B de parâmetros, treinados no corpus DCLM com configurações de pipeline de $P=4$ e $P=8$ estágios.

• Desempenho vs. Compressão: O MAPL supera consistentemente os baselines SSN. Para um modelo de 150M com 4× compressão, o MAPL alcançou uma perda de validação dentro de 0,84% do baseline não comprimido, enquanto os baselines SSN degradaram de 8–12%.
• Escalabilidade: Na escala de 1B de parâmetros com 8× compressão, o MAPL manteve uma degradação de perda de apenas 1,38–2,02%, enquanto os baselines SSN sofreram degradações de até 26%.
• Tarefas de Downstream: Em avaliações zero-shot (HellaSwag, PIQA, ARC), o MAPL acompanhou de perto os baselines não comprimidos (dentro de ~1 ponto de precisão média nas escalas de 500M/1B), enquanto os baselines SSN sofreram quedas de precisão de até 9 pontos.
• Compressão Extrema: Quando combinado com quantização vetorial (MAPL+VQ), o método alcançou 16× de compressão com uma degradação adicional modesta (ex: aumento de ~2,3% na perda em 1B de parâmetros), embora a precisão das tarefas de downstream tenha apresentado uma queda mais pronunciada em comparação ao MAPL sem quantização.

Significância e Alegações
O artigo alega que as ativações de transformadores entre estágios admitem naturalmente estruturas de baixa dimensão adaptáveis que podem ser aprendidas conjuntamente com o objetivo da tarefa, em vez de serem impostas via subespaços globais fixos. Ao tratar a compressão como uma projeção geométrica aprendível na variedade de Stiefel, o MAPL permite uma redução agressiva de comunicação (4× a 16×) preservando a qualidade do treinamento e o desempenho em tarefas de downstream. O trabalho destaca o aprendizado de representação restrito à variedade como um mecanismo prático para reduzir custos de comunicação, ofereando um compromisso superior entre desempenho e compressão em comparação com abordagens de subespaço existentes.

Limitações
Os autores observam que os experimentos são atualmente limitados a modelos de até 1B de parâmetros, e que avaliações adicionais em escalas maiores e sob condições reais de redes heterogêneas são necessárias. Além disso, embora o MAPL reduza significamente a degradação, o desempenho ainda declina sob taxas de compressão extremas, particularmente quando a quantização vetorial é aplicada.