ИИтог

Google опубликовала TorchTPU: бэкенд PyTorch для TPU с открытым кодом

инфраструктураРедакция ИИтогHacker News🔥80Рейтинг ИИтог80Влияние70Новизна65Релевантность100Неожиданность50Участники85Достоверность85
Google опубликовала TorchTPU: бэкенд PyTorch для TPU с открытым кодом

На GitHub опубликован репозиторий google-pytorch/torch_tpu с проектом TorchTPU. По описанию в README, это бэкенд PyTorch для тензорных процессоров Google (TPU), который позволяет запускать модели и операции PyTorch на устройствах TPU. Проект распространяется по лицензии Apache 2.0.

В составе проекта: реализация ядер ATen (базовых операций PyTorch) для TPU, кеш компиляции для оптимизации выполнения вычислений XLA, поддержка распределённого обучения (параллелизм по данным и тензорный параллелизм), интеграция с torch.compile() для графового режима, а также утилиты для отладки и измерения производительности. В качестве примеров приведены модели Llama, Qwen, ResNet, DLRM, minGPT и UNet.

Установить последний публичный релиз можно из PyPI: версия 0.2.0, команда pip install torch_tpu==0.2.0 с дополнительным индексом https://download.pytorch.org/whl/cpu. Без этого индекса установится CUDA-версия torch. Сейчас поддерживается Python 3.11 и новее; политика поддержки версий Python совпадает с матрицей совместимости релизов PyTorch.

Количественных замеров в README нет, как нет и сравнения с другими способами запускать PyTorch на TPU. Документация (руководство пользователя) размещена на torchtpu.readthedocs.io; интерактивные ноутбуки на Marimo (похожем на Jupyter) лежат в папке docs и запускаются с TPU-машины, проще всего на экземпляре GCE v6e.

Собирать проект предлагается через Bazel (через Bazelisk). Колесо (wheel) собирается командой bazel build --config=wheel_common //ci/wheel:torch_tpu_wheel; флаг --config=wheel_common обязателен, иначе сборка падает. Если нет доступа к удалённой сборке (RBE), нужно добавить --config=no_rbe после wheel_common. Раздел с ключевыми Python-зависимостями пока помечен как TBD (будет определено позже). Для навигации по C++ коду рекомендуют VS Code с clangd; первая генерация compile_commands.json может занять более 30 минут. Вклад в проект принимается по правилам CONTRIBUTING.md, включая подписание CLA.

Ключевые факты

  • TorchTPU, бэкенд PyTorch для TPU Google, репозиторий google-pytorch/torch_tpu, лицензия Apache 2.0; последний публичный релиз 0.2.0 на PyPI.
  • В составе: ядра ATen для TPU, кеш компиляции XLA, параллелизм по данным и тензорный параллелизм, интеграция с torch.compile(), утилиты отладки и замеров.
  • Примеры: Llama, Qwen, ResNet, DLRM, minGPT, UNet; ноутбуки на Marimo запускаются с TPU-машины, проще всего на GCE v6e.
  • Поддерживается Python 3.11+; без CPU-индекса установка подтянет CUDA-версию torch.
  • Сборка идёт через Bazel; раздел Python-зависимостей в README пока TBD, бенчмарков нет.

Почему это важно

Проект предлагает запускать привычный код PyTorch на TPU через отдельный бэкенд с ядрами ATen, кешем компиляции XLA и интеграцией с torch.compile(). Репозиторий открыт под Apache 2.0, а в примерах, популярные архитектуры, включая Llama и Qwen. Как TorchTPU соотносится с уже существующими способами запуска PyTorch на TPU, README не объясняет.

Кому это важно

Командам, которые обучают или запускают модели на TPU Google и хотят писать на PyTorch. Также разработчикам, интересующимся внутренним устройством бэкендов: репозиторий открыт, принимает вклад по CONTRIBUTING.md (с подписанием CLA), а сборка описана достаточно подробно.

Как это применить

Создать виртуальное окружение Python 3.11+ и выполнить pip install torch_tpu==0.2.0 --extra-index-url https://download.pytorch.org/whl/cpu; без индекса поставится CUDA-версия torch. Интерактивные ноутбуки запускаются командой marimo edit file.py с TPU-машины, проще всего на экземпляре GCE v6e. Для самостоятельной сборки нужен Bazel (через Bazelisk): bazel build --config=wheel_common //ci/wheel:torch_tpu_wheel, а без доступа к RBE, с добавлением --config=no_rbe после wheel_common. Лицензия, Apache 2.0.

Можно ли доверять

Источник, только README репозитория, то есть собственное описание проекта. В нём нет бенчмарков, сравнений с другими подходами, имён авторов и заявлений о готовности к промышленному применению, поэтому о скорости и зрелости судить по нему нельзя. Версия 0.2.0 и пометка TBD в разделе зависимостей говорят о том, что проект на ранней стадии.

Риски и подводные камни

Раздел ключевых Python-зависимостей в README помечен как TBD. Сборка колеса требует обязательного флага --config=wheel_common, иначе она завершится ошибкой; порядок флагов тоже важен (no_rbe после wheel_common). Забыв CPU-индекс при установке, можно получить CUDA-версию torch. Первая полная сборка при настройке clangd может занять более 30 минут. Документация для запуска ноутбуков предполагает доступ к TPU-машине.