Двенадцать лет TPU были вотчиной TensorFlow и JAX. Если вы писали на PyTorch и хотели использовать тензорные ядра Google, вам предлагали переписать модель под чужой рантайм или терпеть PyTorch/XLA с его компиляцией на минуты. В середине 2025 года в Google начали проект, который ломает это правило: TorchTPU. Это нативный бэкенд PyTorch для TPU, сделанный совместно с командой PyTorch из Meta, и к октябрю 2026 года он прошёл через 4 917 моделей Hugging Face, обучил Qwen3 MoE на 8 192 чипах и получил независимый бенчмарк SemiAnalysis с ценой 2,01 цента за миллион токенов.

Разбираю, что именно они сделали, какие цифры за этим стоят и где всё ещё болит.

Что такое TorchTPU

TorchTPU это открытый бэкенд, который регистрируется в диспетчере PyTorch как устройство tpu. Модель запускается без переписывания: меняется model.to("cuda") на model.to("tpu") и бэкенд процесса с nccl на tpu_dist. Дальше работает обычный PyTorch: eager-режим, torch.compile, DTensor, FSDP, vLLM, SGLang, TorchTitan, Ray и TorchRec. Никаких зависимостей от JAX в рантайме нет.

Именно последняя фраза отличает проект от предыдущих попыток. PyTorch/XLA появился в 2020 году и строился на Lazy Tensors: граф накапливался целиком, а компилировался потом. TorchAX шёл ещё дальше и просто заворачивал PyTorch в JAX. Оба подхода заставляли разработчика подстраиваться под чужую модель исполнения.

Почему PyTorch/XLA не взлетел

Авторы перечисляют причины прямо, и они звучат как список симптомов, знакомых каждому, кто пробовал TPU из PyTorch.

Нужно было вручную размечать в исходном коде регионы, которые компилятор способен обработать. Компиляция занимала столько времени, что итеративный цикл разработки разваливался. Императивного исполнения не было вообще. Новые примитивы распределённых вычислений из PyTorch приезжали с задержкой в месяцы, потому что их сначала должен был поддержать XLA-стек.

В сумме это давало эффект, обратный обещанному: ускоритель, который должен экономить деньги, экономил их только при условии, что вы уже переписали всё под JAX.

Три принципа, которые записали первыми

Команда сформулировала требования до начала работы, и это видно по результату.

Первый блок назвали «гражданством PyTorch»: никакого JAX в рантайме, рабочий eager-режим, torch.compile через Dynamo, а не через Lazy Tensor, полная поддержка идиом распределённых вычислений вроде DDP, FSDP и DTensor за счёт того, что TPU-коллективы реализованы на самом нижнем уровне стека, в Process Group API.

Второй блок про пользователя: портирование должно занимать одну строку, ошибки компилятора и рантайма должны быть внятными, инструменты вроде torch.profiler, DTensor и Ray должны работать как обычно.

Третий блок амбициознее: догнать PyTorch на GPU на больших нагрузках, придумать JIT-компиляцию для eager-исполнения, которая сокращает разрыв между пооперационным запуском и компиляцией, и научить стек работать с ограниченным динамизмом форм, чтобы XLA не перекомпилировал граф при каждом изменении длины последовательности.

Как это устроено внутри

Технически бэкенд зарегистрирован в диспетчере PyTorch как PrivateUse1, а хранилище тензоров реализовано через стандартный механизм at::Tensor с c10::UniqueVoidPtr. Ничего экзотического: это ровно тот путь, которым добавляют новые устройства в PyTorch.

Интереснее система режимов исполнения. В strict eager каждая Aten-операция уходит на TPU по одной, асинхронно, как на GPU. В debug-режиме (TORCH_TPU_LAUNCH_BLOCKING=1) исполнение синхронное, с проверкой границ, и это медленно, но удобно для поиска ошибок.

