Компании Google и Meta представили TorchTPU — бэкенд с открытым исходным кодом для запуска моделей PyTorch на облачных TPU. Инструмент интегрируется через интерфейс PrivateUse1, позволяя работать со стандартными тензорами PyTorch без использования подклассов или оберток. Система поддерживает выполнение одного процесса на устройство, сохраняя при этом представления тензоров и мутации на месте.
В отличие от предшественника PyTorch/XLA, который использовал ленивые тензорные графы и требовал ручной разметки областей компиляции, TorchTPU поддерживает выполнение MPMD. Бэкенд использует torch.compile с аргументом tpu для захвата графов через TorchDynamo. Код транслируется через StableHLO в OpenXLA для генерации машинного кода TPU, а скомпилированные исполняемые файлы кэшируются между процессами и хостами.
TorchTPU позволяет вызывать пользовательские ядра Pallas, а также готовые ядра Tokamax и Helion как стандартные функции PyTorch. Для распределенного обучения предусмотрена группа процессов tpu_dist. Реализована поддержка DDP, FSDP2 и DTensor, а также топологически-зависимая сетка для межсоединений TPU.
Библиотека доступна для установки через pip, текущая версия — 0.2.0. При установке требуется индекс CPU-сборок PyTorch, чтобы исключить загрузку CUDA-версий. Подробные руководства по работе с инструментом опубликованы на официальном ресурсе проекта в ReadTheDocs.
