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

iXBT.com · 22.10.2026, 17:46:21

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