Todo Op Forward Recebe Um Gêmeo Backward

A Disciplina de Pareamento

microgpt_cuda.cu envia dois kernels CUDA para cada operação: um que computa a saída forward, um que computa gradientes de entrada dados gradientes de saída. O pareamento é um-para-um:

Oito pares de operações cobrem o transformer completo. Mais alguns kernels utilitários: k_grad_norm_partial, k_grad_norm_final, k_grad_scale para clipping de gradientes (veja atividade 75).

O Trabalho de um Kernel Backward

Dado o gradiente que flui das camadas posteriores (grad_output), um kernel backward computa:

1. grad_input: o gradiente em relação ao tensor de entrada da operação. Este é passado adiante para trás.

2. grad_weight: o gradiente em relação aos parâmetros aprendíveis na operação. Este vai para o estado do otimizador.

Ambos são computados em um único lançamento de kernel. Threads CUDA cooperam em tiles do tensor de gradiente em paralelo.

Tensores Salvos

O cálculo backward frequentemente precisa de valores do forward pass. Por exemplo, k_layernorm_bwd precisa da média e variância computadas durante o forward; k_mlp_bwd precisa da pré-ativação GELU. O motor de treinamento armazena estes em buffers dedicados durante o forward, depois os lê durante o backward.

Custo de memória: aproximadamente o mesmo formato que a saída do forward para cada tensor salvo. Para ANDREA-120M com batch=8, seq=1024, d_model=768, um tensor salvo é 8 × 1024 × 768 × 4 bytes = 25 MB. Em 12 camadas & múltiplos tensores salvos por camada, as ativações dominam a VRAM durante o treinamento (~5-10 GB em uma placa de 24 GB).

Rastreando Um Passo de Backward

Onde os Gradientes Vivem na Memória

Um Tensor de Gradiente Por Tensor de Peso

Todo tensor de peso aprendível no ANDREA-120M tem um tensor de gradiente correspondente de forma idêntica. Para cada bloco:

W_Q       [768, 768]     ↔   grad_W_Q       [768, 768]
W_K       [768, 768]     ↔   grad_W_K       [768, 768]
W_V       [768, 768]     ↔   grad_W_V       [768, 768]
W_O       [768, 768]     ↔   grad_W_O       [768, 768]
W_1       [768, 3072]    ↔   grad_W_1       [768, 3072]
W_2       [3072, 768]    ↔   grad_W_2       [3072, 768]
LN1.gamma [768]          ↔   grad_LN1.gamma [768]
LN1.beta  [768]          ↔   grad_LN1.beta  [768]
LN2.gamma [768]          ↔   grad_LN2.gamma [768]
LN2.beta  [768]          ↔   grad_LN2.beta  [768]

Mais token embeddings, position embeddings e uma layer norm final. A memória total do buffer de gradientes corresponde à memória dos pesos: ~120M floats, ~480 MB em FP32, ~240 MB em FP16.

Acumulação Através de Microbatches

O batch_size = 8 da ANDREA cabe na VRAM em FP16. Batches efetivos maiores requerem acumulação de gradientes: execute múltiplas passadas forward+backward em batches pequenos, somando gradientes no mesmo buffer, depois dê um passo do otimizador.

for microbatch in range(n_microbatches):
forward(microbatch)
backward()           # ADICIONA aos buffers de gradientes, não sobrescreve
scale_grads(1.0 / n_microbatches)  # média entre microbatches
optimizer_step()
zero_grads()             # reinicia para o próximo passo de treinamento

Os kernels de backward usam semântica +=, não =. Cada chamada adiciona contribuições de gradiente ao buffer existente; o buffer mantém a soma acumulada até que zero_grads() o limpe.

O Estado do Otimizador

AdamW (atividade 73) mantém dois buffers adicionais por peso: primeiro momento m & segundo momento v. Memória total durante o treinamento:

weights:    1× contagem de pesos
gradients:  1× contagem de pesos
Adam m:     1× contagem de pesos
Adam v:     1× contagem de pesos
ativs. salvas: ~2-4× dependendo de camadas & lote
──────────────────────────────────────────
total:      ~6-8× contagem de pesos

ANDREA-120M em FP16: ~240 MB × 4 buffers (peso, grad, m, v) + ~5-10 GB ativações = ~10-12 GB total. Confortavelmente abaixo do limite de 24 GB da RTX 4090. ANDREA-12M treinado em 1.4 GB; o escalonamento de parâmetros 10× traz ~10× memória.

Dimensionando Buffers de Gradiente

