Смена парадигмы: от XLA к нативному PyTorch
20 октября 2026 года на PyTorch Conference 2026 Google анонсировал и открыл исходные коды TorchTPU. Это не новая модель, а программный стек (backend), позволяющий запускать модели PyTorch нативно на облачных TPU. Проект разработан совместно с командой Meta и лицензирован под Apache 2.0. Ключевое отличие от предшественника, PyTorch/XLA, заключается в переходе от ленивых вычислений (lazy tensors) к режиму eager execution через интерфейс PrivateUse1. Теперь TPU являются полноценным устройством PyTorch, доступным через стандартный вызов torch.device("tpu").
Производительность и масштабирование
Google сообщает о значительном росте эффективности. При обучении модели Qwen 3 8B с использованием фреймворка TorchTitan на чипах Ironwood (TPU7x) достигнута метрика 52% MFU (Model FLOPs Utilization). Система масштабируется линейно до 8 192 чипов. Для оптимизации скорости доступны три режима выполнения:
- Strict (по умолчанию): асинхронная диспетчеризация операторов со стандартной семантикой PyTorch.
- Debug: синхронное выполнение с проверкой границ тензоров.
- DeferAndFuse: слияние операторов, снижающее трафик HBM; ускорение до 50%+ по сравнению со Strict.
Для продакшена рекомендуется использовать torch.compile(backend="tpu"), который транслирует граф через StableHLO в OpenXLA. В тесте тонкой настройки (fine-tuning) время шага сократилось с 11.49s до 2.78s (ускорение в 3.8 раза).
Сравнение с альтернативами
TorchTPU позиционируется как прямая замена PyTorch/XLA. Ниже приведено сравнение ключевых характеристик:
| Характеристика | TorchTPU (Google + Meta) | PyTorch/XLA (Legacy) | JAX on TPU |
|---|---|---|---|
| Механизм выполнения | Eager (PrivateUse1) | Lazy-tensor graph | Не PyTorch |
| Поддержка распределения | DDP, FSDP2, DTensor, MPMD | Только чистый SPMD | Не раскрыто |
| Кастомные ядра | Pallas, Tokamax, Helion | Не поддерживаются | Pallas |
| Статус | Open Source (v0.2.0) | Заменяется TorchTPU | Поддерживается |
Экосистема и интеграции
Пакет torch_tpu==0.2.0 доступен в PyPI для Python 3.11+. Он обеспечивает совместимость с major-игроками индустрии:
- Hugging Face: Transformers v5 и Diffusers поддерживают
device_map="tpu". - Inference: Интеграции vLLM-TorchTPU (от Inferact) и SGLang-TorchTPU (от RadixArk) поддерживают FP8, MXFP4, NVFP4 и спекулятивное декодирование (MTP, EAGLE-3).
- Обучение: TorchTitan включает рецепты для Llama 3, Qwen 3 (включая MoE 235B) и FLUX.1.
Google тестирует около 5 000 моделей с Hugging Face в eager и compiled режимах, что делает TorchTPU универсальным решением для разработчиков, переходящих с GPU на TPU.
Источник: MarkTechPost ↗