Третий режим, DeferAndFuse (TORCH_TPU_DEFER_AND_FUSE=1), это гибрид: операции откладываются, группируются и компилируются в более крупные куски. Авторы обещают от 50% ускорения относительно strict eager и заодно снижение пикового потребления HBM, потому что склейка операций уменьшает число временных буферов.

Компилируемый режим сделан по-честному: регистрируется бэкенд для TorchDynamo, который получает FX-граф и примеры входных тензоров, трассирует граф целиком без исполнения на TPU и выдаёт один StableHLO-модуль. Дальше torch.compile(model, backend="tpu") работает так же, как на CUDA.

Отдельная боль TPU это статические формы. XLA исторически требует знать размерности на этапе компиляции, поэтому при смене batch size или длины контекста граф перекомпилируется. TorchTPU добавляет bounded dynamism: буферы выделяются по верхней границе, а реальные размеры становятся известны в рантайме. На 839 моделях с SymInt корректная численность подтверждена у 834, то есть 99%. Пока это не включено по умолчанию, поэтому при вызове torch.compile на TPU флаг dynamic принудительно ставится в False.

С точностью тоже пришлось договариваться. По умолчанию на CUDA PyTorch работает в режиме highest, а на TPU бэкенд ставит medium: умножение идёт в BF16 на MXU, накопление в FP32. Режимы high и highest реализованы алгоритмически, через 3 и 6 проходов по матричным блокам, и стоят примерно 3x и 6x вычислений соответственно. При их включении TorchTPU печатает предупреждение о производительности, а для точечной работы есть контекстный менеджер torch.tpu.precision.

4 917 моделей и 98% прохождений

Самая убедительная часть отчёта это не архитектура, а тестовый корпус. Команда прогнала около 5 000 моделей Hugging Face на поколениях TPU начиная с v5e, в eager и в compiled режимах, и построила двухуровневую выборку. Первый уровень это популярные модели, которые скачивают чаще всего: Qwen, Gemma, Mistral, DeepSeek, GLM, Kimi, Llama, Whisper, Stable Diffusion, FLUX, ViT. Второй уровень это по одному представителю более 500 архитектурных семейств, на которых проверяется численная эквивалентность с GPU-базой.

Результат по 4 917 моделям: 98% проходят smoke-тест, то есть запускаются и не падают (4 821 модель). Численный паритет на архитектурной выборке 97%, это 638 из 657 моделей. Критерии жёсткие: для языковых моделей относительное расхождение perplexity с GPU-базой не должно превышать 0,10, для vision, audio и энкодеров используется относительный MSE с двойным порогом.

Отдельно проверена работа custom-ядер: их можно писать на Pallas через JAX и на Helion DSL. Поддержано около 450 операторов Aten, включая forward и backward для всех базовых операций, нужных основным архитектурам.

Сколько это даёт в цифрах

Пример из репозитория это самодостаточная реализация Muse Glimmer 30B на чистом PyTorch, без единого модельного ядра. На одном чипе TPU v7 в режиме prefill она выдаёт около 25,7 тысячи токенов в секунду при 59% MFU для последовательности 5 120, и около 22,4 тысячи при 51% MFU для 10 240. Обучение на четырёх чипах с FSDP даёт 5,6 тысячи токенов в секунду на чип при 38% MFU.

Любопытная деталь: FSDP здесь запускается с world size 2, потому что один физический чип TPU v7 отдаёт PyTorch два устройства, по одному на каждый TensorCore.

Более честный тест это порт той же модели из TorchTitan, изначально написанной под CUDA. Инженеры заменили только FlexAttention на обычный F.scaled_dot_product_attention и включили torch.compile с SimpleFSDP. Получилось 5,3 тысячи токенов в секунду на чип при 36% MFU и 4,2 тысячи при 29%. То есть код, написанный под GPU, поехал на TPU почти без изменений и без вылизывания под конкретное железо.

Масштабирование проверяли на Qwen3 MoE 235B через TorchTitan. Первый smoke-тест на 64 чипах дал 1 828 токенов в секунду против 1 937 у JAX на MaxText, разрыв около 6%. Дальше шли конфигурации от одного слайса до восьми.

