Como construir um LLM completo do zero em Rust: do pré-treino ao GRPO, Flash Attention e visão

Published on: 2026-07-18
Post image
pt rust llm ia machine-learning candle grpo flash-attention multimodal tutorial

Chamar uma API de LLM é fácil. Entender — e construir — um modelo de linguagem do zero é outro nível. E dá para fazer isso inteiro em Rust puro, sem Python nem PyTorch. Este guia mostra o caminho: quais peças você precisa, em que ordem, e o que cada uma resolve.

Como planta de referência, usamos o aarambh-ai (v2.0.0), um projeto open source que implementa a stack completa em Rust — do tokenizer ao GRPO, da quantização à visão. A ideia não é decorar código, e sim ter o mapa mental para montar (ou ler) um LLM de verdade.

O mapa: as peças de um LLM

Um LLM completo é a soma de blocos bem definidos. Construí-lo é montar esses blocos na ordem certa:

  • Tokenizer (transforma texto em números) e dados;
  • Arquitetura: um transformer decoder-only (a família do GPT/LLaMA);
  • Pré-treino: ensinar o modelo a prever o próximo token;
  • Alinhamento: SFT e reforço (GRPO/DPO) para ele ser útil e raciocinar;
  • Otimização: Flash Attention e kernels para treinar/rodar rápido;
  • Quantização e inferência: fazer caber e gerar token a token;
  • Segurança: os guardrails na entrada e na saída.

Uma dica de organização vale ouro: separe cada peça em um crate próprio. O aarambh-ai usa um workspace em camadas (17 crates + o binário), e é isso que mantém o projeto legível:

Workspace em camadas (17 crates + CLI):
  core          → config, erros, traits
  tokenizer     → BPE implementado do zero
  data          → datasets e batching
  nn / kernel   → camadas e kernels (CUDA + SIMD)
  model         → o transformer decoder-only
  weights/quant → pesos, SafeTensors/GGUF, quantização
  train         → pré-treino (+ multi-GPU)
  finetune      → SFT, LoRA/QLoRA/DoRA, GRPO, DPO
  distill       → distillation on-policy
  inference     → geração com streaming
  vision        → encoder estilo CLIP, VQA
  safety        → guardrails
  selflearn     → auto-aprendizado
  eval / serve  → avaliação e serving
  aarambh-ai    → binário CLI

Antes de escrever qualquer linha, escolha a escala pelo hardware: 25M roda em CPU; 117M em uma T4; 360M em uma P100; 1,3B em uma A100. Comece pequeno.

A arquitetura que você vai montar

O coração é um transformer decoder-only. As escolhas modernas (estilo LLaMA) que valem adotar:

  • RMSNorm no lugar do LayerNorm: normalização mais simples e barata, que estabiliza o treino.
  • RoPE (Rotary Positional Embedding): codifica a posição de cada token “girando” os vetores — ajuda em sequências longas.
  • GQA (Grouped-Query Attention): cabeças compartilham chaves e valores, o que reduz o KV cache e a memória na inferência.
  • SwiGLU na feed-forward: ativação que rende mais qualidade por parâmetro.
  • Tied embeddings: reaproveite a matriz de embedding na saída e economize parâmetros.

Passo 1: tokenizer e dados

Comece por um tokenizer BPE feito do zero — é ele que fatia o texto em tokens. Para os dados, datasets clássicos já dão um bom caminho de aprendizado: WikiText-103 (texto geral), Shakespeare (estilo) e GSM8K (matemática, para treinar raciocínio depois).

Passo 2: o pré-treino

Aqui o modelo aprende a prever o próximo token. Os ajustes que fazem diferença:

  • otimizador AdamW com hiperparâmetros estilo LLaMA (o segundo momento, beta 2, em 0,95 — e não no padrão 0,999);
  • cosine scheduler com warmup linear para a taxa de aprendizado;
  • gradient accumulation e treino multi-GPU em BF16, para escalar com pouca memória por placa.

O progresso ilustra bem o que esperar: em 50 passos, sai uma sopa de letras (“Secondprophebreathety...”); em 5.000 passos, o modelo já completa trechos de Shakespeare de forma coerente. Ter uma suíte de avaliação desde cedo é o que te diz se está no caminho.

Passo 3: alinhar com SFT, GRPO e DPO

