nanochat Карпати перенесли на TPU: качество сошлось, скорость просела вдвое

Разработчик под ником tucan9389 перенёс проект nanochat Андрея Карпати, открытый учебный проект «полного цикла» обучения языковой модели (токенизатор, претрейн, инструктивное дообучение SFT и RL в одном репозитории, известный как «спидран за $100» на связке из 8 видеокарт H100), с PyTorch на фреймворк JAX, чтобы обучать модель на чипах Google TPU. Результат порта выложен отдельным репозиторием nanochat-jax на GitHub. Цель, максимально сохранить архитектуру и конфигурацию оригинала (полное соответствие, «parity») и повторить показатели Карпати и по качеству модели, и по скорости обучения.

По качеству цель достигнута: метрика CORE (усреднённая точность по 22 оценочным задачам) для базовой модели глубиной 24 слоя (~1,4 млрд параметров, рецепт R4) на TPU составила 0,2695, это чуть выше диапазона 0,2512, 0,2677, который сам Карпати получил, семь раз прогнав тот же рецепт на GPU, и выше опорного значения GPT-2 (0,2565). После дообучения (SFT) модель прошла и качественную проверку на конкретных промптах: например, на запрос «как сварить кофе» базовая модель до SFT просто повторяла вопрос, а после SFT давала связный ответ; на просьбу произнести слово по буквам базовая модель ошибалась, SFT-версия справлялась верно.

По скорости разрыв остался. Показатель MFU (доля от пиковой производительности чипа, которую реально использует обучение) для итогового прогона на TPU v6e-8 вышел около 24%, вдвое ниже 47, 48%, которые Карпати измерил на 8×H100. Полный цикл (токенизатор → база → SFT) занял 12,19 часа и обошёлся в $60,8 по spot-цене TPU (около $263 по обычному тарифу), с одной внеплановой остановкой инстанса и 17-минутным восстановлением из чекпоинта.

Чтобы дойти до этого результата, автор потратил три месяца и около $2500, в основном на аренду TPU. Больше всего денег (~$730) ушло на воспроизведение качества на более старом чипе TPU v5p, для проверки каждой гипотезы (баг накопления градиента, точность вычислений и т. д.) требовался полный прогон обучения. Ещё ~$275 и десять дней ушли на поиск регрессии качества после переезда на новый чип v6e, автор сначала подозревал сам чип, потом пониженную точность bf16, а причиной оказался баг в загрузке данных. Пять внештатных инцидентов (потерянные чекпоинты, переобучение из-за нехватки данных, оплата простаивающей машины, падение из-за нулевого чекпоинта, сбой многохостового запуска) стоили ~$585. Оптимизация скорости обучения, ~$375, репетиции перед публикацией (нашли ещё три проблемы, всплывающие только на полном прогоне), ~$165, финальные чистые прогоны, ~$180. Сам по себе перенос кода, по оценке автора, стоил всего $40, почти все траты ушли на воспроизведение результата и отладку.

Автор выделяет пять конкретных находок порта. Первая и самая коварная, тихая порча данных: код собирал мини-батчи для накопления градиента через np.asarray, а эта функция не копирует данные, если на входе уже массив NumPy, в результате все собранные батчи внутри шага оказывались копиями последнего батча из переиспользуемого буфера данных. Ошибка не выдавала ни NaN, ни несовпадения размеров, только просевшее качество: прогоны с этим багом застревали в диапазоне CORE 0,08, 0,13, а после исправления (замена на np.array(..., copy=True), копирование данных сразу при получении) CORE сразу дал 0,274. Вторая находка касается компиляции: сборка всего шага обучения в одну скомпилированную функцию jax.jit (вместо трёх раздельно скомпилированных функций, соединённых Python-циклом) подняла MFU с 16% до 22%, раздельные функции оставляли чип простаивать между вызовами до 26,7% времени шага. Третья, ядро Splash-attention (реализация flash-attention для TPU) по умолчанию использует размеры тайлов 128 (с незакрытым комментарием-напоминанием у самих авторов JAX); подбор тайлов под длину последовательности 2048 (512/512/256) поднял MFU в изолированном эксперименте с 22,4% до 41,2%, почти не изменив результат обучения. Четвёртая, по умолчанию TPU считает даже операции в формате fp32 с пониженной точностью bf16: флаг повышенной точности, который автор считал спасающим качество модели, оказался бесполезным, поскольку веса к тому моменту уже приводились к bf16 и операций в fp32 просто не оставалось. Пятая, оптимальный размер «головы» внимания (head_dim) зависит от поколения TPU: у v6e матричный блок вычислений (MXU) вырос до 256×256 (у более старых чипов, 128×128), и конфигурация с 6 головами по 256 вместо 12 голов по 128, как в оригинале Карпати, дала прирост скорости обучения около 17% при том же числе параметров.