Один слайс на 1 024 чипах с FSDP выдаёт 2 236 токенов в секунду на TensorCore при 28,8% MFU и шаге 29,3 секунды. Восемь слайсов на 8 192 чипах держат 2 104 токена при 27,1% MFU и шаге 31,2 секунды. MFU почти не падает, то есть слабое масштабирование близко к линейному. Причина в том, что XLA полностью перекрывает all-reduce по DCN для слоя t+1 с вычислением слоя t, и в multi-slice HSDP forward вообще не требует межслайсовой коммуникации.

Серьёзные числа есть и по Llama 3 405B: 1 024 чипа Ironwood, контекст 8K, BF16, 500 токенов в секунду на чип и больше 50% MFU. У Qwen3 235B MoE на 8 192 чипах v7x получается 28% MFU при длине последовательности 4K.

Эффект оптимизаций удобно смотреть на FLUX.1-dev 12B. Стартовый eager-бейзлайн даёт шаг 11,49 секунды. Включение DeferAndFuse ускоряет до 9,95 (в 1,15 раза), SimpleFSDP доводит до 6,52 (1,76x), кастомное Pallas-ядро SplashAttention до 5,55 (2,07x), полный torch.compile до 3,02 секунды (3,80x), а peephole-переписывания на целевой модели до 2,78 секунды, то есть 1,08 сэмпла в секунду. Итоговое ускорение в 4,1 раза без смены модели.

Инференс: 2,01 цента за миллион токенов

Самое интересное для тех, кто считает деньги, это независимая проверка SemiAnalysis в рамках InferenceX. TPU v7 вместе со стеком vLLM-TorchTPU показал лучшую стоимость за миллион токенов среди участников: 2,01 цента на профиле 8K/1K, тогда как в сентябрьском дебюте было 3,6 цента. На агентском профиле AgentX, где есть многоходовость, длинный контекст, высокий переиспользование префикса и всплески субагентов, четыре чипа TPU v7 дали 0,83 цента за миллион токенов и лидерство по интерактивности.

Чтобы это получить, пришлось лечить не софт, а стык софта с железом, и три примера показывают, насколько глубоко.

Первый это expert-parallel слой MoE. В vLLM он собирал токены со всех рангов, маршрутизировал, переставлял, запускал два grouped matmul и делал reduce-scatter. На prefill с конкурентностью 256 этот reduce-scatter стоил 1 190 микросекунд на слой, 22% всего времени устройства, и параллельно с ним ничего не выполнялось. В vLLM-TorchTPU весь слой стал одной SPMD-программой на восемь TensorCore: каждый ядро квантует свои строки в FP8 и асинхронно собирает их на SparseCore, выбирает экспертов Pallas-ядром, а результаты отправляет владельцу токена через удалённый DMA, перекрывая отправку со следующими матричными умножениями.

Второй это раскладка KV-кэша. Векторные блоки TPU работают 32-битными словами, и стандартная раскладка KV упаковывала в слово K и V от разных голов. При тензорном параллелизме на 8 ядер каждому достаётся одна KV-голова из двух, и половина слова оказывалась паддингом: 1 024 байта кэша на токен при 512 байтах данных. Раскладка sequence-on-lane кладёт токены страницы по 128 линиям, а размерность головы по саблайнам, и слово теперь содержит четыре значения одного токена. Память под KV падает вдвое, как и объём чтения на каждом шаге декодирования.

Третий это Gated DeltaNet, который занимает 45 слоёв из 60. Он держит фиксированное состояние 128 на 128 на голову вместо растущего кэша, поэтому на конкурентности 64 удаление всей арифметики из ядра ускорило его лишь на 6%: узкое место в памяти, а не в вычислениях. Одно Pallas-ядро склеивает свёртку и рекуррентность, батчит декодирование и обрабатывает prefill блоками.

