TorchTPU теперь с открытым исходным кодом: запуск нативного PyTorch на TPU

Источник: Google Cloud Blog•

TorchTPU теперь с открытым исходным кодом: запуск нативного PyTorch на TPU

Теперь, когда TorchTPU доступен как проект с открытым исходным кодом на GitHub, вы можете обучать, обслуживать и развертывать модели на TPU, используя привычные для вас инструменты.

Вице-президент и генеральный директор по направлению ИИ и вычислительной инфраструктуры

Вице-президент по разработке

Сегодня на конференции PyTorch Conference 2026 мы анонсировали ряд новых продуктов и интеграций, призванных помочь вам эффективно запускать PyTorch в Google Cloud или в собственной инфраструктуре. Стоит отметить, что TorchTPU, наш нативный стек PyTorch для Google TPU, теперь доступен с открытым исходным кодом на GitHub. С появлением TorchTPU чипы TPU наряду с GPU и CPU стали нативным вариантом для PyTorch. Теперь вы можете обучать, обслуживать и развертывать модели на TPU, используя инструменты, на которые вы уже полагаетесь, включая vLLM, SGLang, llm-d, TorchTitan, Ray, TorchRec, Hugging Face и LiteRT. В конечном итоге это предоставляет нашим клиентам свободу выбора среди широкого спектра оптимизированных под рабочие нагрузки вычислительных опций, включая TPU, GPU и CPU, при сохранении единой открытой программной и операционной модели для всего портфолио.

Запуск кода PyTorch в нативном режиме на TPU

Разработанный совместно с командой PyTorch в Meta, TorchTPU позволяет запускать существующие рабочие нагрузки PyTorch на TPU с минимальными изменениями кода, используя привычные API. TorchTPU дает возможность выполнять отладку моделей в режиме eager mode, а затем переходить к производительности уровня продакшена с помощью torch.compile, чтобы оптимизировать использование TPU и снизить накладные расходы на память. Поддерживая стандартные API PyTorch, расширенную диагностику и нативные распределенные примитивы «из коробки», TorchTPU обеспечивает высокую производительность и операционную эффективность с первого дня.

«В Meta продвижение границ искусственного интеллекта требует огромных вычислительных мощностей и глубокого разнообразия аппаратного обеспечения. Мы не можем быть ограничены какой-то одной аппаратной архитектурой. TorchTPU — это огромный шаг вперед, позволяющий нам гибко разрабатывать аппаратно-программные средства совместно с основной командой систем машинного обучения, одновременно поддерживая ключевые рабочие нагрузки для больших языковых моделей (LLM), рекламы и систем рекомендаций в реальных производственных сценариях. Это позволяет нам развертывать рабочие нагрузки PyTorch на Google TPU, предоставлять прямой отклик инженерным командам и создавать гибкую экосистему, необходимую каждому», — отметила Яна ван Грунен (Jana van Greunen), директор по разработке PyTorch в Meta.

Мы также заключили партнерское соглашение с Hugging Face, благодаря которому популярные модели с открытыми весами из репозитория Hugging Face Hub работают на TorchTPU в нативном режиме. В Transformers v5 развертывание моделей на TPU стало настолько простым, что достаточно передать device_map="tpu". Библиотека Diffusers работает на TPU нативно, поддерживая eager mode, torch.compile и тензорный параллелизм для масштабирования крупных генеративных моделей на нескольких ускорителях. Для поддержания неизменной надежности мы также создали специализированные конвейеры CI/CD для непрерывного ночного тестирования пакета Hugging Face Transformers на инфраструктуре TPU.

«Hugging Face стремится быть открытой и нейтральной платформой, где разработчики ИИ могут выбирать модели, фреймворки и аппаратное обеспечение, которые подходят им лучше всего. Аппаратное разнообразие имеет важнейшее значение для того, чтобы сделать ИИ с открытым исходным кодом доступным для каждого. Наше сотрудничество с Google по внедрению нативной поддержки TorchTPU в наши основные библиотеки, включая Transformers, воплощает это обязательство на практике. Вместе мы делаем так, чтобы сообществу было проще эффективно запускать популярные модели из Hub на Google Cloud TPU с использованием привычных рабочих процессов PyTorch. Предоставление разработчикам большего выбора в отношении того, как и где они создают продукты, помогает расти всей экосистеме открытого исходного кода», — заявил Джефф Будье (Jeff Boudier), вице-президент по продукту в Hugging Face.

