Resumo Técnico: DASH – Shampoo Mais Rápido via Pré-condicionamento de Blocos em Lote e Solvers de Raiz Inversa Eficientes

Definição do Problema

O Shampoo é um otimizador de segunda ordem aproximado proeminente, conhecido por produzir modelos com menores outliers de ativação e melhor compressibilidade. No entanto, sua adoção generalizada é dificultada por uma sobrecarga computacional significativa, principalmente devido à operação dispendiosa de computar raízes de matrizes inversas ($A^{-1/4}$) para matrizes de pré-condicionamento. Embora o Shampoo Distribuído tente mitigar isso dividindo os pré-condicionadores em blocos de $B \times B$ para reduzir a complexidade de $O(n^3)$ para $O(Bn^2)$, as implementações existentes sofrem de lacunas críticas de eficiência:

1. Processamento Sequencial: Implementações padrão processam blocos sequencialmente, falhando em aproveitar totalmente o paralelismo das GPUs modernas.
2. Solvers Ineficientes: A dependência padrão da Decomposição de Autovalores (EVD) é difícil de paralelizar em GPUs. Métodos iterativos alternativos como o Coupled-Newton (CN) são frequentemente desativados por padrão devido a preocupações com estabilidade numérica e velocidade de convergência.
3. Escalonamento Subotimizado: Métodos atuais costumam escalar matrizes usando a norma de Frobenius, que é computacionalmente barata, mas pode empurrar os autovalores em direção a zero, exigindo significativamente mais iterações para convergir para uma precisão alvo.

Metodologia

O artigo propõe o DASH (Distributed Accelerated SHampoo), uma implementação de alto desempenho projetada para preencher a lacuna entre a eficiência teórica e as capacidades de hardware. A metodologia baseia-se em dois pilares principais:

1. Pré-condicionamento de Blocos em Lote (Otimização de Sistema)

Em vez de processar blocos de pré-condicionamento sequencialmente, o DASH empilha blocos independentes em tensores 3D.

• Empilhamento de Tensores 3D: Blocos de gradiente $G_{ij}$ são empilhados em uma matriz de forma $(N, B, B)$, onde $N$ é o número total de blocos.
• Execução Paralela: Isso permite que o otimizador aplique procedimentos de raiz inversa (EVD, CN ou novos métodos iterativos) a todo o lote de blocos simultaneamente usando primitivas de multiplicação de matrizes em lote (bmm).
• Utilização de GPU: Essa mudança arquitetural maximiza o throughput dos Tensor Cores da GPU, reduzindo a fragmentação de memória e permitindo o uso eficiente de aritmética de meia precisão (FP16) para certos solvers.

2. Solvers de Raiz Inversa Eficientes (Otimização Algorítmica)

O artigo introduz e analisa abordagens de álgebra linear avançadas para substituir ou complementar o EVD e o CN padrão:

• Iteração de Newton-Denman-Beavers (NDB): Uma abordagem inovadora para o Shampoo que computa tanto a raiz quadrada quanto a raiz quadrada inversa. Os autores otimizam a primeira iteração para eliminar multiplicações de matrizes redundantes, reduzindo o overhead para 3 matmuls por passo (efetivamente menos na prática).
• Aproximações de Polinômios de Chebyshev: Utilizando o algoritmo de Clenshaw para aproximar potências de matrizes, oferecendo um solver iterativo alternativo.
• Multi-Power-Iteration para Escalonamento: Para abordar problemas de convergência, o artigo substitui o escalonamento padrão pela norma de Frobenius por um método de Multi-Power-Iteration.
  • Mecanismo: Ele estima o raio espectral ($\lambda_{max}$) usando um conjunto de vetores iniciais (até 32) processados em paralelo via Power-Iteration.
  • Benefício: Isso fornece um fator de escala mais preciso do que a norma de Frobenius, garantindo que os autovalores permaneçam em um regime onde os solvers iterativos convergem rapidamente (próximo a 1), reduzendo assim o número de iterações necessárias e melhorando a estabilidade numérica.

