Data parallelism, ZeRO e FSDP

01 · Conceito

Conceito

O Qwen3.8-27B é distribuído como safetensors em bf16 que totalizam cerca de 54 GB. Isso é um fato de armazenamento do checkpoint, não uma receita de memória de treino nem prova de que o stack completo de serving cabe num dispositivo nominal de 80 GB: KV cache, ativações, workspaces do runtime e fragmentação também consomem memória. O treino acrescenta categorias persistentes e transitórias cujos dtypes precisam ser declarados, não inferidos do checkpoint lançado.

A lição 9.3 assume a derivação do orçamento de memória deste modelo para todo o curso. Esta lição importa dela um cenário explícito: parâmetros e gradientes bf16, master weights fp32 e dois momentos AdamW fp32, totalizando cerca de 432 GB de estado persistente para 27 bilhões de parâmetros antes de ativações e workspaces. É uma receita conservadora comum, não uma afirmação sobre a run de treino não divulgada do Qwen nem uma exigência universal; existem políticas pure-bf16 e estados de optimizer alternativos. A aritmética de ZeRO abaixo fica condicionada a esse cenário para comparar os estágios de sharding de forma consistente.

Data parallelism é o ponto de partida e não faz nada quanto a isso. Dê a vários workers réplicas idênticas e exemplos diferentes; se o rank $r$ computa o gradiente $g_r$, uma redução coletiva forma a média usada por toda réplica, e partir de pesos iguais com gradientes iguais os mantém sincronizados. Com $P$ ranks, microbatch local $b$, comprimento de sequência $T$ e $a$ passos de acumulação, o batch global contém $PbaT$ tokens — portanto adicionar workers muda a otimização, não só a velocidade. Sob replicação comum, cada rank ainda guarda todo o estado persistente assumido.

ZeRO remove a redundância em três estágios, e vale calcular o efeito num cluster concreto. Tome 64 ranks e o orçamento ilustrativo de 432 GB importado da lição 9.3.

O estágio 1 faz sharding só do estado do optimizer. Cada rank mantém parâmetros e gradientes bf16 completos, mas $1/64$ do optimizer: $324/64 = 5.06$ GB. O estado persistente por rank vira $54+54+5.06 \approx 113$ GB. Ainda impossível em cartões de 80 GB.

O estágio 2 também faz sharding dos gradientes: $54/64 = 0.84$ GB. Agora $54+0.84+5.06 \approx 59.9$ GB. Isso cabe — por pouco. Restam cerca de 20 GB para ativações, temporários, fragmentação e buffers de comunicação num modelo com hidden size 5120 e sequências longas, o que a lição 5.10 vai mostrar que não é muito.

O estágio 3 faz sharding dos parâmetros também, então o estado persistente é simplesmente o orçamento inteiro dividido pelo tamanho do mundo: $432/64 = 6.75$ GB por rank. Os parâmetros completos de um módulo existem só transitoriamente, materializados pouco antes de serem necessários.

Fully Sharded Data Parallel é esse padrão totalmente shardeado implementado em torno de módulos. Antes do forward de um módulo, os ranks fazem all-gather dos shards de parâmetro dele para reconstruir os parâmetros completos localmente; depois do uso, podem voltar a shardear. Durante o backward, a comunicação reconstrói os parâmetros de novo conforme necessário, e um reduce-scatter distribui os shards de gradiente para seus donos. Onde a ordenação e a rede permitem, essa comunicação se sobrepõe à computação.

A política de wrapping decide quão bem o estágio 3 de fato se comporta. Se a unidade shardeada for grande demais, a memória de pico do all-gather é grande e a sobreposição é limitada; pequena demais, e os coletivos ficam numerosos e a latência domina. Fronteiras de bloco de transformer são a escolha usual porque colocam compute substancial entre comunicações. Este modelo oferece uma unidade alternativa incomumente natural: suas 64 camadas são organizadas como 16 super-blocos repetidos de três camadas de Gated DeltaNet seguidas de uma camada de full attention. Fazer wrapping por super-bloco dá 16 unidades uniformes de cerca de 1,7 bilhão de parâmetros cada, o que é grosseiro mas perfeitamente balanceado; fazer wrapping por camada dá 64 unidades mais finas de dois formatos distintos. As duas opções são defensáveis, e só a medição decide.

Duas consequências operacionais decorrem disso. O checkpointing vira distribuído: um gather de state-dict completo pode exceder a memória do host, então os ranks escrevem os tensores que possuem junto com metadados que reconstroem nomes e formatos globais, e retomar com um tamanho de mundo diferente exige um formato e um loader que suportem resharding. Salvar só os pesos não basta para uma continuação exata — momentos do optimizer, posição do scheduler, loss scaler, estado do gerador de números aleatórios e cursor do data loader importam todos, e, sob a receita ilustrativa, o estado do optimizer é a maior categoria do checkpoint.

As falhas também precisam ser coletivas. Se um rank tromba com dados ruins enquanto os outros entram num all-reduce, o job trava em vez de levantar uma exceção limpa. Timeouts, propagação coordenada de erro, logs cientes de rank e identificadores reprodutíveis de batch não são luxos. A topologia de rede também molda o projeto: links rápidos dentro de um nó e links mais lentos entre nós favorecem coletivos hierárquicos, e um job de estágio 3 cujos all-gathers atravessam um tecido lento pode gastar mais tempo movendo parâmetros do que usando-os.

O modelo estável é computação replicada sobre dados diferentes, com estado cada vez mais shardeado por trás. Data parallelism compra throughput. ZeRO e FSDP removem memória redundante. Nenhum dos dois cria capacidade de graça: o preço é volume de comunicação, complexidade de escalonamento e semânticas de checkpoint e de falha muito mais exigentes.