Оптимизация инференса

По мере масштабирования моделей рассуждений и агентных моделей стоимость инференса и задержка становятся главными узкими местами в продакшене. Мы интегрировали TorchTPU в широко используемые движки обслуживания с открытым исходным кодом и объединили их со средствами маршрутизации парка на базе Kubernetes для TPU и GPU, чтобы помочь вам оптимизировать пропускную способность в расчете на доллар.

Эта основа позволяет разработчикам масштабировать рабочие нагрузки на различном оборудовании, сводя к минимуму сложность специфичной для платформ настройки. Чтобы воплотить эти улучшения на уровне фреймворков в реальных продакшен-средах, мы сотрудничаем с ведущими сообществами открытого исходного кода и индустриальными партнерами. В разделах ниже мы рассмотрим, как эти партнерства обеспечивают рекордную эффективность в задачах высокопроизводительного обслуживания с помощью vLLM, структурированного агентного выполнения с помощью SGLang и крупномасштабной распределенной оркестрации парка серверов с помощью llm-d.

vLLM: Высокопроизводительное обслуживание и эффективность памяти

Разработанный в партнерстве с Inferact, нативный бэкенд PyTorch vLLM на базе TorchTPU позволяет разработчикам развертывать модели с открытыми весами напрямую на TPU с оптимизированными, готовыми к продакшену функциями обслуживания и пользовательскими ядрами Pallas. Эта интеграция поддерживает непрерывную батчизацию, квантование и многочиповый параллелизм на оборудовании TPU, сохраняя при этом исходные рабочие модели, планировщики и API-серверы верхнего уровня для улучшения соотношения цены и производительности без необходимости изменения существующих стеков инференса.

«Наша миссия в Inferact всегда заключалась в том, чтобы сделать высокопроизводительный инференс доступным, надежным и экономичным во всей экосистеме ИИ. Мы рады стать партнером по запуску TorchTPU и обеспечить первоклассную поддержку vLLM на TPU. Объединив оптимизированные возможности PyTorch для TPU непосредственно в vLLM, разработчики могут сразу же запускать рабочие нагрузки в продакшене на TPU с высокой пропускной способностью, низкой задержкой и эффективностью использования памяти, которых они ожидают от нашего стека», — сказал Саймон Мо (Simon Mo), генеральный директор и соучредитель Inferact.

SGLang: Ускорение агентных и структурированных рабочих нагрузок

Разработанный в партнерстве с RadixArk, наш нативный бэкенд PyTorch для SGLang ускоряет altamente структурированные и агентные рабочие нагрузки на TPU. SGLang использует RadixAttention для повторного использования KV-кэша между запросами, что позволяет полностью пропускать повторные вычисления для повторяющихся префиксов промптов. SGLang-TorchTPU сохраняет исходные рабочие модели, планировщики и API-серверы верхнего уровня, используя бэкенд передачи с дезагрегацией, который также выполняет роль бэкенда ввода-вывода для HiCache с целью оптимизации задержки и пропускной способности для сложных многошаговых генераций.

«В RadixArk интеграция TorchTPU значительно ускорила нашу работу по внедрению поддержки SGLang для различных процессоров. Это дало нам простой путь к нативному ускорению на TPU, обеспечив колоссальный прирост эффективности и гибкие возможности выбора вычислительных средств непосредственно для наших клиентов», — отметил Бангхуа Чжу (Banghua Zhu), технический директор и соучредитель RadixArk.

llm-d для распределенного инференса LLM

