antonio leandro

ia generativa

PaLM: Scaling Language Modeling with Pathways

paper · Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, et al. ·

a tese

escala densa não tinha saturado em 2022: 540 bilhões de parâmetros mais chain-of-thought no prompt bateram o estado da arte finetunado em raciocínio, sem verificador, sem arquitetura de tarefa

o que fica

  1. Hardware FLOPs utilization mede o que a implementação faz, não o que o modelo precisa; o paper propõe model FLOPs utilization (MFU), que só conta forward e backward e permite comparar sistemas diferentes — PaLM 540B ficou em 46,2% de MFU contra 21,3% do GPT-3.
  2. O mesmo modelo vai de 17% para 54% no GSM8K só trocando o prompt por chain-of-thought: a capacidade de raciocínio já estava nos pesos e o que faltava era a forma de pedir.
  3. Cerca de 25% das tarefas do BIG-bench melhoraram mais de 10 pontos acima da projeção log-linear feita com 8B e 62B — escala produz saltos descontínuos, não só a curva suave da lei de potência.
  4. PaLM empata com o Codex 12B em HumanEval pass@1 tendo visto 2,7 bilhões de tokens de Python contra 100 bilhões: modelo grande transfere de outras linguagens e de texto natural em vez de precisar do dado específico.
  5. Os picos de loss no modelo de 540B não vinham de dado ruim — treinar os mesmos batches a partir de um checkpoint anterior não reproduzia o pico. O gatilho é a combinação de batch com estado específico dos parâmetros, e a correção foi voltar 100 passos e pular de 200 a 500 batches.
  6. A taxa de memorização depende mais de duplicação que de tamanho: exemplos vistos uma vez são reproduzidos literalmente em 0,75% dos casos, e os vistos mais de 500 vezes passam de 40%.

o problema

Em 2022 havia duas paredes distintas. A primeira era de engenharia: quem queria treinar um modelo denso muito grande cabia num único pod de TPU ou aceitava pipeline parallelism para atravessar a rede do datacenter. Pipeline traz a bolha — máquinas paradas enquanto o pipeline enche e esvazia — e obriga a recarregar pesos da memória a cada micro-batch. Gopher usou pipeline entre pods de TPU v3, Megatron-Turing usou pipeline em 2.240 A100. O teto prático era 4.096 chips.

A segunda parede era de capacidade. Modelos de linguagem grandes iam bem em classificação e completamento, mas travavam em qualquer coisa que exigisse vários passos de inferência. O estado da arte em problemas de matemática escolar vinha de empilhar coisas: finetuning no dataset, calculadora externa e um verificador treinado para escolher entre candidatos. A leitura corrente era que raciocínio multi-passo não era um problema de escala, e que a curva de ganhos por parâmetro estava achatando.

a ideia

PaLM ataca as duas com o mesmo movimento: treinar um Transformer decoder-only denso de 540 bilhões de parâmetros em uma passada única sobre 780 bilhões de tokens, e depois avaliar quase tudo em few-shot, sem tocar nos pesos.

A parte de sistema é o Pathways. Em vez de fatiar o modelo ao longo de um pipeline, cada pod de TPU v4 guarda uma cópia inteira dos parâmetros e os dois pods fazem data parallelism entre si. Um único cliente Python despacha metade do batch para cada pod, cada um calcula gradientes na sua metade, os dois trocam gradientes pela rede e aplicam a mesma atualização. O resultado são parâmetros bit a bit idênticos nos dois lados e nenhuma bolha.

como funciona

Dentro de cada pod são 3.072 chips, com cada tensor de peso particionado em 12 vias de model parallelism e 256 vias de data parallelism totalmente fragmentado. No forward, os pesos são reunidos por all-gather no eixo de dados e uma ativação fragmentada é guardada por camada; o resto é rematerializado no backward, porque o batch maior que isso permite compensa o custo.

A transferência entre pods é o ponto delicado. Cada par de hosts troca cerca de 1,3 GB de gradientes por passo, e como todos terminam o backward mais ou menos juntos, isso vira um burst agregado de 81 Tbps. A saída foi quebrar o gradiente em pedaços menores e espalhá-los por vários fluxos e links. Com isso o throughput com dois pods é 1,95 vez o de um só — 97% do escalonamento perfeito. O que falta para 2x é a ausência de sobreposição entre o backward e a redução de gradiente entre pods, e os autores dizem que isso fica para depois.

A arquitetura é padrão com seis desvios: SwiGLU no MLP; camadas paralelas, onde attention e MLP leem o mesmo LayerNorm em vez de encadear (15% mais rápido, porque as multiplicações de entrada se fundem); multi-query attention, com key e value compartilhados entre heads, que não muda qualidade e barateia muito o decoding autorregressivo; RoPE; embeddings de entrada e saída compartilhadas; e nenhum bias em kernel ou layer norm, o que estabilizou o treino. O vocabulário SentencePiece tem 256 mil tokens, preserva whitespace e sempre quebra número em dígitos individuais.

Otimizador Adafactor sem fatoração, sequência de 2.048 tokens sem padding, batch subindo de 512 para 1.024 e depois 2.048 ao longo do treino, z_loss auxiliar para manter o normalizador do softmax perto de zero, sem dropout. Todo o pipeline é bitwise determinístico: reiniciar de um checkpoint reproduz exatamente a mesma trajetória.

o que isso custou

Os próprios autores marcam onde não fecha. No SuperGLUE finetunado, PaLM 540B fica atrás do ST-MoE-32B, um encoder-decoder com objetivo de span corruption e ordens de grandeza menor; em TyDiQA, perde para o ByT5 XXL. A desvantagem arquitetural do decoder-only em finetuning discriminativo é reduzida pela escala, não eliminada.

Geração fora do inglês continua fraca: PaLM processa entrada não-inglesa melhor do que produz saída não-inglesa, e MLSum não alcança o estado da arte nem com finetuning. Em BIG-bench, apesar de superar a média humana no agregado, o humano médio ainda vence em 35% das 150 tarefas, e tarefas como navigate e mathematical_induction quase não melhoram de 62B para 540B.

No MMLU, o ganho sobre o Chinchilla 70B é de menos de 2 pontos — com quase oito vezes mais parâmetros. Os picos de loss ficaram sem explicação de princípio, só com a gambiarra de pular batches. Memorização cresce com escala: 2,4% das amostras no modelo de 540B contra 1,6% no de 8B. E a análise de viés mostra Islam co-ocorrendo com terrorist, violent e radical de forma estável em todas as escalas — não é algo que crescer conserta.

onde isso aparece hoje

MFU virou a métrica padrão para reportar eficiência de treino, justamente porque é independente de implementação. O combo SwiGLU, RoPE, sem bias e attention com key/value compartilhados virou receita de base para modelos decoder-only abertos que vieram depois. E chain-of-thought, que aqui aparece como técnica de prompt avaliada em sete benchmarks, deixou de ser truque e virou parte do vocabulário — a linha direta entre esse resultado e todo o trabalho posterior sobre raciocínio em tempo de inferência é curta.

lido na íntegra por pipeline de llm, revisado por antonio leandro antes de publicar ·