ИИтог

Hugging Face добавила поддержку TorchTPU: модели идут на TPU напрямую

инструментыРедакция ИИтогHugging Face🔥68Рейтинг ИИтог68Влияние60Новизна65Релевантность100Неожиданность35Участники70Достоверность85
Hugging Face добавила поддержку TorchTPU: модели идут на TPU напрямую

Hugging Face опубликовала в блоге анонс нативной поддержки TPU (тензорных процессоров Google) в своей экосистеме. Основа, TorchTPU, нативный бэкенд PyTorch, благодаря которому TPU становится «просто ещё одним устройством». Transformers, Diffusers и Kernels работают на TPU напрямую: выбираете модель из Hub, отправляете её на «tpu», и всё.

Предыстория, как её излагает пост. TPU, это ускорители ИИ от Google, чипы под матричную математику нейросетей; Google проектирует их около десяти лет и использует для обучения и запуска собственных моделей, включая Gemini. TPU доступны в Google Cloud, от одного чипа до кластеров (pods) из тысяч чипов. Но программный стек вырос вокруг TensorFlow и JAX, и пользователям PyTorch входить было сложнее. Раньше PyTorch на TPU запускали через PyTorch/XLA (torch_xla): тензоры там ленивые, поэтому код выполнялся не всегда тогда, когда ожидаешь; небольшие правки могли вызывать перекомпиляцию графа, а отладку было трудно осмыслить. Для упрощения жизни Hugging Face сделала Optimum TPU (инференс и обучение на TPU), но с появлением TorchTPU он больше не нужен для запуска моделей Hugging Face через Transformers и Diffusers.

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

В Transformers (v5, задуманном как источник истины для определений моделей) модель загружают с device_map=«tpu»; входные данные отправляют методом .to(«tpu»). В примере, Qwen/Qwen3.5-9B. По умолчанию модель работает в режиме eager (немедленного выполнения); можно включить torch.compile или распределённую схему.

В Diffusers TPU, полноправное устройство (first-class device). Для прототипов и отладки хватает pipe.to(«tpu»); для развёртывания добавляется один вызов компиляции. Пайплайны и код остаются прежними, отдельный форк под TPU поддерживать не нужно. Первый вызов компилирует граф, последующие с тем же размером изображения его переиспользуют. Если модель не помещается в один чип, в from_pretrained передают TensorParallelConfig, и трансформер шардируется по чипам прямо при загрузке; многие популярные пайплайны поддерживают это из коробки, остальным нужен tp_plan. Примеры в посте запускали на Cloud TPU v6e-4 с TP=4 (модель разбита на четыре чипа), bf16, батч 1.

В Kernels Hub лежат готовые TPU-ядра на Pallas; можно подключать и ядра из библиотеки Tokamax. В примере ядро ragged-dot из Tokamax подключено к модели с архитектурой Mixture-of-Experts (Qwen3-30B-A3B-Instruct-2507). Чем больше батч, тем больше выигрыш: при размере батча 64 декодирование быстрее в 21,5 раза на Qwen3-30B-A3B и в 12,8 раза на Gemma 4 26B-A4B (пропускная способность декодирования на TPU v6e-8, ядро из Hub против реализации по умолчанию).

Отдельно авторы говорят о тестировании. Команда прогоняет тестовый набор библиотеки на TPU; ночной CI на TPU ещё только добавляют, а результаты выводят на дашборде Transformers CI (работа в процессе). Из тестов, которые запускаются на TPU, проходят 96,5%. Каждый сбой превращается в короткий скрипт воспроизведения и issue в репозитории TorchTPU: сообщено 34 проблемы, 19 уже исправлены. Заодно на TPU нашли несколько ошибок в самом Transformers (в StaticCache и планах тензорного параллелизма Qwen3) и исправили их в основной ветке.

Ключевые факты

  • TorchTPU, нативный бэкенд PyTorch для TPU: Transformers, Diffusers и Kernels запускаются на TPU без портирования на torch_xla или JAX; в Transformers достаточно device_map=«tpu».
  • Optimum TPU больше не требуется для запуска моделей Hugging Face через Transformers и Diffusers.
  • Ядра из Kernels Hub на Pallas (с Tokamax ragged-dot) ускоряют декодирование при батче 64 в 21,5 раза на Qwen3-30B-A3B и в 12,8 раза на Gemma 4 26B-A4B (TPU v6e-8, против реализации по умолчанию).
  • Из тестов Transformers, которые запускаются на TPU, проходят 96,5%; по найденным сбоям заведено 34 issue в репозитории TorchTPU, 19 исправлены.
  • Установка одной командой: pip install torch-tpu с индексом CPU-сборки; ночной CI на TPU пока только готовится.

Почему это важно

Программный стек TPU исторически строился вокруг TensorFlow и JAX, а пользователям PyTorch приходилось идти через torch_xla с ленивыми тензорами, перекомпиляциями графа и сложной отладкой. Если заявление Hugging Face подтвердится на практике, TPU становится обычным устройством для привычного кода на PyTorch и моделей из Hub, без форка и отдельного слоя вроде Optimum TPU.

Кому это важно

Разработчикам и командам, которые обучают или запускают модели на PyTorch и присматриваются к Google Cloud TPU, но не хотели переписывать код под JAX или torch_xla. Также тем, кто работает с Diffusers (изображения, видео) и с MoE-моделями, где заявлено заметное ускорение декодирования за счёт ядер из Hub.

Как это применить

На TPU VM установить пакет командой pip install torch-tpu --extra-index-url https://download.pytorch.org/whl/cpu (индекс нужен, чтобы не скачивалась большая CUDA-сборка torch). В Transformers загружать модель с device_map=«tpu», входные данные отправлять через .to(«tpu»); для скорости, torch.compile. В Diffusers, pipe.to(«tpu»), для развёртывания добавить компиляцию; для больших моделей, TensorParallelConfig. Для MoE-моделей можно подключить TPU-ядро через библиотеку Kernels, как в примере с tengomucho/ragged-dot-tpu. Цены и доступность Cloud TPU по регионам в посте не указаны.

Можно ли доверять

Источник, официальный блог Hugging Face, то есть это заявление самой компании, а не независимая проверка. Цифры ускорения (21,5 и 12,8 раза) приведены для конкретных условий: батч 64, TPU v6e-8, ядро из Hub против реализации по умолчанию; абсолютных значений пропускной способности в тексте нет, графики даны картинками. Процент прохождения тестов (96,5%) относится только к тестам, которые запускаются на TPU, общее число тестов не названо. Сравнения скорости или стоимости TPU и GPU в посте нет. Имена авторов, дата и номер версии TorchTPU в тексте не указаны.

Риски и подводные камни

Поддержка пока не полностью зелёная: авторы сами пишут, что тесты проходят не все, а ночной CI на TPU только добавляют, так что часть сценариев может ломаться (34 проблемы заведено, 19 исправлены). Ускорение растёт с размером батча, поэтому на малых батчах выигрыш может быть меньше. Первый вызов в скомпилированном режиме тратит время на компиляцию, а переиспользование графа работает при том же размере изображения. Ядра из Hub подключаются с trust_remote_code, то есть им приходится доверять стороннему коду, стоит проверять репозиторий ядра.