Google открыла TorchTPU для запуска моделей PyTorch на ускорителях TPU

Google представила открытый исходный код TorchTPU — программного стека, который позволяет обучать и запускать модели на тензорных процессорах TPU с использованием PyTorch. Эта разработка создана совместно с командой 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.

Подписывайтесь на наш Telegram-канал, чтобы быть в курсе всех новостей и событий Рунета.