Конец эпохи обходных путей
Ранее запуск PyTorch-моделей на TPU требовал использования torch_xla или перехода на JAX, что создавало барьеры для разработчиков. С выходом TorchTPU, TPU становится таким же устройством, как GPU или CPU. Теперь достаточно передать параметр device_map="tpu" в AutoModelForCausalLM или вызвать pipe.to("tpu") для Diffusers, чтобы модель заработала нативно в экосистеме Hugging Face.
Производительность: Compiled vs Eager
Интеграция с torch.compile дает значительный прирост скорости. На тестовой инфраструктуре Cloud TPU v6e-4 (4 чипа, bf16, batch 1) compiled-режим показывает существенное преимущество перед eager-режимом в задачах генерации изображений:
| Модель / Задача | Разрешение | Шаги | Compiled (с) | Eager (с) |
|---|---|---|---|---|
| Qwen-Image-2.1 (Text-to-Image) | 1344×768 | 49 | 188 | 262 |
| Stable Diffusion 3.5 (Text-to-Image) | 1024×1024 | 40 | 3.5 | 8.0 |
Ускорение LLM через Kernels Hub
Для больших языковых моделей (LLM) доступна библиотека Kernels Hub с оптимизированными ядрами на языке Pallas. Использование ragged-dot kernel от Tokamax для MoE-моделей дает кратный прирост пропускной способности декодирования. На TPU v6e-8 (8 чипов) при размере батча 64:
- Qwen3-30B-A3B: ускорение в 21.5 раз.
- Gemma 4 26B-A4B: ускорение в 12.8 раз.
Стабильность и поддержка
Команда Hugging Face активно тестирует поддержку TPU в CI. На данный момент 96.5% тестового набора проходит успешно. Уже выявлено и исправлено 19 багов, включая проблемы в StaticCache и планах тензорного параллелизма Qwen3. Установка осуществляется одной командой:
pip install torch-tpu --extra-index-url https://download.pytorch.org/whl/cpu
Важно использовать флаг --extra-index-url, чтобы избежать загрузки тяжелого CUDA-билда PyTorch, который не нужен на TPU.
Источник: Hugging Face blog ↗