Data parallelism, ZeRO e FSDP

1

Conceito

Um acelerador processa batch limitado e guarda estado limitado. Data parallelism eleva throughput dando a vários workers réplicas idênticas do modelo e exemplos diferentes. Se o rank $r$ calcula $g_r$, uma redução coletiva forma média ou soma usada por todos. Começar com pesos iguais e aplicar gradientes iguais mantém sincronização.

O global batch combina ranks e gradient accumulation. Com $P$ ranks, microbatch local $b$, comprimento $T$ e $a$ steps acumulados, o batch contém $PbaT$ tokens, ignorando padding. Aumentar workers sem ajustar isso muda otimização, não só velocidade.

Replicação custa memória. Adam mixed-precision pode manter parâmetros de baixa precisão, gradientes, master weights em precisão maior e dois tensors de moments. Bytes exatos dependem de implementação e dtype, mas optimizer state pode superar o peso de inference. Data parallelism comum duplica tudo em todo rank.

ZeRO decompõe redundância em estágios. Stage 1 divide optimizer state. Stage 2 divide também gradientes. Stage 3 divide parâmetros. Cada rank possui uma fração persistente, enquanto collectives disponibilizam peças necessárias. Os nomes descrevem ownership de memória, não novo optimizer.

Fully Sharded Data Parallel envolve módulos nesse padrão. Antes do forward, ranks fazem all-gather dos shards para materializar parâmetros completos necessários localmente. Depois, podem fazer reshard. No backward, a comunicação reconstrói parâmetros quando preciso e reduce-scatter envia shards de gradientes aos donos. Comunicação e compute podem se sobrepor quando ordem e rede permitem.

Wrapping policy importa. Unidade grande gera pico alto de all-gather e pouca sobreposição. Pequena cria muitos collectives e deixa latency dominar. Fronteiras de Transformer blocks são comuns por oferecerem compute substancial. Prefetch do próximo bloco esconde latency, mas aumenta memória simultânea.

Checkpoint também vira distribuído. Reunir full state dict pode exceder host memory. Checkpoints sharded deixam ranks escreverem seus tensores e metadata reconstrói nomes e shapes globais. Retomar com outro world size exige formato que reshard. Salvar só pesos não continua exatamente: optimizer, scheduler, scaler, RNG e data-loader state importam.

Falhas precisam ser coletivas. Se um rank encontra dado ruim enquanto outros entram em all-reduce, o job pode travar. Timeouts, propagação coordenada, logs rank-aware e batch IDs reproduzíveis são necessários. Topologia da rede também conta: links rápidos no node e lentos entre nodes favorecem collectives hierárquicos.

Eficiência precisa ser medida por fase. Registre tempo de forward, backward, all-gather, reduce-scatter, optimizer step e data loading, além do pico de memória por rank. Uma média global esconde stragglers: se um único rank recebe batches mais longos ou hardware degradado, todos esperam no collective. O melhor wrapping é aquele que melhora throughput realizado sem tornar o pico imprevisível.

A visualização deve distinguir ownership de materialização temporária. Mostre shards persistentes coloridos, depois um all-gather em torno de um bloco, compute em dados diferentes e reduce-scatter. Exibir o modelo completo permanentemente em todo rank representaria mal o ganho.

O modelo estável é compute replicado sobre dados diferentes com estado progressivamente sharded. Data parallelism compra throughput; ZeRO e FSDP removem memória redundante. O preço é comunicação, orquestração e semântica mais exigente de checkpoint e falha.