Principais Contribuições

1. Implementação DASH: Um framework eficiente para GPU que substitui loops sequenciais por operações de tensor 3D paralelas, alcançando até 5,6× mais rapidez nos passos do otimizador em comparação ao Distributed Shampoo bem otimizado.
2. Novos Solvers: A introdução da iteração Newton-DB (NDB) e das aproximações de polinômios de Chebyshev como alternativas viáveis e mais rápidas ao EVD e CN. O artigo demonstra que o NDB produz uma perplexidade de validação menor do que o Coupled-Newton e frequentemente supera o EVD.
3. Análise de Convergência: Uma análise comportamental detalhada mostrando que o escalonamento padrão pela norma de Frobenius é subotimizado, levando frequentemente a uma convergência mais lenta. O artigo fornece explicações intuitivas para os distintos comportamentos de convergência do CN versus NDB.
4. Multi-Power-Iteration: Uma implementação eficiente de meia precisão de Power-Iteration que estima robustamente os raios espectrais, permitindo um escalonamento de matriz ideal que satisfaz os critérios de convergência para NDB e CN com menos passos.

Resultados Experimentais

Os autores avaliaram o DASH em um modelo Llama com 953M de parâmetros (tamanho de embedding 2048, comprimento de sequência 1024) treinado no dataset C4.

• Velocidade:
  • O DASH reduz o tempo de execução do passo do otimizador em 4× a 5,6× em comparação ao Distributed Shampoo, dependendo do método de raiz inversa e da precisão utilizada.
  • O aumento de velocidade mais significativo (5,6×) foi observado ao comparar o EVD (FP32) do Distributed Shampoo contra o CN (FP16) do DASH com empilhamento em nível de GPU.
  • Mesmo com atualizações frequentes de pré-condicionador ($f=1$), o DASH mantém vantagens substanciais de velocidade.
• Precisão (Perplexidade de Validação):
  • O DASH iguala o Distributed Shampoo em quase todas as configurações estáveis.
  • O NDB produz consistentemente uma perplexidade de validação menor do que o CN e frequentemente supera o EVD.
  • O uso de escalonamento por Power-Iteration com NDB resultou na menor perplexidade de validação (11,68) em comparação com o escalonamento de Frobenius (11,76), validando a análise teórica.
• Precisão:
  • A implementação suporta FP16 para CN e EVD sem degradar a perplexidade, reduzindo ainda mais o tempo de execução.
  • O NDB em FP16/BF16 foi considerado numericamente instável em seus experimentos, embora os autores notem isso como uma área para trabalhos futuros.

Significância e Alegações

O artigo afirma que o DASH reduz significativamente a barreira computacional para o uso de otimizadores de segunda ordem como o Shampoo, tornando o pré-condicionamento de alta qualidade prático para uso em larga escala.

• Unindo Teoria e Hardware: Ao alinhar o design algorítmico (solvers iterativos, escalonamento de matriz) com as capacidades de hardware (Tensor Cores, operações em lote), o DASH permite que pesquisadores investiguem propriedades de otimizadores sem tempos de execução proibitivos.
• Estabilidade Numérica vs. Velocidade: O trabalho desafia a suposição de que métodos iterativos são inerentemente menos estáveis ou mais lentos que o EVD. Ao introduzir um melhor escalonamento (Multi-Power-Iteration) e iterações otimizadas (NDB), os autores demonstram que métodos iterativos podem alcançar convergência superior e menor perplexidade enquanto são significativamente mais rápidos.
• Direções Futuras: Os autores sugerem que o desempenho superior do Newton-DB abre caminhos para a seleção dinâmica de solvers baseada em números de condição de blocos e para a estabilização do NDB em formatos de menor precisão usando métodos estocásticos ou de correção de erro.

O artigo conclui que o principal gargalo do Shampoo não era a complexidade matemática, mas sim a execução fragmentada e a instabilidade numérica, ambos os quais o DASH aborda de forma eficaz.