Hugging Face интегрирует TPUs в PyTorch через TorchTPU
Hugging Face Blog · оригинал
Материал подготовлен автоматизированной редакционной системой. Факты можно сверить по указанному первоисточнику.
Компания Hugging Face представила TorchTPU — нативный бэкенд, позволяющий использовать ускорители TPU в экосистеме PyTorch без необходимости перехода на JAX или TensorFlow. Теперь модели из библиотеки Transformers и пайплайны Diffusers запускаются на TPU стандартными командами, что устраняет исторические сложности с отладкой и компиляцией графов.
Интеграция TPUs в PyTorch через нативный бэкенд TorchTPU снимает главное барьер для разработчиков: необходимость переписывать код под JAX или использовать обвязку torch_xla. Это делает высокопроизводительные ускорители Google доступными для всей экосистемы Hugging Face, включая Transformers и Diffusers, и открывает путь к эффективному обучению и инференсу больших моделей на масштабируемых кластерах TPU.
Исторически использование ускорителей TPU (Tensor Processing Unit) в связке с PyTorch требовало от разработчиков значительных усилий. Стандартным путем была работа через PyTorch/XLA (torch_xla), что приводило к ленивому вычислению тензоров, сложностям с отладкой и неожиданным перекомпиляциям графов. Альтернативой часто становился переход на JAX или TensorFlow. Hugging Face ранее предлагала решение Optimum TPU для поддержки инференса и обучения своих моделей, но теперь ситуация изменилась благодаря появлению TorchTPU — нативного бэкенда для PyTorch.
TorchTPU добавляет TPU как полноценное устройство в экосистему PyTorch. Это означает, что разработчики могут писать стандартный код PyTorch, используя встроенные ядра Pallas, оптимизированные под TPU, а также поддержку распределенного обучения и интеграцию с torch.compile для повышения производительности. Установка бэкенда на виртуальной машине TPU выполняется одной командой pip, при этом важно указать дополнительный индекс для загрузки CPU-сборки PyTorch, чтобы избежать скачивания ненужных CUDA-библиотек.
Для пользователей библиотеки Transformers v5 процесс запуска модели на TPU стал максимально простым. Достаточно передать параметр device_map="tpu" при загрузке модели из Hugging Face Hub. По умолчанию модель работает в режиме eager, но при необходимости можно активировать torch.compile для ускорения. Это позволяет использовать такие модели, как Qwen3.5-9B, без изменения основной логики кода.
Библиотека Diffusers также получила статус TPU как первого класса устройства. Для прототипирования достаточно переместить пайплайн на устройство с помощью pipe.to("tpu"). Для продакшн-сценариев рекомендуется использовать compiled mode, добавив один вызов compile. Это сохраняет совместимость существующего кода и устраняет необходимость в отдельной ветке разработки для TPU. Например, пайплайн Stable Diffusion 3.5 можно запустить на TPU с минимальными изменениями.
Для моделей, которые не помещаются в память одного чипа, TorchTPU предоставляет механизм тензорного параллелизма. Передавая TensorParallelConfig при загрузке, разработчики могут распределить трансформер между несколькими чипами. Многие популярные пайплайны поддерживают эту функцию из коробки, а для остальных достаточно указать план параллелизма (tp_plan). Это критически важно для работы с большими языковыми моделями и генеративными моделями изображений.
Hugging Face также продвигает использование готовых ядер TPU через Kernels Hub. Эти ядра написаны на языке Pallas, который дает низкоуровневый контроль над аппаратным обеспечением. Разработчики могут загружать оптимизированные ядра, такие как ragged-dot из библиотеки Tokamax, для ускорения конкретных операций, например, в моделях Mixture-of-Experts. По данным компании, при размере батча 64 скорость декодирования на TPU v6e-8 увеличивается в 21,5 раза для Qwen3-30B-A3B и в 12,8 раза для Gemma 4 26B-A4B по сравнению с реализацией по умолчанию.
Команда Hugging Face активно тестирует поддержку TPU, прогоняя тестовый набор библиотеки Transformers. На данный момент 96,5% тестов проходят успешно. Найденные ошибки фиксируются в репозитории TorchTPU; из 34 открытых проблем уже исправлено 19. Кроме того, тестирование на TPU помогло выявить и исправить несколько багов непосредственно в коде Transformers, включая проблемы с StaticCache и планами тензорного параллелизма для Qwen3. Источник: web:Hugging Face Blog — https://huggingface.co/blog/hf-torchtpu