Resumo Técnico: MatrixFSDP

Declaração do Problema

O treinamento de modelos de linguagem de larga escala depende cada vez mais de otimizadores estruturados em matrizes como Muon, Shampoo e SOAP. Diferente dos otimizadores por coordenada (ex: AdamW), esses métodos exigem a matriz de gradiente 2D completa como entrada para realizar operações como a ortogonalização de Newton–Schulz ou o pré-condicionamento de matrizes. Isso cria um descompasso sistêmico fundamental com o FSDP (Fully Sharded Data Parallel) e o ZeRO-3, que fragmentam parâmetros, gradientes e estados de otimização entre os ranks para maximizar a eficiência de memória.

Sob o ZeRO-3 padrão, cada rank detém apenas uma fatia do modelo. Para aplicar um otimizador de matriz, o sistema deve reconstruir a matriz de gradiente 2D completa a cada passo do otimizador. Essa comunicação é um gargalo crítico no caminho crítico que não pode ser sobreposto ao cálculo das camadas, resultando em uma latência significativa que escala mal com o número de nós e o tamanho do modelo.

As soluções existentes são incompletas:

1. Reconstrução (Stock FSDP2-Muon): Mantém a economia de memória do ZeRO-3, mas incorre em um pesado overhead de comunicação a cada passo do otimizador.
2. Posicionamento de Dono ZeRO-1 (ZeRO-1 Owner Placement): Atribui matrizes inteiras a ranks específicos (donos) para tornar o passo do otimizador local, mas exige que os parâmetros completos residam em cada rank, abandonando a eficiência de memória do ZeRO-3.

O objetivo é alcançar entradas de otimizador de matriz completa sem sacrificar a memória de parâmetros da escala ZeRO-3.

Metodologia: MatrixFSDP

O MatrixFSDP introduz um terceiro caminho: fragmentos ZeRO-3 de formato de dono (Owner-shaped ZeRO-3 shards). Em vez de alterar o algoritmo do otimizador ou abandonar o fragmentamento, ele redefine onde os fragmentos residem dentro do grupo de paralelismo de dados.

1. Posicionamento de Formato de Dono (Owner-Shaped Placement)

Para cada matriz de peso 2D no modelo, o MatrixFSDP designa um único rank de paralelismo de dados como o dono (owner).

• Rank Dono: Detém a matriz 2D completa.
• Ranks Não-Donos: Detêm um fragmento vazio para aquele parâmetro específico.
• Tensores Não-Matriciais: Vetores 1D e outros tensores não-2D são agrupados em uma "cauda" (tail) e atribuídos a um dono separado, onde continuam a usar o otimizador AdamW padrão.

Isso garante que, após o backward pass e a redução de gradiente, a entrada completa do Muon (a matriz de gradiente 2D completa) resida naturalmente no rank dono. Consequentemente, a atualização de Newton–Schulz é executada localmente no dono com zero comunicação de matriz no passo do otimizador.

2. Desafios de Tempo de Execução e Soluções

Este layout desigual (alguns ranks detendo matrizes completas, outros detendo fragmentos vazios) viola as suposições de tamanho igual de coletivos padrão do FSDP2. O MatrixFSDP aborda isso através de quatro mecanismos principais de tempo de execução:

• Metadados de MatrixShard & Planejador Global: Um planejador global atribui donos para equilibrar os bytes residentes e a carga de trabalho do otimizador através do cluster. Esses metadados rastreiam qual rank possui qual papel de tensor, garantindo que o sistema permaneça um layout ZeRO-3 coerente em vez de uma coleção de casos especiais.
• Coletivos P2P de Segmento de Dono (Owner-Segment P2P Collectives): Os all-gathers padrão são substituídos por coletivos de segmento de dono determinísticos. Durante a materialização de forward/backward, apenas segmentos não vazios são transferidos do dono para os ranks que os necessitam. Durante a redução de gradiente, os resultados são escritos apenas nos segmentos de dono não vazios. Isso evita o preenchimento (padding) de ranks vazios e elimina os all-gathers de matriz completa.
• Fixação de Buffer de Dono (Owner-Buffer Pinning): Para lidar com o ciclo de vida do FSDP2 (desfragmentar $\to$ computar $\to$ fragmentar novamente), o MatrixFSDP introduz um estado "fixado" (pinned). Quando um buffer é reduzido após o forward pass, o armazenamento não é devolvido ao pool compartilhado, mas fixado. Isso permite que o backward pass "reviva" a matriz completa no mesmo local de memória, prevenindo erros de visualização (view errors) do autograd e garantindo que o dono possa realizar a atualização antes de liberar o buffer.
• Checkpointing e Roteamento: O checkpointing (DCP) e o roteamento de estado do otimizador utilizam os metadados do MatrixShard. Os estados do otimizador são roteados para os donos corretos (Muon para matrizes, AdamW para caudas), e os checkpoints salvam o layout de posicionamento específico para suportar o re-fragmentamento entre diferentes tamanhos de mundo (world sizes).

