Hugging Face ❤️ TorchTPU: TPU нативно

Назад к статьям

Hugging Face ❤️ TorchTPU: TPUs нативно

Опубликовано 20 октября 2026 г.

Обновить на GitHub

Поддержать

Если вы когда-либо пытались запустить PyTorch на TPU, то, вероятно, переносили свой код на torch_xla или в итоге были вынуждены изучать JAX. С TorchTPU, нативным бэкендом PyTorch, TPU теперь является просто ещё одним устройством. Это открывает нативную поддержку TPU в экосистеме Hugging Face. Transformers, Diffusers и Kernels работают на TPU нативно. Выберите модель из Hub, отправьте её на "tpu" — и готово.

Погодите, а что такое TPU? 🤔

TPU расшифровывается как Tensor Processing Unit (тензорный процессор). TPU — это собственные ИИ-ускорители Google. Это чипы, созданные для одной задачи: матричной математики, на которую нейронные сети тратят большую часть времени. Google разрабатывает их около десяти лет и использует для обучения и обслуживания собственных моделей, включая Gemini.

TPU доступны в Google Cloud — от одного чипа до модулей с тысячами чипов. Они особенно эффективны, когда требуется много вычислений для обучения или обслуживания моделей в большом масштабе.

Но всегда существовала проблема с программным обеспечением. Стек TPU развивался вокруг TensorFlow и JAX, а пользователям PyTorch приходилось прилагать больше усилий, чтобы получить к нему доступ.

Прежний подход

До сих пор запуск PyTorch на TPU означал использование PyTorch/XLA (torch_xla). Это работало, но имело свои особенности. Тензоры были ленивыми (тензоры, не вы), поэтому код не всегда выполнялся тогда, когда вы этого ожидали. Небольшие изменения могли запускать повторную компиляцию графов, а отладка была непростой.

Чтобы упростить жизнь, мы создали Optimum TPU. Он обеспечивал поддержку инференса и обучения моделей Hugging Face на TPU. Теперь, когда TorchTPU добавляет нативную поддержку, Optimum TPU больше не требуется для запуска моделей Hugging Face с Transformers и Diffusers.

Встречайте TorchTPU 🚀

TorchTPU добавляет TPU в PyTorch как новый бэкенд. Вы пишете обычный код PyTorch. В него входят ядра Pallas, оптимизированные для TPU, поддержка распределённого обучения и интеграция с torch.compile, если вам нужна дополнительная скорость. Установка на виртуальную машину TPU выполняется одной командой:

pip install torch-tpu --extra-index-url https://download.pytorch.org/whl/cpu

Дополнительный индекс предоставляет CPU-сборку torch. Без него pip загрузит гораздо более объёмную CUDA-сборку, которая на TPU не нужна.

Ваши любимые библиотеки теперь работают на TPU

Transformers

Transformers v5 создавался вокруг одной идеи: быть источником истины для определений моделей, чтобы они могли работать где угодно. Теперь «где угодно» включает и TPU. Чтобы загрузить модель на TPU, передайте device_map="tpu". Вот и всё.

from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "Qwen/Qwen3.5-9B"
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="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-режиме. Для повышения производительности можно использовать torch.compile или запустить модель в распределённой конфигурации.

Изображения? Видео? Diffusers создаёт их на TPU 🎨

Теперь TPU — полноценное устройство в Diffusers. Для прототипирования или отладки в eager-режиме достаточно 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")  

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, и трансформер будет разделён между чипами во время загрузки, так что полной модели никогда не потребуется помещаться в одном месте. Многие популярные пайплайны поддерживают это из коробки, а для остальных достаточно указать tp_plan. В руководстве Diffusers по TPU приведён полный рецепт.

Вот несколько примеров, запущенных на Cloud TPU v6e-4:

_{TP=4 (каждая модель разделена между четырьмя чипами) · bf16 · пакет 1}

MiniMax-H3 (видео + аудио)	Qwen-Image-2.1

_{«Пушистый плюшевый маскот Hugging Face с изображения — круглое жёлтое улыбающееся лицо с румяными щеками и двумя маленькими раскрытыми ладошками — стоит на носу древнегреческой триремы […] Звук разбивающихся волн, скрип древесины, ветер в снастях, далёкие чайки и низкий эпический оркестровый гул». · 1344×768, 5,2 с, 49 шагов · 188 с в скомпилированном режиме, 262 с в eager-режиме}	_{«Капибара в шляпе волшебника читает книгу при свечах, масляная живопись» · 1024×1024, 40 шагов · 3,5 с в скомпилированном режиме, 8,0 с в eager-режиме}	_{«Неоновая вывеска магазина с надписью 'QWEN IMAGE 2.1', дождливая ночь, отражения на мокром асфальте» · 2048×2048, 40 шагов}

Kernels

В Kernels Hub доступны готовые ядра для TPU. Загрузите одно из них с помощью библиотеки Kernels и ускорьте свою модель.

Ядра TPU написаны на Pallas, которая предоставляет низкоуровневый контроль над оборудованием. Также можно использовать ядра из библиотеки Tokamax. Вот ядро ragged-dot из Tokamax, подключённое к модели Mixture-of-Experts:

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=1, 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 декодирование на Qwen3-30B-A3B выполняется в 21,5 раза быстрее, а на Gemma 4 26B-A4B — в 12,8 раза.

Пропускная способность декодирования на TPU v6e-8: ядро из Hub по сравнению со стандартной реализацией.

Мы тестируем, чтобы вам не пришлось 🧪

Мы постоянно улучшаем поддержку TPU в transformers, запуская набор тестов библиотеки. Мы работаем над добавлением ночного CI на TPU и публикацией результатов на (панели мониторинга transformers)[https://huggingface.co/spaces/transformers-community/transformers-ci-dashboard].

Панель CI Transformers с результатами для TPU рядом с другими ускорителями. Работа продолжается!

Пока ещё не всё зелёное, и нас это устраивает. Из тестов, запускаемых на TPU, проходят 96,5 %. Каждая найденная ошибка становится небольшим воспроизводимым скриптом и issue в репозитории TorchTPU. На данный момент мы сообщили о 34 проблемах, 19 из которых уже исправлены. Запуск на TPU также выявил несколько ошибок непосредственно в Transformers — например, в StaticCache и планах tensor-parallel для Qwen3. Мы исправили их в основном репозитории.

Если что-то сломается на TPU, мы хотим узнать об этом раньше вас.