Hugging Face объявляет о нативной поддержке TPU, чипов ускорения, разработанных Google, в своей экосистеме. Благодаря бэкенду TorchTPU для PyTorch, Transformers, Diffusers и Kernels работают напрямую на этих процессорах, как сообщается в заметке компании.
TPU (Tensor Processing Unit) — это чип, созданный для ускорения матричных вычислений, которые занимают большую часть времени работы нейронных сетей. Google разрабатывает эти чипы около десяти лет и использует их для обучения и эксплуатации собственных моделей, включая Gemini. Они доступны в Google Cloud — от одного чипа до кластеров в несколько тысяч единиц.
До сих пор программная экосистема TPU строилась вокруг TensorFlow и JAX, а пользователям PyTorch приходилось прилагать больше усилий. Для запуска PyTorch на TPU требовалось использовать PyTorch/XLA (torch_xla) — решение, которое в заметке описывается как рабочее, но неудобное: тензоры вычислялись лениво, поэтому код не всегда выполнялся в ожидаемый момент, а небольшие изменения могли вызывать перекомпиляцию графа. Hugging Face разработала Optimum TPU для обеспечения инференса и обучения своих моделей на TPU. С TorchTPU этот инструмент больше не нужен для запуска моделей Transformers и Diffusers.
TorchTPU — нативный бэкенд для PyTorch
TorchTPU добавляет TPU как новый бэкенд PyTorch без изменения обычного кода PyTorch. Бэкенд включает ядра Pallas, оптимизированные для TPU, поддержку распределенного обучения и интеграцию с torch.compile для пользователей, которым нужна дополнительная скорость. Установка на виртуальную машину TPU выполняется одной командой: pip install torch-tpu –extra-index-url https://download.pytorch.org/whl/cpu. В заметке уточняется, что этот индекс предоставляет CPU-версию torch, которая легче версии CUDA, загружаемой pip по умолчанию и не нужной для TPU.
Transformers, Diffusers и Kernels на TPU
В Transformers v5, цель которого — стать эталоном для определений моделей, чтобы они работали везде, загрузка модели на TPU осуществляется через параметр device_map= »tpu ». В заметке этот шаг иллюстрируется на модели Qwen/Qwen3.5-9B. По умолчанию модель работает в режиме eager, но также может использовать torch.compile или распределенную конфигурацию.
В Diffusers TPU теперь является устройством первого класса. Режим eager, подходящий для прототипирования, требует только вызова pipe.to(« tpu »), в то время как скомпилированный режим, предназначенный для развертывания, добавляет вызов компиляции. Для моделей, слишком больших для одного чипа, TensorParallelConfig, переданный в from_pretrained, распределяет трансформер между несколькими чипами при загрузке. В заметке представлены примеры, запущенные на Cloud TPU v6e-4, с моделями MiniMax-H3 для видео и аудио и Qwen-Image-2.1.
Библиотека Kernels предоставляет доступ к готовым ядрам TPU, написанным на Pallas, который предлагает низкоуровневый контроль над оборудованием. Ядра из библиотеки Tokamax также могут быть интегрированы. Согласно заметке, для батча размером 64 декодирование выполняется в 21,5 раза быстрее на Qwen3-30B-A3B и в 12,8 раза быстрее на Gemma 4 26B-A4B по сравнению с реализацией по умолчанию, на TPU v6e-8.
Непрерывное тестирование
Hugging Face сообщает о непрерывном запуске набора тестов Transformers на TPU и работе над добавлением ночного непрерывного интеграционного тестирования на этих чипах, результаты которого будут публиковаться на специальной панели мониторинга. Согласно заметке, 96,5% тестов, запущенных на TPU, проходят успешно. Каждый сбой приводит к созданию скрипта для воспроизведения и сообщению в репозитории TorchTPU: было сообщено о 34 проблемах, из которых 19 уже исправлены. Эти тесты также выявили ошибки в Transformers, в частности в StaticCache и в планах тензорного параллелизма Qwen3, которые были исправлены на раннем этапе. При возникновении проблем пользователям предлагается открыть issue.