Для запуска инференс-движков на крупных парках оборудования мы выпускаем llm-d 1.0 совместно с Red Hat, IBM и сообществом, насчитывающим более 500 участников и 25 организаций. Будучи проект-песочницей Cloud Native Computing Foundation (CNCF), llm-d объединяет два сообщества с открытым исходным кодом: PyTorch, где vLLM и SGLang обслуживают модели, и Kubernetes, где распределенные системы работают и масштабируются в продакшене. Маршрутизатор llm-d направляет каждый запрос на реплику, которая уже хранит его контекст, используя маршрутизацию с учетом кэша префиксов. Он также разделяет ресурсоемкий этап префила (prefill) и требовательный к памяти этап декодирования (decode). llm-d 1.0 делает эти возможности готовыми к производственному использованию благодаря многопоточному контролю потоков, высокой доступности маршрутизатора, стабильным API версии v1 и маршрутизации с учетом сессий для агентов. В Google Kubernetes Engine (GKE) вы можете выбирать между запуском на графических процессорах (GPU) или Google TPU с использованием общей конечной точки llm-d, подбирая аппаратное обеспечение, наиболее подходящее для каждой конкретной рабочей нагрузки.

"В Zenteiq мы используем GKE Inference Gateway на базе llm-d для ускорения генерации синтетических данных для наших суверенных научных базовых моделей (Scientific Foundation Models). За счет маршрутизации с учетом нагрузки и кэша префиксов мы добились увеличения совокупной пропускной способности инференса более чем в 4 раза, увеличив ежедневную емкость генерации токенов с 1,47 млрд до 5,97 млрд токенов на 1000 ускорителей TPU v6e. Открытая основа llm-d в интеграции с GKE позволяет нам эффективно обслуживать множество семейств моделей в больших масштабах без необходимости разрабатывать и поддерживать собственный планировщик инференса". — Проф. Сашикумаар Ганесан (Sashikumaar Ganesan), основатель Zenteiq.ai

Масштабируемое обучение

Масштабирование обучения PyTorch на тысячах ускорителей требует отказоустойчивых распределенных фреймворков с поддержкой компиляции. Мы интегрировали TorchTPU в основные фреймворки сообщества PyTorch для обучения, обучения с подкреплением (RL) и оркестрации.

TorchTitan: распределенное высокопроизводительное обучение в масштабе

TorchTitan — эталонная реализация на базе PyTorch для масштабируемого распределенного обучения, включающая интеграцию производственного уровня с TorchTPU. Подключаясь к torch.compile из PyTorch и дружественному к компиляции SimpleFSDP в TorchTitan, специализированным ядрам Pallas и оптимизациям на уровне графов, разработчики могут эффективно создавать и обучать модели на TPU в больших масштабах.

TorchTitan предоставляет рецепты для предварительного обучения крупных плотных моделей (таких как Llama 3 70B/405B, Qwen 3 8B/32B и MuseGlimmer 30B), моделей смещения экспертов (MoE), таких как Qwen 3 235B, и диффузионных трансформеров (таких как FLUX.1 12B), а также прошел валидацию при работе с количеством до 8000 TPU Ironwood (TPU v7x). Помимо первоначального предварительного обучения, TorchTitan служит масштабируемым инструментом обучения в конвейерах пост-обучения, таких как VeRL для алгоритма GRPO (Group Relative Policy Optimization) и выравнивания с помощью RL. Для тестирования крупномасштабной эффективности на требовательных открытых архитектурах пути обучения TorchTPU через TorchTitan продемонстрировали до 52% использования флопсов модели (MFU) на Qwen 3 8B, показав практически линейное масштабирование и стабильность на 8192 TPU Ironwood.

VeRL: масштабируемое пост-обучение и обучение с подкреплением в предварительном доступе

Обучение с подкреплением (RL) — это фундаментальный метод для передовых лабораторий, развивающих возможности моделей, а также для предприятий, настраивающих модели под конкретные задачи предметной области. Тем не менее, оркестрация распределенных рабочих нагрузок RL остается одной из самых сложных задач в современной инфраструктуре ИИ. Устраняя этот разрыв, VeRL быстро стал одним из самых популярных фреймворков RL, поддерживая все: от проверяемых наград за математику и программирование до многошагового использования агентских инструментов в огромных масштабах.