Controle Total de Memória & Precisão

O Custo dos Frameworks Genéricos

PyTorch & JAX tornam o autograd conveniente: escreva código Python, obtenha gradientes automaticamente. O custo: uma camada de despacho genérica entre seu código & CUDA. Toda operação passa por overhead do interpretador Python, bookkeeping do framework, & seleção dinâmica de kernels. Para treinar um pequeno modelo de linguagem em uma GPU, esse overhead importa.

Custos concretos que o ANDREA evita:

1. Latência do interpretador Python. Cada operação PyTorch atravessa a fronteira Python/C++. Para ~100 lançamentos de kernel por passo de treinamento a ~9 passos/min, isso são ~900 travessias de fronteira por minuto. O despacho em nível C elimina isso.

2. Imprevisibilidade do alocador do framework. O alocador de cache do PyTorch oferece bom throughput em média, mas memória máxima imprevisível. O motor de treinamento do ANDREA pré-aloca todos os buffers na inicialização; sem realocação durante o treinamento, sem fragmentação, sem OOMs surpresa no passo 100K.

3. Seleção genérica de kernel. O PyTorch escolhe kernels em tempo de execução via heurísticas. O ANDREA escolhe kernels em tempo de compilação, ajustados aos tamanhos de tile do tensor core RTX 4090.

4. Infraestrutura de precisão mista. O caminho cuBLAS FP16 do ANDREA-120M e os experimentos de tensor core FP8 E4M3 do ANDREA requerem controle preciso sobre quais tensores vivem em qual precisão. Frameworks genéricos expõem esse controle através de APIs em camadas; escritas personalizadas em CUDA o escrevem diretamente.

O Tradeoff

Custos do CUDA personalizado: mais código para escrever, mais bugs para encontrar, sem ecossistema comunitário. O microgpt_cuda.cu do ANDREA tem ~6000 linhas de CUDA escrito à mão que levou meses para depurar. Cada nova operação requer escrever um kernel forward, um kernel backward e testes.

O que o ANDREA ganha:

- Reprodutibilidade total. O pipeline de treinamento é um binário C mais um proxy Python. Sem derivação de versão entre lançamentos do PyTorch, sem incompatibilidades de versão CUDA com wheels do framework.

- Retomadas bit-exatas. SIGTERM aciona a escrita de um checkpoint que captura cada tensor exatamente como a GPU o vê. O retomar pega a mesma trajetória de perda em que a execução estava.

- Memória previsível. ANDREA-120M treinado por 200K passos sem OOMs. A memória foi contabilizada na inicialização do engine.

- Acesso direto ao hardware. Tamanhos de tile do tensor core, configurações FP8 E4M3, cópias de memória assíncronas: tudo diretamente acessível no CUDA, opaco em frameworks genéricos.

Reprodutibilidade Como Missão

A seção 9 do whitepaper ANDREA lista a pilha completa de reprodutibilidade:

Motor de treinamento: microgpt/microgpt_cuda.cu
Proxy de treinamento: microgpt/training_proxy.py
Configurações de experimento: experiments/ANDREA-*-TRAIN.json
Pipeline de dados: scripts/pull-hermes3.py, scripts/prep-megachat.py
Painel: scripts/live-loss-dashboard.html
Especificação do Bandit: docs/FIREHOSE-BANDIT.md
Documentação do modelo: docs/ANDREA.md

Requisito de hardware: uma GPU NVIDIA com ≥8 GB de VRAM (RTX 3060 ou melhor). Qualquer pessoa pode reproduzir o ANDREA-12M a partir desses artefatos. O caminho CUDA personalizado é parte do motivo: sem congelamentos de versão de framework, sem surpresas de dependências daqui a cinco anos.

Sinais & Checkpoints

O loop de treinamento CUDA responde a dois sinais POSIX:

- SIGTERM: grava um checkpoint imediato, depois sai. Usado ao parar o treinamento de forma limpa.

- SIGUSR1: escreve um checkpoint imediato, continua o treinamento. Usado durante o pivot de polimento no v3 para capturar o estado sem interromper a execução.

Formato do checkpoint: [int32 step][int32 n_params][n_params × float32 weights][n_params × float32 m][n_params × float32 v]. Contador de passos, contagem de pesos, depois pesos seguidos pelos momentos de Adam. Retoma bit-exatamente. O proxy arquiva .samples.json & .state.json separadamente no polimento; .loss.json nunca é arquivado (ele acumula o histórico completo do treinamento).

Por Que Não PyTorch