ナノコード：TPU上のJAXでClaude風コーディングエージェントを訓練

Nanocodeは、Anthropicのアプローチに従い、憲法AIを使用して独自のClaude Codeモデルをエンドツーエンドでトレーニングする方法を示すライブラリです。完全にJAXで記述され、TPU向けに最適化されており、Karpathyのnanoprojectプロジェクトのインフラストラクチャを適応させています。

トレーニング設定とコスト

nanocode-d24モデル（1.3Bパラメータ）は、TPU v6e-8上で約9時間、200ドルのコストで再現できます。より小さいnanocode-d20モデル（477Mパラメータ）は、約1.5時間、34ドルでトレーニングされます。このプロジェクトでは、GoogleのTRCプログラムを利用して1か月間無料でプリエンプティブルTPUにアクセスするか、またはGoogle Cloudの新規アカウント向け300ドルのクレジットを使用することを推奨しています。

技術的な実装

トレーニングプロセスには以下が含まれます：

• モデルのアライメントを定義するSOUL.mdファイルの作成
• 世界との相互作用のためのエージェントインターフェースの定義
• 合成データの生成
• SOULとモデルを一致させるための選好最適化の使用

トークン化と事前トレーニングの違い

事前トレーニングとトークナイザーのトレーニングプロセスはnanochatと類似していますが、nanocodeにはThe Stack-V2からの追加のコーディングデータが、事前トレーニングとトークナイザーの混合物の両方で1:5の比率で含まれています。これにより、コーディング性能は向上しますが、一般的なテキストのトークン化効率は低下します。

トークナイザーの比較では、nanocodeはコードに対してnanochatよりも-50.9%優れたトークン化を達成していますが、nanochatは韓国語テキストで優れた性能を示しています（ニュースではnanocodeが+7.9%、韓国語では-27.6%）。

コマンドと設定

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

モデルは、nanochatのスケーリング法則分析に従い、パラメータ:データ比率8でトレーニングされます。TPU向けに最適化されていますが、nanocodeはNVIDIA GPUでもそのまま動作するはずです。