google-pytorch/torch_tpu

Что это?

Ключевые возможности

• 01Собственные реализации ядер ATen для TPU.
• 02Кэш компиляции для повторного использования скомпилированных вычислений XLA.
• 03Параллелизм по данным и тензорам для распределённого обучения.
• 04Интеграция с torch.compile() для выполнения в графовом режиме.
• 05Примеры кода для моделей, в том числе Llama, Qwen, ResNet и UNet.
• 06Утилиты для отладки и бенчмаркинга.

Сценарии использования

• Запуск существующего кода обучения PyTorch на TPU в Google Cloud.
• Распределённое обучение больших моделей на нескольких устройствах TPU.
• Запуск инференса моделей PyTorch на оборудовании TPU.
• Изучение примеров реализации распространённых моделей для TPU.

Быстрый старт

1. 1

   Создайте виртуальное окружение Python

   python -m venv .venv
2. 2

   Установите последнюю публичную версию из PyPI, используя индекс PyTorch для CPU

   pip install torch_tpu==0.2.0 --extra-index-url https://download.pytorch.org/whl/cpu

Что важно знать

• !Проект предназначен для оборудования Google TPU, поэтому полезен только при наличии доступа к TPU, например к TPU VM.
• !Требуется Python 3.11 или новее.
• !Для сборки из исходного кода нужен Bazel, а для полной сборки wheel — учётные данные Remote Build Execution или флаг no_rbe.
• !В метаданных лицензия указана как NOASSERTION, хотя в README указана Apache 2.0.

Темы

Частые вопросы

Что такое torch_tpu?

Бэкенд PyTorch, позволяющий запускать модели и операции PyTorch на Tensor Processing Units (TPU) от Google. TorchTPU предоставляет бэкенд PyTorch для TPU от Google, позволяя запускать существующий код PyTorch на устройствах TPU. В проект входят собственные ядра ATen для ряда операций PyTorch, а также кэш компиляции, который повторно использует скомпилированные вычисления XLA и ускоряет повторные запуски.

torch_tpu можно использовать бесплатно? Какая у него лицензия?

torch_tpu распространяется с открытым исходным кодом под лицензией Other. Перед использованием в коммерческих продуктах проверьте условия лицензии.

torch_tpu всё ещё поддерживается?

Активно поддерживается. Последний коммит в torch_tpu — 22 окт. 2026 г..

Для чего можно использовать torch_tpu?

Запуск существующего кода обучения PyTorch на TPU в Google Cloud. Распределённое обучение больших моделей на нескольких устройствах TPU. Запуск инференса моделей PyTorch на оборудовании TPU. Изучение примеров реализации распространённых моделей для TPU.

Как установить torch_tpu?

Создайте виртуальное окружение Python: python -m venv .venv Установите последнюю публичную версию из PyPI, используя индекс PyTorch для CPU: pip install torch_tpu==0.2.0 --extra-index-url https://download.pytorch.org/whl/cpu

Где скачать torch_tpu?

Исходный код и инструкции по установке — на GitHub: https://github.com/google-pytorch/torch_tpu