나노코드: JAX와 TPU를 사용한 클로드(Claude) 스타일 코딩 에이전트 훈련

Nanocode는 Anthropic의 접근 방식을 따라 Constitutional AI를 사용하여 자신만의 Claude Code 모델을 종단 간 학습시키는 방법을 보여주는 라이브러리입니다. 완전히 JAX로 작성되었으며 TPU에 최적화되어 있으며, Karpathy의 nanochat 프로젝트의 인프라를 적용합니다.

학습 설정 및 비용

nanocode-d24 모델(13억 개의 파라미터)은 TPU v6e-8에서 약 9시간 동안 200달러의 비용으로 재현할 수 있습니다. 더 작은 nanocode-d20 모델(4억 7700만 개의 파라미터)은 약 1.5시간 동안 34달러의 비용으로 학습됩니다. 이 프로젝트는 한 달 동안 선점 가능한 TPU에 무료로 접근하기 위해 Google의 TRC 프로그램을 사용하거나, 새 계정에 대해 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에서도 즉시 작동해야 합니다.