Если вы когда-либо пытались запустить модели PyTorch на ускорителях TPU (Tensor Processing Unit), вы, вероятно, сталкивались с необходимостью портировать свой код на torch_xla или, что еще более вероятно, изучали JAX. Оба пути требовали значительных усилий, понимания внутренних механизмов компиляции графов и адаптации привычного кода под специфические требования аппаратного обеспечения. Однако эра компромиссов подходит к концу. С выходом TorchTPU, нативного бэкенда для PyTorch, TPU теперь становятся просто еще одним устройством в арсенале разработчика. Это открывает двери для полноценной поддержки TPU в экосистеме Hugging Face: библиотеки Transformers, Diffusers и Kernels теперь работают на TPU «из коробки».
Простота использования стала главным девизом этого обновления. Теперь, чтобы запустить модель, вам нужно всего лишь выбрать ее с Hugging Face Hub, отправить ее на устройство "tpu" — и всё готово. Больше не нужно писать обертки, настраивать сложные пайплайны компиляции или изучать новую парадигму программирования. В этой статье мы подробно разберем, что такое TorchTPU, как он работает, какие преимущества дает по сравнению с предыдущими решениями, и как начать использовать его в своих проектах уже сегодня.
01Что такое TPU и почему это важно?
TPU расшифровывается как Tensor Processing Unit (тензорный процессор). Это специализированные чипы искусственного интеллекта, разработанные компанией Google. В отличие от универсальных CPU или даже графических процессоров (GPU), которые могут выполнять широкий спектр вычислений, TPU созданы для одной конкретной задачи: выполнения матричных операций, которые составляют основу вычислений в нейронных сетях.
Google занимается проектированием TPU уже около десяти лет. Эти ускорители используются для обучения и обслуживания собственных моделей Google, включая передовую языковую модель Gemini. TPU доступны через Google Cloud — от одиночных чипов для небольших экспериментов до кластеров (pods) из тысяч чипов для масштабного обучения и высоконагруженного инференса. Именно в таких масштабах TPU показывают свою истинную мощь, обеспечивая высокую пропускную способность и энергоэффективность.
Однако главным препятствием для массового внедрения TPU в сообществе PyTorch всегда было программное обеспечение. Стек TPU исторически развивался вокруг TensorFlow и JAX. Для пользователей PyTorch доступ к этим ускорителям был сопряжен с дополнительными сложностями, требующими глубокого понимания низкоуровневых деталей. TorchTPU устраняет этот барьер, интегрируя TPU непосредственно в ядро PyTorch.
02Проблемы старого подхода: PyTorch/XLA и JAX
До появления TorchTPU единственным официальным способом запуска PyTorch на TPU была библиотека PyTorch/XLA (torch_xla). Хотя этот подход работал, он имел ряд существенных недостатков, которые часто отпугивали разработчиков.
Ключевой особенностью torch_xla является использование ленивых вычислений (lazy execution). Это означает, что ваши тензоры не вычисляются немедленно, как в стандартном PyTorch. Вместо этого строится вычислительный граф, который компилируется и выполняется только при необходимости. Хотя это позволяет оптимизировать производительность, оно создает серьезные проблемы для отладки и понимания потока выполнения программы. Код может вести себя непредсказуемо: небольшие изменения в логике могут вызывать полную перекompиляцию графа, что приводит к резким падениям производительности или неожиданным ошибкам.
Кроме того, отладка таких систем требует специфических знаний. Разработчикам приходилось постоянно помнить о том, что их код выполняется не так, как они его пишут, а так, как его интерпретирует компилятор XLA. Это создавало высокий порог входа и увеличивало время на разработку и тестирование.
Для упрощения работы с моделями Hugging Face на TPU ранее была создана библиотека Optimum TPU. Она предоставляла поддержку инференса и обучения для моделей Hugging Face, но требовала дополнительных настроек и интеграции. Теперь, с появлением TorchTPU, Optimum TPU перестает быть обязательным звеном для запуска моделей Transformers и Diffusers на TPU, хотя и продолжает развиваться для других задач.
03Введение TorchTPU: Нативная интеграция
TorchTPU добавляет TPU как новый бэкенд непосредственно в PyTorch. Это означает, что вы пишете обычный, стандартный код на PyTorch, и он работает на TPU так же естественно, как и на CPU или GPU. Библиотека включает в себя оптимизированные ядра Pallas (Pallas kernels), поддержку распределенного обучения и интеграцию с torch.compile для максимальной производительности.
Установка TorchTPU на виртуальной машине TPU (TPU VM) в Google Cloud предельно проста. Достаточно выполнить одну команду:
pip install torch-tpu --extra-index-url https://download.pytorch.org/whl/cpuВажно отметить параметр --extra-index-url. Он указывает pip использовать репозиторий с версией PyTorch для CPU. Без этого параметра pip попытается установить версию для CUDA, которая значительно больше по размеру и не нужна на TPU, так как TPU не используют архитектуру NVIDIA CUDA. Это упрощает процесс установки и экономит дисковое пространство.
--extra-index-url при установке torch-tpu. Это гарантирует, что вы получите легковесную сборку, совместимую с TPU, и избежите конфликтов зависимостей с CUDA-версией PyTorch.04Transformers на TPU: Простота использования
Библиотека Transformers версии 5 была построена вокруг идеи быть «единым источником истины» для определений моделей, позволяя им запускаться где угодно. Теперь «где угодно» включает в себя и TPU. Загрузка модели на TPU стала тривиальной задачей.
Для загрузки модели достаточно передать параметр device_map="tpu" в метод from_pretrained. Модель будет загружена непосредственно на TPU, и вы можете работать с ней, используя привычный интерфейс PyTorch.
from transformers import AutoModelForCausalLM, AutoTokenizer
model_id = "Qwen/Qwen3.5-9B"
model = AutoModelForCausalLM.from_pretrained(
model_id,
device_map="tpu" # модель загружается напрямую на TPU
)
tokenizer = AutoTokenizer.from_pretrained(model_id)
inputs = tokenizer("The future of AI is", return_tensors="pt").to("tpu")
out = model.generate(**inputs, max_new_tokens=64)
print(tokenizer.decode(out[0], skip_special_tokens=True))По умолчанию модель работает в режиме eager mode (режим немедленного выполнения), что обеспечивает максимальную совместимость и простоту отладки. Однако для достижения максимальной производительности вы можете использовать torch.compile или настроить распределенное обучение. Код остается чистым и понятным, без необходимости писать специфичные для TPU обертки.
05Diffusers на TPU: Генерация изображений и видео
TPU теперь является устройством первого класса в библиотеке Diffusers. Это открывает возможности для генерации изображений и видео на TPU с той же простотой, что и на GPU.
Для прототипирования или отладки достаточно переместить пайплайн на TPU с помощью метода pipe.to("tpu"). Для развертывания в продакшене можно добавить вызов compile, что значительно ускорит генерацию. Ваш код и пайплайны остаются неизменными, и вам не нужно поддерживать отдельную ветку кода для TPU.
import torch
from diffusers import StableDiffusion3Pipeline
pipe = StableDiffusion3Pipeline.from_pretrained(
"stabilityai/stable-diffusion-3.5-medium",
dtype=torch.bfloat16
)
pipe.to("tpu") # eager mode: это единственное изменение
# опционально: compiled mode
pipe.transformer.compile(backend="tpu", fullgraph=True, dynamic=False)
image = pipe("a corgi astronaut planting a flag on a pizza moon").images[0]Первый вызов генерации запускает компиляцию графа, и все последующие вызовы с тем же размером изображения используют уже скомпилированный граф, что обеспечивает значительное ускорение. Если модель слишком велика для одного чипа, вы можете передать TensorParallelConfig в метод from_pretrained. В этом случае трансформер будет разделен (sharded) между несколькими чипами во время загрузки, и полная модель никогда не должна помещаться в память одного устройства. Многие популярные пайплайны поддерживают это из коробки, а для остальных достаточно указать tp_plan. Подробные инструкции доступны в руководстве по TPU для Diffusers.
Примеры производительности
Давайте посмотрим на реальные примеры, запущенные на кластере Cloud TPU v6e-4 (4 чипа). Все модели запущены в режиме Tensor Parallelism (TP=4), с точностью bf16 и размером батча 1.
- MiniMax-H3 (видео + аудио): Запуск мультимодальной модели.
- Qwen-Image-2.1: Генерация изображения «пушистый плюшевый маскот Hugging Face из картинки, круглое желтое улыбающееся лицо с румяными щеками и двумя маленькими открытыми руками, стоит на носу древнегреческой триеры […] Звук разбивающихся волн, скрип дерева, ветер в такелаже, отдаленные чайки и низкое эпическое оркестровое звучание».
Результаты для Qwen-Image-2.1:
- Разрешение 1344×768, 5.2 секунды, 49 шагов: 188 секунд в compiled режиме против 262 секунд в eager режиме.
- Разрешение 1024×1024, 40 шагов: 3.5 секунды в compiled режиме против 8.0 секунд в eager режиме.
- Разрешение 2048×2048, 40 шагов: Генерация знака «QWEN IMAGE 2.1» в дождливую ночь с отражениями на мокром асфальте.
06Kernels Hub: Ускорение через Pallas
Библиотека Kernels предоставляет готовые к использованию ядра, оптимизированные для TPU. Загрузив ядро с Kernels Hub, вы можете значительно ускорить работу своей модели. Ядра TPU написаны на языке Pallas, который дает низкоуровневый контроль над аппаратным обеспечением. Вы также можете использовать ядра из библиотеки Tokamax.
Например, можно использовать ядро ragged-dot из Tokamax для ускорения работы модели Mixture-of-Experts (MoE). Вот как это выглядит в коде:
from kernels import get_kernel
from transformers import AutoModelForCausalLM
from transformers.integrations.moe import ALL_EXPERTS_FUNCTIONS
repo_id = "tengomucho/ragged-dot-tpu"
kernel = get_kernel(repo_id, version=0, trust_remote_code=[repo_id])
ALL_EXPERTS_FUNCTIONS.register("tokamax_ragged_dot", kernel.experts_forward)
model = AutoModelForCausalLM.from_pretrained(
"Qwen/Qwen3-30B-A3B-Instruct-2507",
experts_implementation="tokamax_ragged_dot"
)Чем больше размер батча, тем больше выигрыш в производительности. При размере батча 64 декодирование становится в 21.5 раза быстрее на Qwen3-30B-A3B и в 12.8 раза быстрее на Gemma 4 26B-A4B по сравнению с реализацией по умолчанию. Это демонстрирует огромную пользу от использования специализированных ядер для конкретных архитектур моделей.
07Тестирование и надежность
Hugging Face непрерывно улучшает поддержку TPU в библиотеке Transformers, запуская полный набор тестов. В настоящее время команда работает над добавлением ночных сборок (nightly CI) на TPU и публикацией результатов в Transformers CI Dashboard.
На данный момент не все тесты проходят успешно, и команда открыто признает это. Из тестов, которые запускаются на TPU, 96.5% проходят успешно. Каждая неудача превращается в небольшой воспроизводимый скрипт и issue в репозитории TorchTPU. На данный момент сообщество уже сообщило о 34 проблемах, и 19 из них уже исправлены. Важно отметить, что запуск тестов на TPU также помог выявить некоторые баги в самой библиотеке Transformers, например, в StaticCache и планах тензорного параллелизма для Qwen3. Эти баги были исправлены в основной ветке.
08Что это значит на практике
Для разработчиков и исследователей в области ИИ TorchTPU означает радикальное упрощение работы с TPU. Больше не нужно выбирать между экосистемой PyTorch и производительностью TPU. Теперь вы можете:
- Использовать знакомые инструменты: Продолжайте писать код на PyTorch, используя Transformers и Diffusers, без необходимости изучать JAX или разбираться в сложностях XLA.
- Масштабироваться легко: Благодаря нативной поддержке распределенного обучения и тензорного параллелизма, вы можете легко масштабировать свои модели на кластеры TPU.
- Достигать высокой производительности: Использование
torch.compileи специализированных ядер Pallas позволяет достичь производительности, сопоставимой с оптимизированными решениями на JAX, но с гораздо меньшими усилиями. - Экономить время: Установка и запуск моделей занимают считанные минуты. Простой API
device_map="tpu"делает интеграцию в существующие проекты минимальной.
TPU больше не являются «побочным квестом» для пользователей PyTorch. С TorchTPU они становятся полноценной частью экосистемы. Установите TorchTPU, выберите модель с Hugging Face Hub и отправьте ее на "tpu". Мы с нетерпением ждем, что вы построите с помощью этих технологий. Если вы столкнетесь с проблемами, откройте issue в репозитории — команда Hugging Face активно работает над улучшением поддержки TPU вместе с сообществом.
Для начала работы вы можете посетить:
Добро пожаловать в новую эру нативной поддержки TPU в PyTorch!
Источник: Hugging Face ↗