Um modelo pré-treinado sabe completar texto, mas ainda não é “útil”. O alinhamento resolve isso em camadas:

  • SFT (Supervised Fine-Tuning): ensina o formato de resposta a partir de exemplos.
  • GRPO (Group Relative Policy Optimization), o reforço popularizado pelo DeepSeek-R1: compara várias respostas do próprio modelo para a mesma pergunta e reforça as melhores — sem “modelo de recompensa” separado, ideal para tarefas com resposta verificável (como matemática). Na referência, rendeu de 15% a 30% mais acerto que a base de SFT.
  • DPO (Direct Preference Optimization): alinha a partir de pares “resposta preferida vs. rejeitada”.

É nessa fase que entra o “motor de pensamento”, com orçamentos de raciocínio configuráveis (256, 1.024 ou 4.096 tokens).

Otimizar a atenção: Flash Attention v2 e kernels

Atenção é a operação mais cara do transformer. Implemente Flash Attention v2 com online softmax — ela calcula a atenção sem materializar a matriz inteira, economizando VRAM. Na referência, isso rendeu 3,5× no tempo de treino ponta a ponta. Escrever kernels próprios (CUDA e SIMD) também compensa: um RMSNorm fundido ficou 2,8× mais rápido que a versão ingênua.

Fazer caber em hardware modesto: quantização e LoRA

Para rodar e ajustar o modelo sem uma GPU enorme, use quantização: INT8, GPTQ INT4 (com decomposição de Cholesky para estabilizar a Hessiana e evitar NaN), AWQ e o formato GGUF Q4_K_M — que reduz o modelo Tiny a 13 MB. Para fine-tuning barato, use LoRA, QLoRA ou DoRA: dá para treinar QLoRA em 8 GB de RAM com pico de 400 MB. E na inferência, KV-cache streaming gera token a token.

Ir além do texto: visão (multimodal)

Quer um modelo que enxerga? O caminho clássico é congelar um bom encoder de imagem (estilo CLIP) e treinar só um projector que liga a imagem ao espaço do texto; depois é só fundir a imagem no contexto e treinar em VQA (perguntas sobre imagens). O aarambh-ai já traz esse crate de visão — em Rust puro também.

Bônus: auto-aprendizado

Um passo mais avançado é o laço de auto-aprendizado: o modelo melhora a partir das próprias saídas, guardadas em um replay buffer, com reforço online durante a inferência e fine-tuning periódico. É poderoso — e exige cuidado para o modelo não degenerar.

Não esqueça a segurança

Um LLM que vai para produção precisa de guardrails: detecção de prompt injection e jailbreak, detecção e redação de PII (e-mail, telefone, CPF, cartão com verificação de Luhn), varredura de toxicidade na saída e log de auditoria com hash SHA-256. Melhor construir isso desde o começo do que remendar depois.

Os detalhes que fazem funcionar (gotchas)

O que separa um brinquedo de um projeto sério são os detalhes que só aparecem na prática:

  • o beta 2 do AdamW em 0,95 rende de 2 a 3 pontos de perplexity a menos que o padrão 0,999;
  • use Cholesky na inversão da Hessiana para evitar NaN na quantização GPTQ;
  • o backward do QLoRA precisa desquantizar dentro do grafo de autograd;
  • o KV cache precisa ter teto no tamanho máximo da sequência, senão estoura a memória (OOM).

Por que Rust para isso

Dá para construir um LLM em várias linguagens — então por que Rust? Porque ele encaixa bem no ciclo de ML: memória determinística (sem pausas de garbage collector), erros pegos em tempo de compilação, embutibilidade (o modelo quantizado roda em CPU, num binário único, sem runtime de Python) e controle full-stack numa linguagem só — do tokenizer ao kernel SIMD. É a opção natural quando o objetivo é empacotar e rodar modelos com controle.

Por onde começar

Teoria sem prática não gruda. Um roteiro concreto para sair do zero:

  • Clone o aarambh-ai (link nas referências) — ele é open source e vem com um binário de linha de comando.
  • Leia os crates na ordem das camadas: coretokenizermodeltraininference. É a espinha dorsal.
  • Treine o modelo Tiny (25M) na sua CPU para ver o pipeline rodando de ponta a ponta.
  • Depois avance para SFT e GRPO e experimente quantizar para rodar leve.

Conclusão

Construir um LLM do zero deixou de ser mistério: é um conjunto de peças bem definidas — tokenizer, transformer, pré-treino, alinhamento, otimização, quantização e segurança — montadas na ordem certa. Fazer isso em Rust puro te dá um mapa completo e um artefato que roda em qualquer lugar. Se você quer parar de só “chamar a API” e entender de verdade, comece pequeno, treine o Tiny na CPU e vá subindo. 🚀

Referências