Nanocode : Entraînement d'agents de codage de type Claude avec JAX sur des TPU

Nanocode est une bibliothèque qui montre comment entraîner votre propre modèle Claude Code de bout en bout en utilisant l'IA Constitutionnelle, suivant l'approche d'Anthropic. Écrite entièrement en JAX et optimisée pour les TPU, elle adapte l'infrastructure du projet nanochat de Karpathy.

Configuration et Coûts d'Entraînement

Le modèle nanocode-d24 (1,3 milliard de paramètres) peut être reproduit en environ 9 heures sur un TPU v6e-8 pour un coût de 200 $. Le plus petit modèle nanocode-d20 (477 millions de paramètres) s'entraîne en environ 1,5 heure pour 34 $. Le projet recommande d'utiliser le programme TRC de Google pour un accès gratuit aux TPU préemptibles pendant un mois, ou les crédits de 300 $ de Google Cloud pour les nouveaux comptes.

Implémentation Technique

Le processus d'entraînement comprend :

• Écrire un fichier SOUL.md pour définir l'alignement du modèle
• Définir une interface agentique pour l'interaction avec le monde
• Générer des données synthétiques
• Utiliser l'optimisation des préférences pour aligner le modèle avec SOUL

Différences de Tokenisation et de Pré-entraînement

Bien que le processus de pré-entraînement et d'entraînement du tokeniseur soit similaire à nanochat, nanocode inclut des données de codage supplémentaires de The Stack-V2 dans un ratio de 1:5 dans les mélanges de pré-entraînement et de tokeniseur. Cela donne de meilleures performances en codage mais réduit l'efficacité de la tokenisation du texte général.

La comparaison des tokeniseurs montre que nanocode obtient une tokenisation -50,9 % meilleure pour le code par rapport à nanochat, tandis que nanochat est meilleur sur le texte coréen (+7,9 % pour nanocode sur les actualités, -27,6 % sur le coréen).

Commandes et Configuration

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

Les modèles sont entraînés avec un ratio paramètres:données de 8, suivant l'analyse de la loi d'échelle de nanochat. Bien qu'optimisé pour les TPU, nanocode devrait également fonctionner sur les GPU NVIDIA sans modification.