Google открыла исходный код TorchTPU, обеспечив нативную поддержку PyTorch на своих TPU

20 октября 2026 года на конференции PyTorch Conference 2026 Google объявила, что TorchTPU — ее нативный стек PyTorch для тензорных процессоров (TPU) — теперь доступен в открытом исходном коде на GitHub, сообщила компания в блоге Google Cloud. Программное обеспечение разработано совместно с командой PyTorch в Meta и предназначено для запуска существующих рабочих нагрузок PyTorch на TPU с минимальными изменениями кода.

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

Джана ван Грюнен, директор по разработке PyTorch в Meta, заявила в анонсе, что TorchTPU позволяет Meta гибко совместно разрабатывать аппаратные решения с основной командой по системам машинного обучения, поддерживая при этом крупные рабочие нагрузки систем LLM, рекламы и рекомендаций. «Это позволяет нам развертывать рабочие нагрузки PyTorch на TPU Google, напрямую передавать отзывы инженерным командам и создавать необходимую всем гибкость экосистемы», — сказала она.

Интеграции с Hugging и для обслуживания моделей

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

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

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

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

Профессор Сашикумаар Ганесан, основатель Zenteiq, заявил в анонсе, что GKE Inference Gateway на базе llm-d обеспечил более чем четырехкратное увеличение совокупной пропускной способности инференса при генерации синтетических данных компании: суточная емкость генерации выросла с 1,47 млрд до 5,97 млрд токенов на каждые 1 000 ускорителей TPU v6e.

Обучение и оркестрация

TorchTitan — эталонная реализация распределенного обучения, созданная для PyTorch, — теперь включает интеграцию TorchTPU, готовую к промышленной эксплуатации и построенную на torch.compile, SimpleFSDP из TorchTitan, специализированных ядрах Pallas и оптимизациях на уровне графа. TorchTitan предлагает рецепты для обучения TorchTPU, в том числе предварительного обучения плотных моделей Llama 3 70B и 405B, Qwen 3 8B и 32B, а также MuseGlimmer 30B, модели со смешением экспертов Qwen 3 235B и диффузионного трансформера FLUX.1 12B. Google сообщила, что интеграция прошла проверку на системах с числом ускорителей Ironwood TPU (TPU v7x) до 8 000, а при обучении TorchTPU через TorchTitan достигла на Qwen 3 8B показателя использования вычислительных ресурсов модели до 52% и обеспечила почти линейное масштабирование на системах с числом ускорителей Ironwood TPU до 8 192.

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

Ray добавляет встроенную поддержку TorchTPU для организации обработки данных, обучения, обучения с подкреплением и развёртывания моделей в продуктивной среде на CPU, GPU и TPU в рамках единого рабочего процесса. Эдвард Оукс, главный архитектор Anyscale, сказал, что компании работали над тем, чтобы TPU стали ускорителями первого класса в Ray, добавив: «TorchTPU встраивает TPU Google непосредственно в этот отраслевой стандартный стек».

Предшествующая архитектура и доступность

TorchTPU был представлен публике 7 апреля 2026 года в публикации в блоге Google Developers, где описывалась архитектура «Eager First», построенная на интерфейсе PrivateUse1 в PyTorch и включающая режимы Debug, Strict и Fused Eager с общей кэш-памятью компиляции. Тогда Google заявил, что Fused Eager обеспечивает прирост производительности от 50% до 100% и более по сравнению со Strict Eager, а torch.compile в TorchTPU захватывает графы FX с помощью Torch Dynamo и преобразует их посредством XLA и StableHLO; пользовательские ядра написаны на Pallas и JAX. В публикации говорилось, что стек изначально поддерживает Distributed Data Parallel, FSDPv2 и DTensor и спроектирован для обработки различающегося выполнения на разных рангах. Среди планов развития TorchTPU на 2026 год также значился запуск общедоступного репозитория на GitHub.

Репозиторий torch_tpu теперь открыт для публики в организации google-pytorch на GitHub и распространяется по лицензии Apache 2.0. В файле README перечислены пользовательские ядра ATen, кэш компиляции для вычислений XLA, поддержка распределённого обучения с параллелизмом по данным и тензорам, интеграция с torch.compile, утилиты для отладки и тестирования производительности, а также примеры для Llama, Qwen, ResNet, DLRM, minGPT и UNet. Проект поддерживает Python 3.11 и более поздние версии, а пакет torch-tpu 0.2.0 был опубликован на PyPI 19 октября 2026 года. Он включает wheel-пакеты для Linux x86-64 с CPython версий с 3.11 по 3.14. Руководство пользователя размещено на torchtpu.readthedocs.io, а в README отмечается, что интерактивные блокноты проекта необходимо запускать на виртуальной машине TPU; проще всего это сделать на экземпляре v6e в GCE.

В том же объявлении Google сообщил, что LiteRT, ранее известный как TensorFlow Lite, присоединяется к экосистеме PyTorch: библиотека LiteRT-Torch преобразует модель PyTorch одним litert-torch.convert() вызовом, а преобразованные модели работают на устройствах под управлением Android, iOS, macOS, Windows и Linux, в веб-среде и на устройствах Интернета вещей. Набор TPU-агентов, предназначенный для автоматизации подключения, диагностики и настройки производительности инференса моделей PyTorch, доступен в ранней предварительной версии на GitHub.

В объявлении разработчикам предлагается обратиться к документации Google Cloud по TPU, чтобы начать запускать PyTorch на TPU.