Google відкриває вихідний код TorchTPU, переносячи PyTorch на власні TPU

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

За словами Google, завдяки цьому випуску TPU стають ще одним нативним варіантом для PyTorch поряд із GPU та CPU, підтримуючи навчання моделей, їх обслуговування та розгортання за допомогою таких інструментів, як 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 забезпечив для генерування синтетичних даних компанії сукупну пропускну здатність інференсу, вищу більш ніж у 4 рази, збільшивши щоденну кількість токенів, яку можна згенерувати, з 1,47 мільярда до 5,97 мільярда на кожні 1 000 прискорювачів TPU v6e.

Навчання та оркестрація

TorchTitan — еталонна реалізація розподіленого навчання, створена на основі PyTorch, — тепер містить, за описом Google, інтеграцію 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 повідомила, що інтеграцію перевірено на системах із кількістю до 8 000 TPU Ironwood (TPU v7x), а навчання TorchTPU через TorchTitan досягло на Qwen 3 8B використання обчислювальної потужності моделі на рівні до 52% із майже лінійним масштабуванням на системах із кількістю до 8 192 TPU Ironwood.

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

Ray додає вбудовану підтримку TorchTPU для оркестрації обробки даних, навчання, навчання з підкріпленням і обслуговування моделей PyTorch на CPU, GPU та TPU в межах одного робочого процесу. Едвард Оукс, головний архітектор Anyscale, сказав, що компанії спільно працювали над тим, щоб TPU стали повноцінними прискорювачами в Ray, додавши: «TorchTPU безпосередньо інтегрує TPU Google у цей галузевий стандартний стек».

Попередня архітектура та доступність

TorchTPU представили широкому загалу 7 квітня 2026 року в дописі в блозі Google Developers, де описано архітектуру «Eager First», побудовану на інтерфейсі PrivateUse1 у PyTorch, із режимами eager-виконання Debug, Strict і Fused, які використовують спільний кеш компіляції. Тоді 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 року; він містить колеса для 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, у вебсередовищі та на пристроях IoT. Набір агентів TPU, призначений для автоматизації початкового налаштування, діагностики та оптимізації швидкодії інференсу моделей PyTorch, доступний на GitHub у ранній попередній версії.

В анонсі розробникам пропонують скористатися документацією Google Cloud про TPU, щоб почати запускати PyTorch на TPU.