Google представила программный комплекс TorchTPU с открытым исходным кодом, предназначенный для запуска ИИ-моделей с использованием фреймворка PyTorch на ускорителях TPU. Проект призван упростить использование ускорителей Google в обучении и инференсе ИИ; поддерживаются популярные программные средства — создавать отдельную программную среду для TPU не требуется.
Разработанный совместно с командой PyTorch в Meta✴ инструментарий TorchTPU позволяет запускать существующие проекты в PyTorch с минимальными изменениями кода, используя привычный API. Поддерживаются отладка в обычном режиме PyTorch и механизмы распределенных вычислений. Google заручилась поддержкой Hugging Face, благодаря чему представленные на ИИ-платформе открытые модели могут запускаться нативно через TorchTPU — поддерживаются библиотеки Transformers v5 и Diffusers.
Источник изображения: youtube.com/@googlecloudtech
Программный модуль vLLM для TorchTPU поддерживает одновременную обработку множества запросов, квантизацию и распределение вычислений между чипами TPU; стандартные средства управления вычислениями и обработки запросов сохранены. Для TPU адаптирован инструмент SGLang — он позволяет ускорять сложные многоэтапные задачи, в том числе выполняемые ИИ-агентами; поддерживаются повторные обращения к KV-кешу и повторяющиеся префиксы запросов, за счёт чего одни и те же вычисления могут не выполняться заново. Версия SGLang-TorchTPU сохранила механизмы выполнения задач, распределения запросов и API. Средства обмена данными с HiCache помогают сократить задержку и оптимизировать вывод при сложной многоэтапной генерации.
Совместно с Red Hat, IBM и сообществом разработчиков выпущен роутер llm-d 1.0 — он направляет запросы к экземплярам модели, где уже сохранен контекст, а также разделяет обработку входного запроса и генерацию ответа по вычислительным ресурсам, что позволяет использовать их эффективнее. В Google Kubernetes Engine этот инструмент поддерживает запуск моделей как на GPU, так и на TPU. Google интегрировала TorchTPU с TorchTitan — инструментом распределённого обучения крупных моделей; поддерживаются, в частности, Llama 3, Qwen 3 и FLUX.1. Поддержка появилась в VeRL — одном из наиболее популярных фреймворков для обучения с подкреплением, позволяющий улучшать модели в математических задачах, написании кода и многоэтапных агентских задачах.
Заявлено о поддержке TorchTPU в Ray — инструмент для распределения вычислительных задач позволяет объединить обработку данных, обучение, и инференс в рабочий процесс с использованием центральных и графических процессоров, а также TPU. В экосистему PyTorch вошла библиотека LiteRT, предназначенная для запуска моделей на пользовательских устройствах. Google также представила предварительную версию TPU agent suite — при помощи плагинов этот набор инструментов помогает автоматизировать первоначальную настройку, диагностику и оптимизацию моделей в PyTorch.
Поддержку PyTorch на своих ускорителях реализовали также Amazon, Intel и Meta✴.
Источник:
