Google открыла TorchTPU: модели PyTorch теперь можно запускать на ускорителях TPU
iXBT.com ·
Google открыла исходный код TorchTPU — программного стека, который позволяет обучать и запускать модели на тензорных процессорах TPU с использованием PyTorch. Разработка создана совместно с командой PyTorch и призвана сделать TPU полноценным вариантом ускорителя наряду с центральными и графическими процессорами. Существующие программы на PyTorch можно переносить с минимальными изменениями кода, используя привычные интерфейсы для отладки, компиляции и распределённых вычислений.
TorchTPU поддерживает пошаговый режим выполнения для отладки моделей, а также torch.compile, который компилирует вычисления для повышения производительности и эффективности использования памяти. Google также добавила поддержку в экосистему Hugging Face: в Transformers 5 запуск моделей на TPU можно указать через параметр 'device_map="tpu"', а библиотека Diffusers поддерживает выполнение и распределение вычислений между несколькими ускорителями. Для проверки совместимости Transformers с TPU компания организовала регулярные автоматические тесты.
Отдельное направление — запуск больших языковых моделей в открытых движках инференса. Для vLLM и SGLang подготовлены бэкенды, позволяющие выполнять модели на TPU. SGLang, в частности, использует повторное применение кэша ключей и значений внимания (KV cache), чтобы не пересчитывать уже обработанные части повторяющихся запросов. Google также развивает интеграцию с llm-d — системой распределения запросов между вычислительными узлами, которая учитывает наличие сохранённого контекста и позволяет разделять этапы обработки входного запроса и генерации ответа.
TorchTPU интегрировали и в инструменты обучения моделей. В связке с TorchTitan он поддерживает распределённое обучение больших нейросетей, а также специализированные вычислительные ядра Pallas. По данным Google, при обучении модели Qwen 3 8B система достигла утилизации вычислительных возможностей ускорителей на уровне до 52% и масштабировалась до 8192 TPU Ironwood. Эти результаты относятся к конкретной конфигурации и не означают аналогичного ускорения для любой модели или задачи.
Открытие TorchTPU расширяет возможности разработчиков выбирать аппаратную платформу без необходимости полностью перестраивать программный стек под оборудование одного производителя. Вместе с поддержкой vLLM, SGLang, Hugging Face, Ray и других проектов Google формирует более единую среду для обучения и запуска ИИ-моделей на CPU, GPU и TPU. Сам по себе открытый код не гарантирует одинаковой производительности на разных ускорителях, но снижает барьер для переноса существующих проектов на TPU.