Нанокод: обучение кодирующих агентов, подобных Claude, с использованием JAX на TPU

Nanocode — это библиотека, которая демонстрирует, как обучить собственную модель Claude Code сквозным образом с использованием Constitutional AI, следуя подходу Anthropic. Написана полностью на JAX и оптимизирована для TPU, адаптирует инфраструктуру из проекта nanochat от Карпати.

Настройка обучения и затраты

Модель nanocode-d24 (1,3 млрд параметров) можно воспроизвести примерно за 9 часов на TPU v6e-8 стоимостью $200. Меньшая модель nanocode-d20 (477 млн параметров) обучается около 1,5 часов за $34. Проект рекомендует использовать программу Google TRC для бесплатного доступа к прерываемым TPU на месяц или кредиты Google Cloud в размере $300 для новых аккаунтов.

Техническая реализация

Процесс обучения включает:

• Написание файла SOUL.md для определения согласованности модели
• Определение агентского интерфейса для взаимодействия с миром
• Генерацию синтетических данных
• Использование оптимизации предпочтений для согласования модели с SOUL

Различия в токенизации и предварительном обучении

Хотя процесс предварительного обучения и обучения токенизатора похож на nanochat, nanocode включает дополнительные данные по программированию из The Stack-V2 в соотношении 1:5 как в предварительном обучении, так и в смесях токенизатора. Это приводит к более высокой производительности в кодировании, но снижает эффективность токенизации общего текста.

Сравнение токенизаторов показывает, что nanocode достигает на -50,9% лучшей токенизации для кода по сравнению с nanochat, в то время как nanochat работает лучше с корейским текстом (+7,9% для nanocode на новостях, -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

Модели обучаются с соотношением параметров:данных 8, следуя анализу закона масштабирования nanochat. Хотя оптимизированы для TPU, nanocode также должен работать на GPU NVIDIA из коробки.