Что это меняет на рынке

До сих пор выбор ускорителя означал выбор экосистемы. CUDA означала PyTorch и весь его зоопарк, TPU означала JAX и переписывание, Trainium означала собственные обвязки. TorchTPU убирает этот пункт из уравнения: железо становится заменяемым, а стек остаётся одним и тем же.

Для Google это способ продавать TPU тем, кто никогда не согласится уйти с PyTorch, то есть почти всем. Для Meta это усиление позиции PyTorch как нейтрального слоя, который не принадлежит производителю чипов. Для команд, которые упёрлись в стоимость инференса на GPU, это второй поставщик с подтверждённой ценой за токен.

Есть и неприятный для Nvidia слой. Модель, которая запускается на любом ускорителе одной строкой, обесценивает не только CUDA-lock-in, но и целые команды, чья работа состояла в портировании.

Где пока больно

Честности отчёту добавляет раздел с ограничениями. FlexAttention не поддерживается, и в примерах его заменяют обычным scaled dot product attention. Один процесс на устройство вместо модели JAX, где один процесс на хост: это решение принято ради интеграции с PyTorch, но оно означает, что RL-вычисления с совместным размещением политики и среды из коробки не работают, и команда ищет обходные пути.

Динамические формы включены не полностью, torch.compile на TPU пока принудительно статичен. Режимы high и highest стоят 3x и 6x по матричным блокам, что для чувствительных к точности моделей означает серьёзный удар по производительности.

В дорожной карте на 6–12 месяцев перечислены FP8 и FP4 для TPU v8, производственная готовность bounded dynamism, упаковка под GKE, обучение с RL через сэмплирование в vLLM и SGLang под управлением Ray, поддержка диффузионных и MoE-архитектур в день релиза, а также Trace Viewer 2.0 на WebGPU и постоянное профилирование XProfz.

Часто задаваемые вопросы

Нужно ли переписывать модель, чтобы запустить её на TPU?

Нет, если модель написана на стандартном PyTorch. Достаточно поменять устройство и бэкенд процесса: model.to("tpu") вместо model.to("cuda") и backend="tpu_dist" вместо nccl. Модели на чистом PyTorch запустились на 98% из 4 917 проверенных. Исключения это модели с FlexAttention и редкими операторами вне списка 450 поддержанных Aten-опов.

TorchTPU заменяет JAX?

Нет. JAX остаётся основным способом выжать из TPU максимум, и авторы сами сравнивают свои результаты именно с JAX-бейзлайнами. TorchTPU снимает требование переходить на JAX ради доступа к железу. Для production-нагрузок рекомендуются готовые стеки: vLLM-TorchTPU, SGLang-TorchTPU и TorchTitan-TorchTPU с Pallas-ядрами.

Насколько это быстрее, чем PyTorch/XLA?

Отчёт не даёт прямого сравнения с PyTorch/XLA, но разница в подходе видна по режимам. Появился рабочий eager, гибридный DeferAndFuse даёт от 50% ускорения относительно strict eager, torch.compile приносит до 3,8x на модели FLUX. Главный выигрыш не в процентах, а в том, что цикл разработки перестал упираться в минуты компиляции.

Итог

TorchTPU это не ещё один мост между PyTorch и TPU, а попытка сделать TPU обычным устройством PyTorch с точки зрения пользователя. Цифры за первый открытый релиз серьёзные: 4 917 моделей, 98% запусков, 2 104 токена в секунду на TensorCore при масштабировании на 8 192 чипа и 2,01 цента за миллион токенов в независимом бенчмарке.

Ограничения тоже названы прямо: нет FlexAttention, нет полного динамизма форм, RL с совместным размещением пока не работает. Если вы держите обучение или инференс на PyTorch и считаете стоимость токена, самое время посмотреть на TPU как на реальную опцию, а не как на чужую планету. Четыре репозитория открыты, начните с прогона своей модели на одном чипе и сравните с текущим счётом за GPU.