Представьте: вы написали сложный пайплайн обучения на PyTorch, и вам нужно перенести его с GPU на TPU. Раньше это означало недели боли — переписывание кода, изучение новых API, борьба с JAX или XLA. Теперь достаточно одной строки: .to("tpu") вместо .to("cuda"). Именно это обещает Google с выходом TorchTPU — нативного бэкенда PyTorch для TPU, открытого 20 октября 2026 года на конференции PyTorch Conference в Сан-Хосе.

Что такое TorchTPU и почему это важно

TorchTPU — это результат открытого сотрудничества между Google и Meta, который вводит нативный и принципиально выстроенный бэкенд PyTorch для Google TPU. Проект не появился из ниоткуда: работа над ним началась в середине 2025 года, а официальный открытый релиз состоялся на PyTorch Conference North America в Сан-Хосе 20 октября 2026 года.

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

ℹ Контекст релиза
За день до анонса TorchTPU Amazon выпустил TorchNeuron для своих чипов Trainium. Оба проекта направлены на снижение зависимости разработчиков от NVIDIA CUDA — де-факто стандарта AI-вычислений последнего десятилетия.

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

Архитектура: как это работает под капотом

TorchTPU — это новый инженерный стек, спроектированный для обеспечения нативного высокопроизводительного опыта запуска PyTorch-нагрузок на TPU-инфраструктуре Google с минимальными изменениями кода. Он реализует подход «Eager First» с несколькими режимами выполнения и использует XLA-компилятор для оптимизации распределённого обучения на массивных кластерах.

Три режима выполнения

Система поддерживает три режима Eager: Strict, Debug и DeferAndFuse, а также интегрируется с torch.compile через Dynamo.

  • Strict Eager — строгое последовательное выполнение операций, полный контроль.
  • Debug — расширенная диагностика для отладки.
  • DeferAndFuse — гибридный режим: DeferAndFuse откладывает и сливает несколько ATen-операций в более крупные блоки перед выполнением. Он может быть на 50% и более быстрее режима Strict Eager, одновременно снижая пиковое потребление HBM-памяти за счёт уменьшения временного хранения при слиянии операций.

Стек компилятора

Стек захватывает FX-графы с помощью Torch Dynamo и понижает их через XLA и StableHLO, с кастомными ядрами, написанными на Pallas и JAX.


graph TD
    A[PyTorch модель] --> B[Torch Dynamo / FX Graph]
    B --> C{Режим выполнения}
    C -->|Eager| D[Strict / Debug]
    C -->|Compiled| E[DeferAndFuse / torch.compile]
    D --> F[XLA Runtime]
    E --> G[StableHLO]
    G --> F
    F --> H[TPU Hardware]
    H --> I[Результат]

Поддержка операций

Под капотом релиз поддерживает около 450 ATen-операций. ATen — это базовая библиотека тензорных операций PyTorch, основные математические строительные блоки, из которых собираются модели.

Стек поддерживает Distributed Data Parallel, FSDPv2 и DTensor из коробки.

Производительность: цифры и бенчмарки

Google заявляет, что интеграция была проверена на конфигурациях до 8 000 чипов Ironwood TPU (TPU v7x) и что обучение через TorchTitan достигло 52% утилизации Model Flops (MFU) на Qwen 3 8B при масштабировании до 8 192 чипов Ironwood TPU.

Проект прошёл проверку на 4 821 из 4 917 протестированных Hugging Face моделей (98%) и совпал с GPU-числами на 638 из 657 семейств архитектур (97%).

Обучение на кластере из 64 чипов TPU v7 достигло 1 828 токенов в секунду против 1 937 у JAX/MaxText. Разница минимальна, но речь идёт о PyTorch-коде без каких-либо переписываний.

ПараметрTorchTPUPyTorch/XLAJAX на TPU
ВендорGoogle + MetaGoogle / PyTorchGoogle
PyTorch devicetpuxlaНе PyTorch
Eager-режим✅ Да⚠️ Ограниченно❌ Нет
torch.compile✅ Да⚠️ Частично❌ Нет
JAX dependency❌ Нет❌ Нет✅ Требуется
ЛицензияApache 2.0BSDApache 2.0
Версия на PyPI0.2.0Отдельный пакет—
💡 Совет по производительности
Для максимальной скорости используйте режим DeferAndFuse в сочетании с torch.compile. Это снижает overhead на запуск ядер и уменьшает потребление HBM-памяти — особенно критично при работе с большими языковыми моделями.

Интеграции: экосистема PyTorch на TPU

Одна из главных сильных сторон TorchTPU — глубокая интеграция с популярными инструментами экосистемы PyTorch.

Hugging Face

Google сотрудничал с Hugging Face, чтобы open-weight модели из Hugging Face Hub запускались нативно на TorchTPU. В Transformers v5 деплой модели на TPU выполняется передачей device_map="tpu" при загрузке модели, а библиотека Diffusers работает нативно на TPU с поддержкой eager-режима, torch.compile и тензорного параллелизма.

