veScale-FSDP: Flexible and High-Performance FSDP at Scale

O artigo apresenta o veScale-FSDP, um sistema redesenhado que combina o formato de fragmentação flexível "RaggedShard" com um algoritmo de planejamento consciente da estrutura para superar as limitações de eficiência e compatibilidade dos sistemas FSDP atuais, permitindo o treinamento escalável de modelos de grande porte com quantização em blocos e otimizadores não elementares.

O essencial

Imagine que você está tentando construir o prédio mais alto do mundo (um modelo de Inteligência Artificial gigante) usando milhares de equipes de pedreiros (GPUs) espalhadas por todo o planeta. O desafio é que o prédio é tão grande que nenhuma equipe consegue carregar todos os tijolos sozinha. Elas precisam dividir o trabalho.

Até agora, a maneira de dividir esse trabalho (chamada de FSDP) era como se fosse um "corte de pizza" muito rígido. Se você tinha uma pizza com 10 fatias e 3 pessoas, a pizza era cortada em fatias iguais, mas se a fatia caísse em cima de uma borda de uma peça especial (como um queijo especial ou uma borda crocante), o corte estragava a peça. Para consertar isso, os engenheiros tinham que parar a obra, recortar tudo manualmente ou jogar fora pedaços inteiros, desperdiçando tempo e espaço.

O veScale-FSDP é a nova ferramenta que a ByteDance criou para resolver exatamente esse problema. Vamos entender como funciona com analogias simples:

1. O Problema: O Corte Rígido vs. A Realidade

Os modelos modernos de IA não são mais apenas "fatias de pizza". Eles têm estruturas complexas:

• Otimizadores Inteligentes: Alguns métodos de aprendizado (como o Muon ou Shampoo) precisam olhar para o "bloco inteiro" de uma vez, não apenas para um pedacinho solto.
• Quantização em Blocos: Para economizar memória, os modelos são comprimidos em blocos de 32x32. Se o sistema de divisão cortar um desses blocos ao meio, a compressão quebra e o modelo fica lento ou errado.

Os sistemas antigos (como o FSDP padrão do PyTorch ou o DeepSpeed) cortavam os dados de forma "elementar" (pedaço por pedaço) ou "por linha". Isso criava bordas bagunçadas onde os blocos especiais eram cortados, obrigando os computadores a fazerem cópias extras, preencherem espaços vazios (padding) e trocarem dados desnecessários. Era como tentar montar um quebra-cabeça onde metade das peças foi cortada ao meio por engano.

2. A Solução: O "RaggedShard" (O Cortador Flexível)

O grande trunfo do veScale-FSDP é uma nova forma de dividir os dados chamada RaggedShard.

• A Analogia: Imagine que, em vez de usar uma régua para cortar a pizza em fatias iguais, você tem um cortador de pizza inteligente e flexível. Ele pode cortar a pizza em pedaços de tamanhos diferentes, exatamente onde as "bordas dos blocos" (as peças especiais) estão.
• Na prática: Se um bloco de dados precisa de 32 linhas, o sistema garante que o corte nunca aconteça no meio dessas 32 linhas. Ele pode dar 10 linhas para a Equipe A, 15 para a Equipe B e 7 para a Equipe C, desde que cada bloco inteiro fique intacto em algum lugar. Isso permite que os métodos de aprendizado mais avançados funcionem sem que os programadores precisem reescrever todo o código do modelo.

3. O Planejamento: O Mestre de Obras (Algoritmo de Planejamento)

Agora que temos cortes flexíveis, como organizamos isso para não criar uma bagunça na comunicação? Se cada equipe tiver um pedaço de tamanho diferente, como elas trocam informações sem ficar esperando?

O veScale-FSDP usa um Algoritmo de Planejamento (como um Mestre de Obras genial).

• O Problema: É matematicamente muito difícil (um problema "NP-difícil") descobrir a melhor maneira de empacotar esses pedaços desiguais para que todos os caminhões (GPUs) saiam ao mesmo tempo e sem espaço vazio.
• A Solução: O algoritmo do veScale reorganiza os pedaços como se fosse um jogo de Tetris. Ele move os blocos de um lugar para outro e adiciona apenas o mínimo necessário de "areia" (padding) para preencher os buracos, garantindo que todos os caminhões saiam cheios e sincronizados. Isso economiza tempo de comunicação e memória.

4. O Caminhão de Carga: O "DBuffer" (Buffer Distribuído)

Antes, quando os dados eram movidos entre as equipes, eles precisavam ser copiados de um lugar para outro várias vezes (como mover caixas de um caminhão para um armazém e depois para outro caminhão). Isso gastava energia e tempo.

