Nanocode: Treinando agentes de codificação semelhantes ao Claude com JAX em TPUs

Nanocode é uma biblioteca que demonstra como treinar seu próprio modelo Claude Code de ponta a ponta usando Constitutional AI, seguindo a abordagem da Anthropic. Escrito inteiramente em JAX e otimizado para TPUs, ele adapta a infraestrutura do projeto nanochat de Karpathy.

Configuração e Custos de Treinamento

O modelo nanocode-d24 (1,3B de parâmetros) pode ser reproduzido em aproximadamente 9 horas em uma TPU v6e-8 a um custo de US$ 200. O modelo menor nanocode-d20 (477M de parâmetros) treina em cerca de 1,5 horas por US$ 34. O projeto recomenda usar o programa TRC do Google para acesso gratuito a TPUs preemptíveis por um mês, ou os créditos de US$ 300 do Google Cloud para novas contas.

Implementação Técnica

O processo de treinamento inclui:

• Escrever um arquivo SOUL.md para definir o alinhamento do modelo
• Definir uma interface agentica para interação com o mundo
• Gerar dados sintéticos
• Usar otimização de preferência para alinhar o modelo com o SOUL

Diferenças de Tokenização e Pré-treinamento

Embora o processo de pré-treinamento e treinamento do tokenizador seja semelhante ao nanochat, o nanocode inclui dados adicionais de codificação do The Stack-V2 em uma proporção de 1:5 tanto nas misturas de pré-treinamento quanto do tokenizador. Isso resulta em um desempenho de codificação mais forte, mas reduz a eficiência de tokenização de texto geral.

A comparação do tokenizador mostra que o nanocode alcança -50,9% melhor tokenização para código em comparação com o nanochat, enquanto o nanochat tem desempenho melhor em texto coreano (+7,9% para nanocode em notícias, -27,6% em coreano).

Comandos e Configuração

export NANOCODE_BASE_DIR="$HOME/.cache/nanocode"
export MODEL_TAG=d24
python -m data.pretrain -d fineweb-edu -n 300
python -m data.pretrain -d the-stack-v2-dedup -n 60
python -m scripts.tok_train --max-chars=2000000000
python -m scripts.tok_eval

Os modelos são treinados com uma proporção param:dados de 8, seguindo a análise da lei de escala do nanochat. Embora otimizado para TPUs, o nanocode também deve funcionar em GPUs NVIDIA prontas para uso.