Инференс-серверы

Интеграции инференса включают vLLM с ядрами Pallas, SGLang с RadixAttention и llm-d 1.0 (совместно с Red Hat и IBM) для маршрутизации на уровне флота.

Среди фич сервинга — FP8 KV-кэши, чекпоинты FP8, MXFP4 и NVFP4. Спекулятивное декодирование включает MTP и EAGLE-3.

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

TorchTitan, референсная реализация распределённого обучения на основе PyTorch, теперь включает production-grade интеграцию с TorchTPU, построенную на torch.compile, SimpleFSDP от TorchTitan, специализированных Pallas-ядрах и оптимизациях на уровне графов.

VeRL получает TPU-бэкенд для обучения с подкреплением, Ray добавляет нативное TPU-планирование, а LiteRT (ранее TensorFlow Lite) присоединяется к экосистеме PyTorch для on-device инференса на телефонах, в браузерах и встраиваемом железе.

Как начать работу с TorchTPU

Проект устанавливается через pip командой pip install torch_tpu==0.2.0 --extra-index-url https://download.pytorch.org/whl/cpu, которая подтягивает CPU-сборку torch вместо стандартной CUDA-версии.

# Установка
# pip install torch_tpu==0.2.0 --extra-index-url https://download.pytorch.org/whl/cpu

import torch
import torch_tpu

# Минимальное изменение кода: меняем cuda на tpu
device = torch.device("tpu")

model = MyModel().to(device)
input = torch.randn(8, 512).to(device)
output = model(input)

На практике разработчик может перенести нагрузку, изменив инициализацию устройства с .to("cuda") на .to("tpu").

Для production-использования с torch.compile:

import torch
import torch_tpu

device = torch.device("tpu")
model = MyModel().to(device)

# Компиляция для максимальной производительности
compiled_model = torch.compile(model)

# Дальше работает как обычный PyTorch-код
output = compiled_model(input_tensor.to(device))

Разработчики могут отлаживать модели в eager-режиме, а затем переключаться на production-производительность через torch.compile для оптимизации утилизации TPU и снижения memory overhead.

⚠ Важно знать
PyTorch на TPU по умолчанию использует «medium» точность — умножение в BF16 с аккумуляцией в FP32, тогда как PyTorch на CUDA по умолчанию использует «highest» (чистый FP32). Это может влиять на числовые результаты. Проверяйте совпадение результатов между GPU и TPU, особенно для задач, чувствительных к точности вычислений.

Поддерживаемые TPU-поколения: v5e, v5p, v6e и v7. Для распределённого обучения доступны DDP, FSDP2, SimpleFSDP, DTensor и multi-slice конфигурации до 8K чипов.

Стратегический контекст: борьба за AI-инфраструктуру

Выход TorchTPU — это не просто технический релиз, это стратегический манёвр.

Обе компании — Google и Amazon — нацелены на одно и то же: снижение стоимости переноса существующего PyTorch-кода с чипов NVIDIA, под которые написано большинство AI-приложений по умолчанию.

Отделяя код модели от конкретного железа, TorchTPU стремится снизить friction при перемещении нагрузок между различными вычислительными средами. Это прямой удар по концепции CUDA Lock-in — многолетнему конкурентному преимуществу NVIDIA.

Стратегия open-source Google с TorchTPU разработана так, чтобы сделать TPU надёжным вторым или третьим ускорителем в любом серьёзном AI-инфраструктурном портфолио.

Релиз охватывает четыре GitHub-репозитория под лицензией Apache 2.0 — разрешительной лицензией, допускающей коммерческое использование.

📝 Пример реального применения
Джана ван Грёнен, директор по инженерии PyTorch в Meta, заявила, что TorchTPU позволяет Meta гибко сотрудничать при разработке железа с командой core machine learning systems, поддерживая крупные LLM, рекламные и рекомендательные системы. Это означает, что проект уже зарекомендовал себя в условиях реальной production-нагрузки одной из крупнейших AI-компаний мира.

Заключение

TorchTPU — это технологический сдвиг, который меняет правила игры в AI-инфраструктуре. Впервые за всю историю Google TPU разработчики на PyTorch получают нативный, полноценный и open-source путь к этому железу — без переписывания кода, без изучения JAX, без болезненной миграции.

Ключевые выводы:

  • Минимальные усилия по миграции: меняем .to("cuda") на .to("tpu") — и готово.
  • Зрелый проект: 98% совместимость с Hugging Face моделями, проверено на 8 192 TPU-чипах.
  • Богатая экосистема: vLLM, SGLang, TorchTitan, Ray, Hugging Face из коробки.
  • Открытый код: Apache 2.0, активное сообщество, прозрачная разработка.
  • Стратегическая ставка: первый реальный инструмент для снижения зависимости от NVIDIA CUDA.

Для команд, уже работающих на PyTorch, TorchTPU открывает новое измерение гибкости: возможность выбирать железо под задачу и бюджет, а не под экосистему. И это, пожалуй, самое ценное, что принёс этот релиз.