O veScale introduz o DBuffer, que é como um caminhão de carga com uma esteira rolante integrada.

• Em vez de descarregar e recarregar as caixas (copiar dados), o sistema permite que as equipes acessem os dados diretamente onde eles estão, sem cópias desnecessárias. Isso reduz o "atrito" e a memória usada, permitindo que o sistema escale para dezenas de milhares de GPUs sem explodir o orçamento de memória.

Os Resultados: O que isso muda no mundo real?

Graças a essas inovações, o veScale-FSDP conseguiu:

• Mais Velocidade: Treina modelos entre 5% a 66% mais rápido do que os sistemas atuais. É como ter um time de pedreiros que trabalha 60% mais rápido sem contratar mais gente.
• Menos Memória: Usa 16% a 30% menos memória. Isso significa que você pode treinar modelos maiores no mesmo hardware ou usar hardware mais barato.
• Flexibilidade Total: Permite usar técnicas de ponta (como os otimizadores Muon e a quantização de 8 bits) sem precisar "hackear" o código do modelo. O cientista de dados pode focar na inteligência do modelo, e o sistema cuida da logística complexa.

Resumo Final:
O veScale-FSDP é a evolução da logística de treinamento de IA. Ele troca o "corte rígido e desperdiçador" por um "corte inteligente e flexível", organiza a troca de informações como um jogo de Tetris perfeito e elimina cópias desnecessárias de dados. O resultado é que podemos construir "arranha-céus" de Inteligência Artificial mais altos, mais rápidos e mais baratos do que nunca.

Resumo técnico

Resumo Técnico: veScale-FSDP

1. O Problema

O Fully Sharded Data Parallel (FSDP), também conhecido como ZeRO, é uma técnica fundamental para treinar modelos de linguagem grandes (LLMs) em milhares de GPUs. No entanto, os sistemas FSDP atuais (como DeepSpeed ZeRO e PyTorch FSDP1/FSDP2) enfrentam limitações críticas que impedem a eficiência em modelos de última geração:

• Incompatibilidade com Treinamento "Structure-Aware" (Consciente da Estrutura): Modelos modernos (ex: Gemini, DeepSeek-V3) utilizam otimizadores não elementares (como Shampoo e Muon) e técnicas de quantização em blocos (Block-wise Quantization). Esses métodos exigem que tensores sejam manipulados como blocos atômicos (matrizes 2D ou blocos de quantização).
• Conflito de Fragmentação: Os sistemas FSDP existentes utilizam formatos de particionamento fixos (elemento a elemento ou linha a linha). Isso cria fronteiras de particionamento que não se alinham com os blocos estruturais necessários, forçando os desenvolvedores a modificar intrusivamente o código do modelo/otimizador ou o sistema a lidar com complexos mecanismos de preenchimento (padding), verificação de limites e comunicação extra.
• Ineficiência de Memória e Comunicação: Implementações atuais sofrem com:
  • Cópia de Dados Interleaved: O FSDP2 (PyTorch) exige cópias de dados intercaladas para alinhar endereços de memória, consumindo até 14% do tempo de iteração.
  • Fragmentação de Memória: A alocação não determinística e o padding excessivo (especialmente em modelos MoE) levam ao desperdício de memória GPU, limitando a escalabilidade a dezenas de milhares de GPUs.
  • Comunicação Desalinhada: Buffers de comunicação desalinhados causam ineficiência nas operações coletivas (AllGather, ReduceScatter).

2. Metodologia e Arquitetura

O veScale-FSDP reinventa o sistema FSDP do PyTorch combinando flexibilidade de particionamento com um algoritmo de planejamento estruturado. A arquitetura baseia-se em três pilares principais:

A. RaggedShard (Formato de Particionamento Flexível)

• Introduz um novo formato de placement para DTensor (Tensor Distribuído) chamado RaggedShard.
• Diferente do particionamento uniforme (elemento a elemento), o RaggedShard permite granularidade arbitrária e distribuição desigual de blocos contíguos entre dispositivos.
• Benefício: Permite que blocos de quantização ou matrizes inteiras de otimizadores (como Muon) permaneçam intactos em um único dispositivo, eliminando a necessidade de comunicação para acessar partes de um bloco.
• Composabilidade: É ortogonal a outras estratégias de paralelismo (TP, EP) e pode ser combinado com elas sem quebrar a semântica do DTensor.

B. Algoritmo de Planejamento (Planning Algorithm)

