Google открыла код TorchTPU для запуска PyTorch на TPU

Подробно

Что такое TorchTPU и что именно выпустила Google?

TorchTPU - это нативный бэкенд PyTorch для тензорных процессоров Google. Его совместно разработали Google и команда PyTorch в Meta. Об открытии кода 20 октября 2026 года объявили Марк Ломейер, вице-президент и генеральный менеджер AI & Computing Infrastructure в Google Cloud, и Билл Цзя, вице-президент по разработке.

Репозиторий google-pytorch/torch_tpu опубликован на GitHub под лицензией Apache 2.0. На момент проверки README указывал публичный пакет torch_tpu версии 0.2.0 и поддержку Python 3.11 и новее. Среди заявленных возможностей:

• ATen-ядра для TPU;
• кеш компиляции XLA;
• распределённое обучение с параллелизмом данных и тензоров;
• интеграция с torch.compile;
• примеры для Llama, Qwen, ResNet, DLRM, minGPT и UNet.

Отдельно Google открыла TPU agent suite в статусе раннего предварительного доступа (early preview). Это набор подключаемых инструментов, которые автоматизируют перенос моделей, диагностику, ответы на вопросы и настройку производительности инференса.

Как запустить PyTorch на Google TPU с помощью TorchTPU?

TorchTPU позволяет запускать операции и модели PyTorch на TPU, отлаживать их в eager mode, компилировать через torch.compile и масштабировать через torch.distributed с бэкендом tpu_dist. В примере из документации TPU выбирается вызовом torch.device("tpu"). В Hugging Face Transformers v5 для этого служит параметр device_map="tpu".

По формулировке Google, существующие нагрузки PyTorch запускаются с минимальными изменениями кода и без структурного переписывания. Это не значит, что любое приложение, завязанное на CUDA, переносится одной настройкой. Как минимум нужно указать устройство, а собственные CUDA-ядра и сторонние библиотеки могут потребовать адаптации. Насколько полно TorchTPU поддерживает произвольные операторы, CUDA extensions и Triton-ядра вне перечисленных Google интеграций, пока не проверено.

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

Раньше PyTorch на TPU работал через PyTorch/XLA. Он использовал ленивую запись вычислительного графа, и разработчикам порой приходилось вручную размечать компилируемые участки. TorchTPU по умолчанию выполняет операции сразу, в eager mode, по модели «один процесс на устройство».

В режиме графа TorchTPU захватывает граф через TorchDynamo и компилирует его в машинный код TPU через OpenXLA и StableHLO. Скомпилированные графы кешируются между процессами, хостами и перезапусками заданий. Google рекомендует TorchTPU для новых проектов и заявляет, что он заменяет PyTorch/XLA.

Какие модели и библиотеки поддерживает TorchTPU?

Google заявила интеграцию с Hugging Face Transformers и Diffusers, vLLM, SGLang, llm-d, TorchTitan, VeRL, Ray, TorchRec и LiteRT. Для TorchTitan с TorchTPU есть готовые рецепты для Llama 3 70B и 405B, Qwen 3 8B, 32B и 235B MoE, MuseGlimmer 30B и FLUX.1 12B.

Партнёры Google положительно оценили релиз в официальном анонсе.

«TorchTPU позволяет нам переносить нагрузки PyTorch на Google TPU, напрямую давать обратную связь инженерам и создавать необходимую экосистеме гибкость», - Яна ван Грёнен, директор по разработке PyTorch в Meta.

«Поддержка TorchTPU в Transformers помогает сообществу запускать модели из Hugging Face Hub на Google Cloud TPU привычными средствами PyTorch», - Джефф Будье, вице-президент по продукту Hugging Face.

«TorchTPU встроил Google TPU непосредственно в стандартный стек PyTorch, моделей, Ray и Kubernetes», - Эдвард Оукс, главный архитектор Anyscale.

Ранняя версия не лишена недочётов. В публичном отчёте об ошибках на GitHub пользователь указал на лишний диагностический вывод и предупреждения. Он также сообщил, что нет ожидаемых функций torch.tpu.memory_allocated() и torch.tpu.max_memory_allocated(). Сопровождающий проекта отложил эти замечания в бэклог, сославшись на ограниченные ресурсы и текущие приоритеты.

Насколько быстр TorchTPU и есть ли независимые тесты?

По собственным испытаниям Google, обучение Qwen 3 8B через TorchTitan достигало до 52% Model FLOPs Utilization (MFU). Масштабирование при этом было близким к линейному вплоть до 8192 чипов TPU Ironwood. TorchTitan с TorchTPU, по данным компании, проверяли на конфигурациях до 8000 TPU Ironwood.

Все эти показатели получены самой Google. Независимых тестов и сравнения с GPU Nvidia на CUDA в анонсе нет, и к моменту публикации они не найдены.

Чем TorchTPU похож на TorchNeuron от Amazon и при чём тут Nvidia CUDA?

AWS представила TorchNeuron - нативный бэкенд PyTorch для своих чипов Trainium. TechCrunch датирует его выпуск 19 октября. В материале Amazon при этом говорилось, что код откроют на PyTorch Conference 20 октября. Поэтому формулировка «Google выпустила TorchTPU на следующий день после Amazon» опирается на данные TechCrunch, а точную последовательность релизов ещё предстоит уточнить.

TechCrunch описывает оба релиза как попытку уменьшить зависимость разработчиков от CUDA. PyTorch создали в Facebook (ныне Meta) в 2016 году, сейчас проектом управляет PyTorch Foundation в составе Linux Foundation. По оценке, которую издание приводит со ссылкой на компании, PyTorch используется примерно в 85% ИИ-приложений. Первичное исследование и методика этой оценки не раскрыты.

Редактор TechCrunch по венчурной тематике Джули Борт отмечает, что открытие TorchTPU само по себе не означает массового ухода разработчиков от Nvidia. По её мнению, Google и Amazon прежде всего пытаются ослабить преимущество Nvidia, которое даёт связка CUDA и тесной интеграции с PyTorch. Официальной реакции Nvidia к моменту публикации найти не удалось.

Сколько стоит TorchTPU и доступен ли он в России?

За код TorchTPU платить не нужно, но использование Cloud TPU оплачивается отдельно и требует квоты. По опубликованному прайс-листу чип Trillium (v6e) по требованию стоит от 2,70 доллара США за чип-час в регионах us-east1 и us-east5. В europe-west4 цена составляет 2,97 доллара, в asia-northeast1 - 3,24 доллара. Итоговая цена зависит от региона и модели потребления.

Cloud TPU v6e доступны в зонах Северной Америки, Европы, Азиатско-Тихоокеанского региона и Южной Америки. Российской зоны в официальном перечне нет, а наличие нужного количества ускорителей не гарантируется. С 10 марта 2022 года Google Cloud не принимает новые регистрации пользователей из России, хотя документация по-прежнему указывает рубль как валюту самостоятельного биллинга для России. Единого актуального статуса доступа и оплаты для уже существующих российских аккаунтов установить не удалось.