AIREITERTorchTPU позволяет относительно безболезненно перенести PyTorch-код с cuda на tpu. Но одного изменения устройства недостаточно: всё упирается в доступ к TPU, правильный пакет, компиляцию и формы тензоров.
Что TorchTPU действительно предлагает
TorchTPU — это нативный бэкенд PyTorch для Google TPU, а не облачный API моделей и не замена TPU VM. В анонсе PyTorch заявлены eager-исполнение, torch.compile, распределённые примитивы и поддержка примерно 450 операций ATen. В публичном репозитории также перечислены ядра ATen, кэширование компиляции, распределённые примеры и интеграция с torch.compile.
Сам перенос выглядит привычно:
# Before
model.to("cuda")
# First TPU port
model.to("tpu")Но это лишь проверка переносимости, а не обещание высокой производительности. Как отмечает PyTorch, пиковая скорость на TPU по-прежнему зависит от форм моделей, шардирования и того, насколько используемые операции подходят этому железу.
Что подготовить до установки
TorchTPU имеет смысл только там, где выделен Google TPU. Репозиторий рекомендует запускать ноутбуки на TPU VM и называет экземпляр Google Compute Engine v6e самым простым вариантом для старта. На ноутбуке без TPU пакеты установить можно, но проверить выполнение на TPU — нет.
В публичном README заявлена поддержка Python 3.11 и новее, а также указан пакет torch_tpu==0.2.0. Используйте изолированное окружение, чтобы эксперимент не унаследовал CUDA-сборку или случайные зависимости из основной системы.
Как установить torch_tpu без CUDA
Создайте виртуальное окружение и установите CPU-колесо PyTorch вместе с зафиксированной версией TorchTPU:
python -m venv .venv
source .venv/bin/activate
pip install torch_tpu==0.2.0 \
--extra-index-url https://download.pytorch.org/whl/cpuСогласно README, если не указать дополнительный индекс, pip может подтянуть CUDA-колесо PyTorch — для TPU VM это неподходящий вариант.
Для первого переноса лучше взять готовое опубликованное колесо. Сборка из исходников требует Bazel или Bazelisk, а также описанных в репозитории флагов --config=wheel_common и, при необходимости, --config=no_rbe.
Проверяем бэкенд простым PyTorch-тестом
Запустите этот код на TPU VM:
import torch
x = torch.randn(2, 2, device="tpu")
y = torch.randn(2, 2, device="tpu")
z = x @ y
print(z)
print(z.device)Если программа успешно завершилась, а устройство обозначено как tpu, значит выделение тензоров и матричное умножение действительно прошли через TPU-бэкенд. Если устройство недоступно, сначала проверьте окружение VM и установку пакетов — менять код модели пока рано.
TorchTPU поддерживает eager-режим для отладки и torch.compile для компилируемого выполнения. В анонсе PyTorch сказано, что режим DeferAndFuse в тестах проекта был как минимум на 50% быстрее строгого eager-режима. Это результат конкретного проекта, а не универсальный бенчмарк.
Переносим рабочую PyTorch-модель
Начните с минимальных изменений. Специальные настройки распределённого запуска для TPU добавляйте только после того, как заработает однопроцессорный вариант.
- Замените выбор устройства на
torch.device("tpu")или используйтеtorch.accelerator.current_accelerator(), если эта функция поддерживается. - Перенесите модель и входные тензоры на одно и то же устройство.
- Сначала выполните один короткий батч, не включая компиляцию и распределённый режим.
- При измерениях сохраняйте размеры батча и формы тензоров: их изменение может снова запустить компиляцию.
- До масштабного переноса проверьте типичный путь выполнения модели и используемую схему распределения.
В распределённом CUDA-коде обычно используется NCCL. В опубликованном примере TorchTPU вместо него задействован TPU-бэкенд для группы процессов:
import torch
import torch.distributed as dist
device = torch.device("tpu")
dist.init_process_group(backend="tpu_dist")
model = model.to(device)Точная команда распределённого запуска зависит от окружения TPU и конкретной задачи. Замена строки с устройством сама по себе не настраивает шардирование, размещение процессов и конвейер входных данных.
Первый запуск почти наверняка будет отличаться от последующих. В документации TorchTPU описано кэширование компиляции, благодаря которому повторяющиеся графы можно использовать повторно. Поэтому время прогрева измеряйте отдельно от производительности в установившемся режиме.
Где TorchTPU уже готов к работе, а где ещё требует осторожности
Google документирует интеграции TorchTPU с vLLM, SGLang, TorchTitan и инструментами Hugging Face. В анонсе об инференсе описаны нативные TPU-сценарии для vLLM и SGLang, включая перенос KV-кэша и возможности сервинга. В анонсе Google Cloud приводятся результат TorchTitan с Qwen 3 8B на уровне 52% MFU и заявления об использовании в продакшене.
На практике ограничения выглядят менее впечатляюще, чем заголовки анонсов:
| Область | Что подтверждают источники | Что нужно проверить самостоятельно |
|---|---|---|
| Установка | Python 3.11+, torch_tpu==0.2.0, индекс CPU-версии PyTorch | Корректно ли разрешаются зависимости на вашем образе VM и с вашим wheel |
| Переносимость | Eager-режим, torch.compile, устройство tpu, поддержка распределённого выполнения | Какие операторы не поддерживаются и где появятся откаты |
| Производительность | Около 450 операций ATen и заявленный проектом результат 52% MFU для Qwen 3 8B в TorchTitan | Время прогрева, формы тензоров, размеры батчей и стоимость на вашей модели |
| Рабочий процесс разработчика | Публичные примеры, ноутбуки, инструменты для отладки и бенчмарков | Доступ к TPU, сложность Bazel и доступные средства наблюдаемости |
| Текущие пробелы | Набор возможностей всё ещё развивается | Не зависит ли ваш стек от функций, которые пока не описаны для вашего сценария |
Когда стоит переходить на TorchTPU
TorchTPU уже стоит попробовать, если у вас есть доступ к Google TPU, модели работают на PyTorch, а вам нужен единый путь для обучения или инференса. Наиболее очевидный ранний сценарий — команда, готовая профилировать и подбирать формы тензоров, а не ожидающая, что одна замена устройства автоматически даст эффективность на уровне GPU.
С критически важной нагрузкой лучше не спешить, если вам нужны зрелая матрица совместимости, предсказуемое время холодного старта или функция, от которой зависит ваш стек, но которую TorchTPU пока не предоставляет. Разворачивайте миграцию поэтапно: одна модель, один срез TPU, один фиксированный набор для оценки. Отдельно измеряйте время установки, первой компиляции, стабильную пропускную способность и корректность результатов.
FAQ по TorchTPU
TorchTPU — это то же самое, что PyTorch/XLA?
Нет. TorchTPU — более новый нативный бэкенд PyTorch, рассчитанный на eager-режим и torch.compile. PyTorch/XLA — прежний путь интеграции PyTorch с TPU. Существующий код на XLA, скорее всего, потребует отдельной оценки при миграции.
Можно ли установить TorchTPU на ноутбук без TPU?
Пакеты установить можно, но проверить выполнение на TPU без совместимой инфраструктуры Google TPU не получится. Практическая конфигурация из официального репозитория предполагает TPU VM.
Почему первый шаг в TorchTPU оказался медленным?
Во время первого выполнения граф может компилироваться, а кэши — заполняться. Последующие итерации измеряйте отдельно и при поиске повторной компиляции сохраняйте формы тензоров неизменными.
Гарантирует ли замена cuda на tpu такую же производительность?
Нет. Замена устройства подтверждает корректность переноса, но перед сравнением производительности нужно профилировать формы тензоров, шардирование, покрытие операторами и время прогрева.
Готов ли TorchTPU для продакшена?
Google описывает использование TorchTPU в продакшен-инференсе и обучении, однако зрелость публичных компонентов и отдельных функций различается. Перед переходом проверьте именно вашу модель, сервинг-движок и схему распределённого запуска.
Установите зафиксированную версию wheel на TPU VM, выполните проверку с матрицами 2×2 и сравните одну прогретую нагрузку с текущим GPU-сценарием. Так вы разделите вопросы переносимости кода, прогрева компилятора, покрытия операторов и экономики выбранного железа.