• Para maximizar a eficiência da comunicação com tensores RaggedShard, o sistema agrupa esses tensores em buckets (agrupamentos).
• O problema de agrupar tensores de tamanhos e granularidades variadas para minimizar o padding e equilibrar a carga é formulado como um problema de otimização NP-difícil (reduzível ao problema de Partição).
• Solução Heurística: O authors desenvolveram um algoritmo de programação dinâmica (DP) em tempo polinomial que reorganiza os tensores e insere padding apenas entre eles (e não dentro dos tensores), garantindo que os blocos permaneçam contíguos e alinhados. Isso minimiza o volume de comunicação e o uso de memória.

C. Distributed Buffer (DBuffer)

• Uma primitiva de baixo nível que fornece um buffer global unificado para o grupo de tensores.
• Acesso Zero-Copy: Mapeia diretamente os dados dos tensores para fatias do buffer global, eliminando cópias de dados antes e depois das comunicações coletivas.
• Fusão de Operações: Agrupa operações de kernel (como add, scale, zero) em nível de grupo antes da comunicação, reduzindo a sobrecarga de bloqueio.
• Gerenciamento de Memória: Utiliza alocação em lote (batched allocation) e dependências de stream determinísticas para reduzir a fragmentação e liberar memória de forma previsível.

3. Principais Contribuições

1. Suporte Nativo a Otimizadores e Quantização Estruturais: O veScale-FSDP suporta nativamente otimizadores baseados em matriz (Muon, Shampoo) e quantização em blocos (ex: 8-bit Adam) sem exigir modificações intrusivas no código do modelo.
2. Flexibilidade de Granularidade: O formato RaggedShard permite definir tamanhos de bloco personalizados, alinhando-se perfeitamente com as necessidades de algoritmos de ponta.
3. Alta Performance em Escala: O sistema foi projetado para escalar eficientemente para dezenas de milhares de GPUs (até 10K+), mantendo a eficiência de comunicação e memória.
4. Integração Transparente: Funciona como um módulo Python plug-and-play, substituindo o backend do FSDP2 do PyTorch, mantendo a mesma API (fully_shard).

4. Resultados Experimentais

Os autores avaliaram o sistema em modelos densos (LLaMA-3-70B, GPT-OSS-120B) e modelos MoE (até 2.4 Trilhões de parâmetros) em clusters de até 10.000 GPUs.

• Throughput (Vazão): O veScale-FSDP alcançou 5% a 66% de aumento na vazão em comparação com DeepSpeed ZeRO, FSDP1, FSDP2 e Megatron-FSDP.
  • Em modelos MoE, a melhoria foi de 11% a 66%, devido à eliminação de padding excessivo e comunicação fragmentada.
• Uso de Memória: Redução de 16% a 30% no uso de memória de pico por GPU.
  • Isso permite treinar modelos maiores no mesmo hardware ou evitar erros de Out-Of-Memory (OOM) que paralisam outros sistemas (ex: FSDP2 falha em 256 GPUs para GPT-OSS, enquanto veScale-FSDP escala).
• Escalabilidade:
  • Weak Scaling: Escalabilidade quase linear até 8.000 GPUs.
  • Strong Scaling: Escalabilidade linear até 10.000 GPUs para batches globais grandes.
  • Model Scaling: Sucesso no treinamento de modelos de 2.4T parâmetros em apenas 1.000 GPUs sem degradação de desempenho.
• Casos de Uso Específicos:
  • 8-bit Adam: Implementado com poucas linhas de código, mantendo a convergência e reduzindo drasticamente o estado do otimizador.
  • Muon Distribuído: Suporte nativo para o algoritmo Newton-Schulz, alcançando 47.3% de MFU (Model FLOPS Utilization) em 256 GPUs.

5. Significado e Impacto

O veScale-FSDP representa um avanço significativo na infraestrutura de treinamento de IA:

• Desacoplamento de Modelo e Sistema: Permite que pesquisadores inovem em arquiteturas de modelos e otimizadores sem se preocupar com as restrições de particionamento do sistema de treinamento.
• Viabilidade de Modelos de Escala Trilhões: Torna economicamente e tecnicamente viável treinar modelos com trilhões de parâmetros em clusters massivos, resolvendo gargalos de memória e comunicação que limitavam os sistemas anteriores.
• Ecossistema Open Source: O código do RaggedShard foi aberto no GitHub, e a tecnologia já está sendo considerada para integração no roadmap oficial do PyTorch, prometendo beneficiar toda a comunidade de IA.

Em resumo, o veScale-FSDP resolve o conflito entre a necessidade de flexibilidade estrutural dos novos algoritmos de IA e a necessidade de eficiência de hardware em escala massiva, estabelecendo um novo padrão para o treinamento distribuído de LLMs.