Improved Scaling for Fast Mode of Ozaki Scheme II

Resumo Técnico: Melhoria de Escalonamento para o Modo Rápido do Esquema de Ozaki II

Problema
O Esquema de Ozaki II é um método para emular multiplicação de matrizes de alta precisão (ex: FP64/FP32) utilizando hardware de baixa precisão de inteiros (ex: INT8 Tensor Cores) baseado no Teorema Chinês dos Restos (CRT). O esquema requer o escalonamento das matrizes de entrada de alta precisão em matrizes de inteiros de tal forma que o produto desses inteiros, quando reconstruído via CRT, não exceda um limiar específico ($P/2$). Se este limiar for excedido, a recuperação via CRT falha, levando a uma perda catastrófica de precisão.

O esquema oferece dois modos para determinar os fatores de escalonamento ($\mu$ e $\nu$):

1. Modo Preciso (Accurate Mode): Utiliza uma multiplicação de matrizes (GEMM) auxiliar de INT8 para estimar os fatores de escalonamento, garantindo alta precisão, mas incorrendo em um overhead computacional significativo.
2. Modo Rápido (Fast Mode): Utiliza a desigualdade de Cauchy–Schwarz para estimar os fatores de escalonamento sem a GEMM auxiliar, oferecendo maior vazão (throughput), mas historicamente menor precisão.

Este artigo identifica uma falha crítica na fórmula de escalonamento existente do Modo Rápido: ela carece de invariância de escala. Quando as matrizes de entrada são multiplicadas por uma constante escalar $\alpha$, a largura de bits efetiva das matrizes de inteiros escalonadas resultantes muda de forma imprevisível.

• Para $\alpha$ grande, a largura de bits efetiva diminui, causando degradação da precisão.
• Para $\alpha$ pequeno, o produto de inteiros escalonados pode exceder o limiar do CRT ($P/2$), causando falha total na recuperação.
  Experimentos numéricos em matrizes de todos os uns e matrizes aleatórias demonstraram que, para certas faixas escalares, o Modo Rápido falha em recuperar o resultado correto, enquanto o Modo Preciso permanece estável.

Metodologia
Os autores propõem uma fórmula de escalonamento revisada para o Modo Rápido que é derivada diretamente da condição de unicidade do CRT via desigualdade de Cauchy–Schwarz.

1. Derivação Teórica:
   Os autores reformulam a condição para os fatores de escalonamento $\mu_i$ e $\nu_j$. Em vez da heurística anterior, eles definem uma condição suficiente onde:
   $$2\mu_i \|a_{i,:}\|_2 \|b_{:,j}\|_2 \nu_j \leq P - 1$$
   Ao resolver para $\mu_i$ e $\nu_j$ usando logaritmos de base 2 e contabilizando os limites de erro absoluto da função log2f fornecida pelas APIs matemáticas de GPU, eles derivam novas fórmulas:
   $$\mu_i := 2^{\lfloor P'_{prop} - 0.5 \cdot (\widehat{\log_2}(\sum a_{ih}^2) + 4u_{32}) \rfloor}$$
   onde $P'_{prop}$ é uma constante pré-computada e $\widehat{\log_2}$ é uma aproximação limitada do logaritmo.

2. Prova de Invariância de Escala:
   Os autores demonstram analiticamente que a nova fórmula é invariante de escala. Se a matriz de entrada $A$ for escalonada por $\alpha = 2^s$, o termo $\log_2(\alpha)$ se cancela no cálculo da largura de bits efetiva da matriz de inteiros escalonada. Consequentemente, a largura de bits permanece constante independentemente da magnitude da entrada, prevenindo a degradação da precisão e a falha do CRT observadas no Modo Rápido original.

3. Implementação:
   O método proposto (OS II-prop) substitui a fórmula de escalonamento no Modo Rápido da biblioteca GEMMul8. Ele mantém a complexidade computacional $O(mk)$e$O(kn)$ do Modo Rápido original, introduzindo nenhum overhead adicional em comparação ao Modo Rápido original e permanecendo mais rápido que o Modo Preciso (que requer uma GEMM extra).

Principais Contribuições

• Identificação de Limitação: O artigo demonstra que a fórmula de escalonamento do Modo Rápido existente carece de invariância de escala, levando à degradação da precisão e falha de recuperação do CRT para escalonamentos de entrada específicos.
• Fórmula Revisada: Uma nova fórmula de escalonamento é proposta que garante a invariância de escala e assegura que a condição de unicidade do CRT seja sempre satisfeita.
• Zero Overhead: A nova fórmula não introduz custo computacional adicional sobre o Modo Rápido original.
• Avaliação Abrangente: O método é avaliado em uma GPU NVIDIA GH200 para multiplicações de matrizes de Precisão Dupla (DGEMM) e Precisão Simples (SGEMM).

Resultos
Experimentos foram conduzidos usando matrizes aleatórias com variando dispersões de magnitude de elementos ($\phi$) e dimensões de matriz ($k \in \{1024, 16384\}$).

• Precisão: O método proposto (OS II-prop) atinge precisão comparável ao Modo Preciso (OS II-accu) em todos os valores de $\phi$. Em contraste, o Modo Rápido original (OS II-fast) mostrou degradação significativa de precisão e falha para valores de $\phi$ elevados (casos mal condicionados).
• Vazão (Throughput): O OS II-prop mantém uma vazão comparável ao Modo Rápido original. Embora o OS II-prop seja ligeiramente mais lento que o Modo Rápido original em cenários específicos com poucos módulos (devido a menos elementos zero nos inteiros escalonados), ele supera significativamente o Modo Preciso em termos de vazão.
• Compromisso (Trade-off): Na análise de compromisso entre precisão e vazão, o OS II-prop domina o OS II-accu (maior vazão para a mesma precisão) e supera as limitações de precisão do OS II-fast. É o único método capaz de superar simultaneamente o cuBLAS em precisão e vazão para uma ampla gama de condições.

Significância
O artigo afirma que o método proposto resolve a limitação fundamental de precisão do Modo Rápido no Esquema de Ozaki II sem sacrificar sua vantagem de desempenho. Ao garantir a invariância de escala, o método garante que a recuperação via CRT tenha sucesso independentemente do escalonamento de entrada, oferecendo um equilíbrio superior entre precisão e desempenho. Isso permite que os usuários obtenham resultados de alta precisão em hardware de baixa precisão com a velocidade do Modo Rápido e a confiabilidade do Modo Preciso, superando efetivamente as restrições de vazão do Modo Preciso e as restrições de precisão do Modo Rápido original.

Os autores observam que, embora o presente trabalho foque em motores de matriz INT8 para DGEMM e SGEMM, a abordagem é aplicável a motores FP8 e multiplicações de matrizes complexas (ZGEMM/CGEMM), que ficam para trabalhos futuros.