Principais Contribuições

1. ZeRO-3 de Formato de Dono (Owner-Shaped ZeRO-3): Uma estratégia de posicionamento inovadora onde matrizes 2D são inteiras em um único dono e vazias em outros lugares, permitindo a otimização de matriz local dentro de um regime de memória fragmentada.
2. Passos de Otimizador Sem Comunicação: Ao alinhar o destino da redução de gradiente com o requisito de entrada do otimizador, o MatrixFSDP elimina o gargalo de all-gather de tamanho de peso anteriormente inerente aos otimizadores de matriz no ZeRO-3.
3. Infraestrutura de Tempo de Execução: Um conjunto de componentes de tempo de execução (planejador, coletivos customizados, fixação de buffer) que tornam os layouts desiguais de dono eficientes e corretos dentro do framework FSDP2.
4. Generalidade: A abordagem é agnóstica ao otimizador, demonstrando funcionar com Muon, Shampoo e SOAP sem exigir variantes distribuídas dos pré-condicionadores.

Resultados Experimentais

Avaliado em 64 GPUs NVIDIA A100 (8 nós) com vários tamanhos de modelo (até 32B de parâmetros):

• Latência do Passo do Otimizador: O MatrixFSDP reduz a latência do passo do otimizador em 4,2× em um único nó e 54,6× em oito nós em comparação ao FSDP2-Muon padrão. O ganho de velocidade cresce com a escala porque a comunicação da linha de base cruza o fabric inter-nó, enquanto o passo do MatrixFSDP permanece local.
• Ganho de Velocidade End-to-End: Alcança até 2,15× de aceleração de treinamento end-to-end.
• Eficiência de Memória: O MatrixFSDP mantém os níveis de memória do ZeRO-3. Para um modelo de 32B, ele utiliza 10 GB/rank, enquanto uma abordagem de posicionamento de dono ZeRO-1 excederia o limite de 80 GB da GPU (projetado >54 GB/rank em 14B). É 3–7× mais eficiente em termos de memória do que as abordagens FSDP2-Muon de "reunir uma vez" (gather-once).
• Corretude: A trajetória de treinamento corresponde a uma referência de Muon DDP de matriz completa com alta precisão (diferença de perda impressa $\leq 0$), confirmando que o dono recebe exatamente a matriz de gradiente necessária para a atualização.

Significância e Alegações

O artigo afirma que o MatrixFSDP resolve o "descompasso sistêmico" entre otimizadores de matriz e o fragmentamento de parâmetros. Sua principal significância reside em permitir o treinamento em larga escala com otimizadores de matriz (que oferecem melhor eficiência de tokens e convergência) sem incorrer nos custos de comunicação proibitivos ou nos overheads de memória das abordagens anteriores.

Os autores enfatem que o MatrixFSDP é particularmente eficaz para execuções ZeRO-3 multi-nó com grandes spans de fragmentação, onde o custo de reconstruir matrizes a cada passo dominaria o tempo de treinamento. O artigo observa modestamente que os ganhos são menores em execuções de nó único ou sob escalonamento forte de modelo fixo, onde o número de papéis de dono se torna esparso, potencialmente tornando o fanout do dono um novo gargalo. O trabalho não pretende substituir os otimizadores por coordenada para parâmetros não matriciais, mas sim fornecer um caminho unificado e eficiente para lidar com os componentes de matriz 2D dos modelos modernos.