Чтобы расширить выбор оборудования для этого процветающего сообщества, мы добавили бэкенд TPU и сквозной эталонный рецепт в VeRL. Наш эталонный рецепт GRPO для наборов данных по математическим рассуждениям, созданный на основе дезагрегированной архитектуры "тренер-семплер", призванной максимизировать эффективность обучения, демонстрирует, как разработчики могут развертывать рабочие нагрузки VeRL на TPU, сохраняя существующий аппаратно-независимый код моделей и алгоритмов.

Ray: оркестрация гетерогенных вычислений

Мы также добавили нативную поддержку TorchTPU в Ray, что позволяет вам оркестрировать обработку данных PyTorch, обучение, RL и инференс на CPU, GPU и TPU в рамках единого рабочего процесса. С помощью всего нескольких строк конфигурации ресурсов вы можете планировать каждый этап гетерогенного конвейера ИИ на том ускорителе, который лучше всего подходит для этой задачи.

"На саммите Ray Summit мы говорили о стеке PARK — PyTorch, модели ИИ, Ray и Kubernetes, — потому что именно на этом команды создают промышленный ИИ сегодня. TorchTPU встраивает Google TPU непосредственно в этот стандартный для отрасли стек. Мы также тесно сотрудничали с командой Google, чтобы обеспечить поддержку TPU в качестве ускорителя первого класса в Ray. В сочетании с нативной поддержкой PyTorch команды теперь могут беспрепятственно разрабатывать, масштабировать и мониторить свои рабочие нагрузки на TPU". — Эдвард Оукс (Edward Oakes), главный архитектор, Anyscale

Гибкость развертывания и операционная эффективность

Истинная переносимость требует стандартизации инструментов развертывания на различных профилях оборудования, включая периферийные устройства и облачные ускорители.

LiteRT пополняет экосистему PyTorch

Многим моделям PyTorch необходимо работать на телефонах, ноутбуках, в браузерах и на встроенных устройствах. Сегодня LiteRT (ранее TensorFlow Lite), инструмент высокопроизводительного локального ИИ-инференса от Google, присоединяется к экосистеме PyTorch.

С помощью LiteRT-Torch вы можете конвертировать модель PyTorch одним вызовом litert-torch.convert(). Решение основано на torch.export() и базовых операторах ATen, поэтому код вашей модели остается в PyTorch. Для больших языковых моделей (LLM) LiteRT Torch экспортирует модели Hugging Face Transformer для работы с LiteRT-LM. Сконвертированные модели работают на устройствах Android, iOS, macOS, Windows, Linux, в веб-среде и на IoT-устройствах, обеспечивая аппаратное ускорение на процессорах (через XNNPack), графических процессорах (через открытый ML Drift) и NPU от основных производителей чипсетов через единый API.

Набор TPU-агентов в раннем предварительном доступе

Для расширения внедрения PyTorch на TPU мы объявляем, что набор TPU-агентов доступен в режиме предварительного просмотра на GitHub. Интегрируясь с помощью плагинов с различными комбинациями кодинг-оснастки и моделей, набор TPU-агентов помогает автоматизировать подключение (onboarding), диагностику, пользовательские запросы и настройку производительности инференса для многих моделей PyTorch.

Начало работы с PyTorch в Google Cloud

В совокупности эти релизы подтверждают наше стремление сделать программную экосистему Google TPU открытой и удобной для разработчиков средой для сообщества PyTorch. Внедрив нативную поддержку TPU на верхнем уровне стека PyTorch в TorchTPU, Hugging Face, vLLM, SGLang, llm-d, TorchTitan, VeRL и Ray, а также добавив LiteRT в экосистему PyTorch, мы предоставляем вам гибкость для создания, адаптации и масштабирования ИИ на ваших условиях.

Посетите документацию Google Cloud, чтобы начать работу с PyTorch на TPU уже сегодня.

  • AI инфраструктура
  • TPU
  • Разработка приложений
  • Открытый исходный код

О чём эта статья

Ещё в разделе «Кибербезопасность»

Все →

Ещё от Mandiant