Nanocode: Training Claude-like coding agents with JAX on TPUs

Nanocode ist eine Bibliothek, die zeigt, wie man sein eigenes Claude Code-Modell end-to-end mit Constitutional AI nach Anthropics Ansatz trainiert. Vollständig in JAX geschrieben und für TPUs optimiert, adaptiert es die Infrastruktur von Karpathys nanochat-Projekt.

Trainingssetup und Kosten

Das nanocode-d24-Modell (1,3 Milliarden Parameter) kann in etwa 9 Stunden auf einem TPU v6e-8 zu Kosten von 200 US-Dollar reproduziert werden. Das kleinere nanocode-d20-Modell (477 Millionen Parameter) trainiert in etwa 1,5 Stunden für 34 US-Dollar. Das Projekt empfiehlt die Nutzung von Googles TRC-Programm für einen Monat kostenlosen Zugang zu pre-emptible TPUs oder die 300-US-Dollar-Gutschriften von Google Cloud für neue Konten.

Technische Umsetzung

Der Trainingsprozess umfasst:

• Erstellen einer SOUL.md-Datei zur Definition der Modellausrichtung
• Definition einer agentenbasierten Schnittstelle für Weltinteraktion
• Generierung synthetischer Daten
• Verwendung von Präferenzoptimierung zur Ausrichtung des Modells an SOUL

Tokenisierung und Pretraining-Unterschiede

Während der Pretraining- und Tokenizer-Trainingsprozess nanochat ähnelt, enthält nanocode zusätzliche Coding-Daten von The Stack-V2 im Verhältnis 1:5 sowohl in Pretraining- als auch Tokenizer-Mischungen. Dies führt zu besserer Coding-Leistung, reduziert aber die Effizienz der allgemeinen Text-Tokenisierung.

Ein Tokenizer-Vergleich zeigt, dass nanocode eine um -50,9 % bessere Tokenisierung für Code erreicht als nanochat, während nanochat bei koreanischem Text besser abschneidet (+7,9 % für nanocode bei Nachrichten, -27,6 % bei Koreanisch).

Befehle und Setup

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

Die Modelle werden mit einem Parameter:Daten-Verhältnis von 8 trainiert, entsprechend nanochats Skalierungsgesetz-Analyse. Obwohl für TPUs optimiert, sollte nanocode auch auf NVIDIA-GPUs ohne Anpassungen funktionieren.