Отдельная головная боль, не техническая, а инфраструктурная: получить сами TPU в аренду оказалось труднее, чем решить перечисленные проблемы. Из 2885 попыток создать TPU-инстанс с апреля по июль 2026 года 1632 закончились ошибкой; даже в самом успешном регионе (us-central1) доля успеха составила 36% (208 из 575 попыток), в большинстве других регионов, 0% из-за нулевых квот или отсутствия свободных чипов.

Код порта и обученные веса (базовая и SFT-версии модели) автор выложил в открытый доступ, репозиторий nanochat-jax на GitHub и чекпоинты на Hugging Face. Что именно ограничивает оставшуюся разницу в скорости, простои между компилируемыми блоками, размеры тайлов внимания, выравнивание head_dim или пропускная способность памяти, автор называет предметом следующего этапа работы.

Ключевые факты

Качество модели (CORE 0,2695) воспроизвелось на уровне Карпати, чуть выше его диапазона, скорость обучения (MFU ~24%) вдвое ниже, чем у него на H100 (47, 48%)
Полный цикл обучения на TPU v6e-8 занял 12,19 часа и обошёлся в $60,8 по spot-цене; весь эксперимент с отладкой, три месяца и около $2500
Ключевой скрытый баг, NumPy-функция np.asarray не копирует данные при накоплении градиента, из-за чего CORE проседал до 0,08, 0,13; исправление, одна строка кода
Подбор размеров тайлов ядра Splash-attention почти удвоил MFU (с 22,4% до 41,2%); сборка шага обучения в один jax.jit подняла MFU с 16% до 22%
Получить сами TPU оказалось сложнее технической части: из 2885 попыток аренды успешными были лишь 36% даже в лучшем регионе, в большинстве остальных, 0%

Почему это важно

Перенос показывает: путь с PyTorch и CUDA на JAX и TPU, не просто смена библиотеки, а источник скрытых нестандартных ловушек, которые не описаны в документации фреймворков. Автор задокументировал их с конкретными числами (MFU до/после, код фикса, стоимость в долларах), редкий уровень детализации для инженерного отчёта такого рода, восполняющий пробел: готовых пошаговых разборов переноса конкретного открытого проекта обучения LLM на TPU почти нет.

Кому это важно

ML- и инфраструктурным инженерам, которые обучают модели на TPU или рассматривают переход с GPU на TPU; разработчикам на JAX/Flax, которые пишут собственные шаги обучения и ядра Pallas; тем, кто считает бюджет обучения и хочет заранее понимать реальную стоимость эксперимента с TPU, включая квоты и простои.

Как это применить

Готовые к переиспользованию выводы: собирать весь шаг обучения в один jax.jit вместо цепочки раздельно скомпилированных функций; при накоплении градиента вручную копировать данные сразу при чтении, если источник переиспользует буфер (np.array(..., copy=True) вместо np.asarray); для Splash-attention на TPU задавать тайлы вручную (в примере, 512/512/256 для длины последовательности 2048), не полагаясь на дефолт 128; не рассчитывать, что флаг «высокая точность» в JAX что-то даёт при обучении в bf16; при выборе head_dim ориентироваться на размер матричного блока конкретного поколения TPU (256×256 у v6e, значит head_dim кратен 256). Код и веса открыты (GitHub nanochat-jax, Hugging Face).

Можно ли доверять

Источник, личный технический отчёт автора порта в обсуждениях его же репозитория на GitHub, без независимой проверки. В пользу достоверности, высокая детализация (полные логи обучения выложены отдельно, приведены точные команды и код, числа сверяются с официальным лидербордом nanochat Карпати) и то, что похожие по масштабу цифры MFU (15, 52% в зависимости от конфигурации) независимо получали несколько других авторов, портировавших похожие проекты на TPU. Общая стоимость ~$2500 и разбивка по статьям, самооценка автора, без чеков или внешнего аудита.

Риски и подводные камни

Главный урок автора: без острой необходимости такой перенос лучше не затевать, три месяца и порядка $2500 ушли не на сам перенос кода (~$40), а на отладку и повторные прогоны. Отдельные системные ловушки: NumPy-функция np.asarray молча не копирует данные, если на входе уже массив, источник целого класса скрытых багов в пайплайнах данных (автор ссылается на исследование, где похожий баг найден более чем в 95% случайно проверенных репозиториев с похожим паттерном загрузки данных); TPU по умолчанию занижает точность fp32-вычислений до bf16 без явной ошибки; получить сами TPU в нужном регионе по квоте может быть сложнее, чем решить любую из технических проблем.

«Самое главное, чему меня научил этот перенос: если можно обойтись без него, не делайте этого.»

— tucan9389, автор порта nanochat-jax