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

Google опубликовала исходный код проекта torch_tpu — бэкенда для PyTorch, который даёт возможность исполнять модели и операции этого фреймворка непосредственно на тензорных процессорах (TPU). До сих пор связка PyTorch с фирменными ускорителями Google требовала обходных путей; теперь предлагается более прямой маршрут.

Что внутри

Проект включает собственную реализацию ATen-ядер для различных операций PyTorch, кэш компиляции для ускорения исполнения XLA-вычислений, а также интеграцию с torch.compile() для работы в графовом режиме. Поддерживается распределённое обучение — как параллелизм по данным, так и параллелизм по тензорам.

В репозитории собраны примеры запуска популярных моделей: Llama, Qwen, ResNet, DLRM, minGPT и UNet. Отдельно упомянуты утилиты для отладки и замера производительности.

Установка и сборка

Текущая публичная версия на PyPI — 0.2.0; поддерживается Python 3.11 и выше, политика версий совпадает с матрицей совместимости самого PyTorch. Пакет ставится обычным pip, причём документация настойчиво советует подключать дополнительный индекс со сборкой torch под CPU, иначе по умолчанию подтянется вариант под CUDA.

Для сборки из исходников используется Bazel (точнее, обёртка Bazelisk). Здесь кроется немало тонкостей: сборка колеса требует конфигурации wheel_common, без которой процесс просто падает. Причина техническая — без этого флага каждая версия тащит собственную копию XLA/MLIR-бэкенда, и готовый пакет затем аварийно завершается прямо на import torch из-за дублирующихся статических регистраций. Документация даже предупреждает о порядке передачи флагов: --config=no_rbe должен идти после --config=wheel_common, поскольку «побеждает» тот конфиг, что указан позже.

Документация

Пользовательское руководство доступно на readthedocs, а интерактивные блокноты выполнены на Marimo (аналог Jupyter) и лежат в каталоге docs. Запускать их удобнее всего прямо на TPU-инстансе в облаке Google.

Проект распространяется под лицензией Apache 2.0. На момент публикации новости обсуждение на Hacker News собрало 189 баллов и 47 комментариев.

Источник: Hacker News