Google открыла исходники TorchTPU: нативный бэкенд PyTorch для облачных TPU

На конференции PyTorch Conference 2026 Google совместно с командой PyTorch из Meta анонсировала и открыла исходники TorchTPU. Это нативный бэкенд, который подключается через интерфейс PrivateUse1 и позволяет переносить обычный PyTorch-код на TPU простым указанием torch.device("tpu"). Тензоры остаются стандартными, без обёрток и подклассов, а сам бэкенд работает в одном процессе на устройство, сохраняя представления тензоров и in-place операции.

TorchTPU приходит на смену PyTorch/XLA, который требовал дополнительной разметки для компиляции и поддерживал только SPMD-код. Новый бэкенд умеет обрабатывать и MPMD-исполнение, предлагает три eager-режима: Strict (по умолчанию), Debug с проверкой границ и DeferAndFuse, который сливает операции и сокращает трафик HBM — в некоторых случаях быстрее Strict на 50%. Для максимальной производительности используется torch.compile(backend="tpu"), который через TorchDynamo, StableHLO и OpenXLA генерирует машинный код TPU и кэширует исполняемые файлы между процессами и перезапусками.

Google отчитывается, что в обучении Qwen 3 8B с TorchTitan на Ironwood достигнуто 52% MFU, а масштабирование почти линейное вплоть до 8192 чипов. В тесте дообучения время шага упало с 11,49 секунд до 2,78 секунд: DeferAndFuse дал ускорение 1,15x, а полный torch.compile — 3,80x. Серверная часть поддерживает FP8, MXFP4 и NVFP4, а около 5000 моделей Hugging Face непрерывно тестируются в eager и compiled режимах.

Релиз включает четыре открытых репозитория: TorchTPU, vLLM-TorchTPU, SGLang-TorchTPU и TorchTitan. Версия 0.2.0 уже доступна на PyPI, лицензия Apache 2.0, требуется Python 3.11+. Это шаг к тому, чтобы TPU стали полноценным устройством наравне с GPU и CPU в экосистеме PyTorch.