TatuEngine: o retorno do lm_head — a auditoria que destravou a Fase 4

A pendência que ninguém esqueceu

No post do codec, eu deixei um spoiler no final: “o lm_head volta a aparecer”. E voltou.

A história curta: o codec de blocos ternários encolheu o modelo de 27GB em 245MB, mas a validação por tensor acusou divergência no lm_head — a camada que mapeia os estados internos para o vocabulário. Enquanto in_proj saiu perfeito (diferenças no nível de ULP) e out_proj foi explicado por ordem de acumulação, o lm_head ficou como pendência aberta no backlog.

Duas semanas depois, eu voltei para ele com um script de auditoria — e o que eu achei mudou o rumo do projeto inteiro.

A auditoria

O script lm_head_fix.py nasceu com uma pergunta simples: o que tem dentro do lm_head depois de toda essa compressão?

A primeira coisa que ele faz é carregar o checkpoint original e medir estatísticas da matriz de pesos:

w = model.lm_head.weight.data
print(f"lm_head.weight: mean={w.float().mean().item():.4f} "
      f"std={w.float().std().item():.4f} "
      f"max|w|={w.float().abs().max().item():.4f}")
print(f"zeros absolutos: {(w.float().abs() < 1e-12).sum().item()} / {w.numel()}")

E depois compara com a embedding:

ew = model.embed.weight.data
print(f"embed.weight: mean={ew.float().mean().item():.4f} "
      f"std={ew.float().std().item():.4f}")
print(f"compartilhados (tied): {torch.equal(ew, w)}")

O que a auditoria revelou: a distribuição dos pesos do lm_head estava desalinhada com a da embedding. Em modelos causais com peso compartilhado (tied embeddings), a cabeça de vocabulário deveria acompanhar a escala da entrada. Depois das transformações do codec e dos checkpoints intermediários, a escala tinha derivado — e era isso que fazia a validação divergir.

A correção: re-escala com Xavier

A decisão foi cirúrgica: em vez de mexer no modelo inteiro, re-inicializar apenas o lm_head com uma escala pequena, como se fosse uma cabeça nova. Usando Xavier uniform — a mesma inicialização clássica de uma camada recém-nascida:

new_w = torch.empty_like(w)
fan_in, fan_out = w.shape[1], w.shape[0]
bound = math.sqrt(6.0 / (fan_in + fan_out))
new_w.uniform_(-bound, bound)
lh.weight.data.copy_(new_w)

O teste de sanidade veio na sequência: rodar um forward com uma frase real e olhar os 5 maiores logits — para conferir que a cabeça ainda produzia saídas com magnitude razoável, e não um mar de NaN:

top5_vals, top5_idx = logits[-1].topk(5)
for v, i in zip(top5_vals.tolist(), top5_idx.tolist()):
    print(f"  {v:10.2f} -> {repr(tok.decode([i]))}")

Depois, o reteste de loss com um batch de treino — a expectativa era um valor perto do aleatório (~10.8), porque uma cabeça re-inicializada ainda não aprendeu nada:

loss = F.cross_entropy(shift.reshape(-1, shift.size(-1)), sample[:, 1:].reshape(-1))
print(f"LOSS APOS RE-ESCALA: {loss.item():.4f} (random esperado ~10.8)")

E por fim, um teste de backward com gradientes — conferir que o gradiente flui, que não tem NaN, e que a norma total faz sentido:

loss3.backward()
total_norm = sum(g**2 for g in grad_norms) ** 0.5
print(f"grad_norm total: {total_norm:.4f}")
print(f"NaN nos grads: {any(math.isnan(g) for _, g in grad_norms)}")

O checkpoint que destravou tudo

A auditoria gerou um checkpoint novo: bitmamba_1b_lmfix.pt — o mesmo modelo, mas com a cabeça re-escalada e saudável.

E foi exatamente esse arquivo que o run_fase4.py escolheu como base. Olha a linha de configuração:

# Base .pt (modelo nu — sem treino, pesos do HF)
MODEL_PT = PROJECT_ROOT / "models" / "bitmamba_cpp" / "bitmamba_1b_lmfix.pt"

A Fase 4 — batizada de Ressurreição, autorizada pelo arquiteto em 28/08 — precisa de uma carga limpa: um checkpoint de pré-treino sem SFT contaminando a memória. O bitmamba_1b.pt original tinha a cabeça com escala derivada; o _lmfix é a versão corrigida. Sem a auditoria, a Fase 4 teria começado com uma cabeça doente e os primeiros steps de treino estariam brigando contra a escala errada.

A ordem de serviço da Fase 4:

1. Carga limpa: full_warmup/best — pré-treino, sem SFT
2. AdamW (Fused) + Parameter Groups: A_log/D @ 10x, resto @ 0.3x
3. Loop Híbrido: 0-300 pretrain puro, 300-2000 co-training pretrain:SFT
4. Telemetria: losses separadas (Pretrain vs SFT) a partir do step 301

E o treino ainda ganhou uma inspeção extra de gradientes por grupo — incluindo o lm_head separado do resto:

gnorm_lm = 0.0
for n, p_g in model.named_parameters():
    if p_g.grad is not None and "lm_head" in n:
        gnorm_lm += p_g.grad.float().norm().item() ** 2
gnorm_lm = gnorm_lm ** 0.5

Agora o lm_head tem telemetria própria no log a cada step. Depois de ser o vilão da validação do codec, ele virou uma das métricas mais observadas do treino.

Métricas

Aprendizados

1. Uma pendência anotada vira trilha, não dívida — o “lm_head divergindo” ficou documentado como spoiler no post do codec. Duas semanas depois, virou o checkpoint base da fase mais importante do projeto. Pendência registrada com contexto vale ouro.
2. Auditar a escala antes de treinar — a maioria dos problemas de treino não está na arquitetura, está na escala dos pesos de entrada. Uma cabeça fora de escala faz a loss divergir de um jeito que parece bug de código.
3. Re-inicializar é cirúrgico — não precisei reconstruir nada. Um tensor novo, uma cópia, um save. A menor intervenção que destrava o próximo passo é quase sempre a certa.
4. Telemetria separada conta a história — separar o gradiente do lm_head do resto do modelo transformou o ex-vilão em sinal de saúde diário.

A Ressurreição começou com a cabeça certa. E o spoiler do codec, enfim, se pagou.

╔══════════════════════════════════════╗
║  Samuel Medeiros                    ║
║  Senior Software Engineer           ║
║  Stack: Python · TypeScript · Rust  ║
║  Projetos: Arachne, Dogwalk,        ║
║            Capivara, TatuEngine      ║
╚══════════════════════════════════════╝