Если вы когда-либо пытались запустить PyTorch на TPU, вы, вероятно, портировали свой код в torch_xla или были вынуждены изучать JAX. Благодаря TorchTPU, родному бэкенду PyTorch, TPU теперь представляет собой просто еще одно устройство. Это открывает нативную поддержку TPU в экосистеме Hugging Face. Transformers, Diffusers и Kernels работают на TPU напрямую. Выберите модель из Хаба, отправьте ее на "tpu" — и все готово.
Стоп, а что такое TPU? 🤔
TPU расшифровывается как Tensor Processing Unit (тензорный процессор). TPU — это собственные AI-ускорители от 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 🚀
добавляет TPU в качестве нового бэкенда для PyTorch. Вы пишете обычный код на PyTorch. Он поставляется с ядрами Pallas, оптимизированными для TPU, поддержкой распределенного обучения и интеграцией torch.compile для тех случаев, когда вам нужна большая скорость. Установка на TPU VM выполняется одной командой:
Дополнительный индекс предоставляет CPU-сборку torch. Без него pip загружает гораздо большую CUDA-сборку, которая вам на TPU не нужна.
Ваши любимые библиотеки теперь на TPU
Transformers
Transformers v5 создавалась вокруг одной идеи: быть главным источником правды для определений моделей, чтобы они могли работать где угодно. Теперь «где угодно» включает и TPU. Чтобы загрузить модель на TPU, передайте device_map="tpu". Вот и все.
По умолчанию модель запускается в режиме eager. Вы также можете использовать torch.compile для повышения производительности или запустить ее в распределенной конфигурации.
Картинки? Видео? Diffusers отрисовывает их на TPU 🎨
Теперь TPU является устройством первого уровня в Diffusers. Для прототипирования или отладки в режиме eager достаточно вызвать pipe.to("tpu"). Для развертывания скомпилированный режим добавляет всего один вызов compile. Ваши пайплайны и код остаются прежними, и нет необходимости поддерживать отдельный форк для TPU.
Первый вызов компилирует модель, а каждый последующий вызов с тем же размером изображения переиспользует граф.
Модели слишком велики для одного чипа? Передайте TensorParallelConfig в from_pretrained, и трансформер будет разделен между чипами по мере загрузки, поэтому всей модели никогда не придется помещаться в одном месте. Многие популярные пайплайны поддерживают это из коробки, а для остальных достаточно параметра tp_plan. Руководство по Diffusers TPU содержит полное описание.
Вот несколько примеров, запущенных на Cloud TPU v6e-4:
TP=4 (каждая модель разделена на четыре чипа) · bf16 · размер батча 1
Ядра (Kernels)
В Kernels Hub доступны готовые к использованию ядра TPU. Загрузите одно из них с помощью библиотеки Kernels и ускорьте свою модель.
Ядра TPU написаны на Pallas, что дает вам низкоуровневый контроль над железом. Вы также можете использовать ядра из библиотеки Tokamax. Вот ядро ragged-dot от Tokamax, встроенное в модель Mixture-of-Experts:
Чем больше батч, тем выше выигрыш. При размере батча 64 декодирование на Qwen3-30B-A3B работает в 21,5 раза быстрее, а на Gemma 4 26B-A4B — в 12,8 раза быстрее.
Пропускная способность декодирования на TPU v6e-8, ядро из Хаба против реализации по умолчанию.
Мы тестируем все сами, чтобы вам не пришлось 🧪
Мы постоянно улучшаем поддержку TPU в transformers, запуская набор тестов библиотеки. Мы работаем над добавлением ежедневного CI на TPU и предоставлением результатов в (панели мониторинга transformers)[https://huggingface.co/spaces/transformers-community/transformers-ci-dashboard].
Панель мониторинга Transformers CI с результатами TPU рядом с другими ускорителями. Работа в процессе!
Здесь еще не все идеально зелено, и нас это устраивает. Из тестов, которые запускаются на TPU, проходят 96,5%. Каждая найденная нами ошибка превращается в небольшой скрипт воспроизведения и ишсю в репозитории TorchTPU. На данный момент мы сообщили о 34 проблемах, и 19 из них уже исправлены. Запуск на TPU также выявил несколько багов в самой Transformers, например, в StaticCache и планах тензорного параллелизма Qwen3. Мы исправили их в upstream.
Если что-то ломается на TPU, мы хотим узнать об этом раньше вас.
Начало работы
TPU больше не являются второстепенной задачей для пользователей PyTorch. Установите TorchTPU, возьмите модель из Хаба и отправьте ее на "tpu".
- TorchTPU на GitHub
- Ядра TPU в Хабе
- Документация по Cloud TPU
Нам не терпится увидеть, что вы создадите 🤗. Если вы столкнетесь с проблемой, откройте ишсю и расскажите нам о ней.
