Google и Meta выпустили TorchTPU для нативного запуска PyTorch на TPU
Бэкенд с открытым исходным кодом, разработка которого началась в Google в середине 2025 года вместе с командой PyTorch из Meta, уже используется в продакшене некоторыми из крупнейших клиентов Google Cloud, сообщила его команда.
Иллюстрация, созданная ИИGoogle и Meta представили TorchTPU — бэкенд с открытым исходным кодом, позволяющий запускать PyTorch нативно на чипах TPU от Google. Команда проекта сообщила, что работа началась в Google в середине 2025 года совместно с командой PyTorch в Meta и что проект находится в активной разработке.
По словам команды, TorchTPU уже используется в продакшене для задач инференса и обучения некоторыми из крупнейших клиентов Google Cloud. Команда отметила, что TPU исторически были привязаны к экосистемам TensorFlow и JAX.
Команда описала TorchTPU как отход от прежних попыток. В 2020 году Google выпустила PyTorch/XLA — первую компилируемую версию PyTorch для TPU на основе Lazy Tensors. По словам команды, эта интеграция требовала навязчивых изменений в исходном коде, страдала от серьёзных задержек при компиляции, не поддерживала режим eager-выполнения PyTorch и отставала по новым возможностям распределённых вычислений. TorchTPU же считает eager-выполнение первоочередной задачей и не имеет зависимостей от JAX в своём рантайме.
Заявленная цель — чтобы перенос рабочей нагрузки на TPU требовал немногим больше, чем смены целевого устройства в коде с GPU на TPU. Команда предупредила, что одно такое изменение не обеспечит пиковую производительность TPU, поскольку формы, шардирование и операции модели могут быть настроены под GPU, но заверила, что модель всё равно должна работать «из коробки».
По словам команды, на данный момент в проекте:
- добавлена поддержка eager- и компилируемого режимов, включая новый eager-режим DeferAndFuse;
- добавлена поддержка около 450 операций ATen;
- включена поддержка пользовательских ядер, написанных с помощью JAX/Pallas и Helion DSL;
- подтверждён запуск «из коробки» на 5,000 самых загружаемых моделей Hugging Face;
- подтверждена интеграция с vLLM, SGLang, TorchTitan, Ray и TorchRec.
Команда также назвала ограничения. К устройству TPU может обращаться только один процесс в пользовательском пространстве, поэтому совместное размещение для обучения с подкреплением нельзя поддержать «из коробки»; команда сообщила, что изучает способы обойти это ограничение. Кроме того, добавляется поддержка ограниченной динамичности, поскольку базовые компилятор и рантайм предполагают, что формы тензоров известны на этапе компиляции.
Что касается точности, по умолчанию TorchTPU использует средний режим, который выполняет умножение в bfloat16 и накапливает в FP32. Более высокие режимы восстанавливают больше точности FP32 при вычислительных затратах примерно в 3 или 6 раз выше, и команда рекомендует режим по умолчанию для большинства задач глубокого обучения.