Google открыла исходный код TorchTPU — нативного бэкенда PyTorch для Cloud TPU

Google открыла исходный код TorchTPU — нативного бэкенда PyTorch для Cloud TPU. Он был совместно разработан с командой PyTorch в Meta и представлен на конференции PyTorch Conference 2026 20 октября 2026 года. Существующий код PyTorch обращается к TPU через стандартный torch.device("tpu"). Теперь TPU находятся в PyTorch на одном уровне с GPU и CPU как нативное устройство. Команда Google заявляет, что TorchTPU заменяет PyTorch/XLA и рекомендует его для всех новых рабочих нагрузок.

Кратко

• Размер: Программный стек, а не модель. 4 открытых репозитория: TorchTPU, vLLM-TorchTPU, SGLang-TorchTPU, TorchTitan. TorchTPU доступен в версии v0.2.0 на PyPI, требуется Python 3.11+, лицензия Apache 2.0.
• Работает на: Google Cloud TPU. Протестирован на поколениях начиная с v5e, масштабирован до 8 192 чипов Ironwood (TPU7x). Обслуживание поддерживает чекпойнты FP8, MXFP4 и NVFP4.
• Производительность: Eager-режим для отладки, torch.compile для пропускной способности производственного уровня.
• Лучший результат: 52% MFU при обучении Qwen 3 8B с TorchTitan на Ironwood (по собственным данным).
• Худший результат: В режиме eager «из коробки» один шаг в сценарии дообучения занимал 11,49 с против 2,78 с после полной оптимизации.

Что такое TorchTPU?

TorchTPU — это открытый бэкенд PyTorch от Google и Meta, который нативно запускает модели PyTorch на TPU. Он подключается через интерфейс PrivateUse1 PyTorch. Тензоры являются обычными тензорами PyTorch, без подклассов и обёрток. Он запускает отдельный процесс для каждого устройства и сохраняет представления тензоров и мутации «на месте».

Предшественник, PyTorch/XLA, записывал графы ленивых тензоров. Разработчикам часто приходилось добавлять код для обозначения компилируемых областей. PyTorch/XLA также поддерживал только чистый код SPMD. TorchTPU также обрабатывает расходящееся выполнение MPMD.

Как TorchTPU выполняет код PyTorch?

TorchTPU поставляется с 3 eager-режимами:

• Strict (по умолчанию): асинхронно отправляет каждый оператор, соблюдая стандартную семантику PyTorch.
• Debug: выполняет операции синхронно с проверкой границ тензоров.
• DeferAndFuse: объединяет последовательности eager-операторов, чтобы сократить трафик HBM. В статье о TorchTPU говорится, что этот режим может быть как минимум на 50% быстрее Strict.

Для максимальной скорости torch.compile(backend="tpu") захватывает граф с помощью TorchDynamo. Затем он понижает его через StableHLO до OpenXLA, который генерирует машинный код TPU. Скомпилированные исполняемые файлы кэшируются между процессами, хостами и перезапусками заданий.

Пользовательские ядра Pallas можно вызывать как стандартные функции PyTorch. Также работают предварительно собранные ядра Tokamax и тайловые ядра Helion.

Для распределённых вычислений используется группа процессов tpu_dist. Она поддерживает DDP, FSDP2 и DTensor, а также учитывающую топологию mesh-структуру для межсоединений TPU.

Насколько быстр TorchTPU?

Все приведённые ниже показатели представлены Google и партнёрами самостоятельно.

• Масштаб обучения: TorchTitan достиг показателя до 52% MFU на Qwen 3 8B. Google сообщает о почти линейном масштабировании до 8 192 TPU Ironwood.
• Этапы оптимизации: В сценарии дообучения TorchTitan время шага сократилось с 11,49 до 2,78 с. Один только DeferAndFuse дал ускорение в 1,15 раза. Полный torch.compile дал ускорение в 3,80 раза.
• Охват моделей: Google постоянно тестирует около 5 000 моделей Hugging Face в eager- и скомпилированном режимах.
• Обслуживание в масштабах парка: Zenteiq увеличила суточный объём вывода с 1,47 млрд до 5,97 млрд токенов на 1 000 чипов TPU v6e. Этот прирост обеспечила маршрутизация llm-d в GKE, а не только TorchTPU.

Какие интеграции поставляются с TorchTPU?

• Hugging Face: Transformers v5 разворачивается на TPU с помощью device_map="tpu". Diffusers поддерживает eager-режим, torch.compile и параллелизм по тензорам.
• vLLM-TorchTPU: создан с помощью Inferact. Сохраняет вышестоящие планировщики и API-серверы, используя пользовательские ядра Pallas.
• SGLang-TorchTPU: создан с помощью RadixArk. Сохраняет повторное использование префиксов RadixAttention.
• Возможности обслуживания: KV-кэши FP8, чекпойнты FP8, MXFP4 и NVFP4. Спекулятивное декодирование включает MTP и EAGLE-3.
• TorchTitan: рецепты для Llama 3 70B/405B, Qwen 3 8B/32B, Qwen 3 235B MoE и FLUX.1 12B.
• Также: llm-d 1.0, Ray, TorchRec, VeRL (предварительная версия) и LiteRT-Torch для экспорта на устройства.

Чем TorchTPU отличается от PyTorch/XLA, TorchNeuron и JAX?

Как установить TorchTPU?

Установите пакет на TPU VM, затем выберите устройство tpu:

# pip install torch_tpu
import torch
device = torch.device("tpu")
x = torch.randn((4, 4), device=device)

В README репозитория зафиксирована версия torch_tpu==0.2.0 с индексом колёс CPU-версии PyTorch. Без этого pip устанавливает CUDA-сборку torch. Полные руководства доступны на сайте torchtpu.readthedocs.io.

Основные выводы

• TorchTPU распространяется с открытым исходным кодом по лицензии Apache 2.0 и создан совместно с Meta.
• Он заменяет PyTorch/XLA, предлагая выполнение на основе PrivateUse1 с приоритетом eager-режима.
• torch.compile(backend="tpu") понижает код через StableHLO и XLA.
• TorchTitan достиг 52% MFU на Qwen 3 8B и масштабируется до 8 192 чипов Ironwood.
• vLLM, SGLang и Hugging Face теперь работают на TPU через стандартный PyTorch.

Изучите репозиторий на GitHub, статью, документацию и объявление Google Cloud. Вся заслуга принадлежит исследователю этого проекта. Также подписывайтесь на нас в Twitter и не забудьте присоединиться к нашему сабреддиту о машинном обучении с более чем 150 тыс. участников и подписаться на нашу рассылку. Стоп! Вы есть в Telegram? теперь вы также можете присоединиться к нам в Telegram.