Google сделала важный шаг для разработчиков, работающих с искусственным интеллектом, представив проект TorchTPU. Это решение выступает в роли нативного бэкенда для популярного фреймворка PyTorch, обеспечивая прямую поддержку тензорных процессоров Google (TPU). Исходный код проекта уже доступен на платформе GitHub под лицензией Apache 2.0.
Основная задача TorchTPU — упростить интеграцию PyTorch с аппаратным обеспечением Google. Инструментарий включает собственные ядра ATen, специально адаптированные для работы на TPU, а также систему кеширования компиляции для оптимизации выполнения вычислений. Разработчики получают возможность использовать преимущества аппаратного ускорения Google как для обучения моделей, так и для их инференса.
Ключевые возможности и поддержка
Проект предлагает широкий набор инструментов для работы с современными архитектурами нейросетей. Среди основных функций можно выделить:
- Поддержку распределенного обучения (Data Parallelism и Tensor Parallelism).
- Интеграцию с torch.compile() для выполнения в графовом режиме.
- Готовые примеры работы с популярными моделями, включая Llama, Qwen, ResNet, DLRM, minGPT и UNet.
- Утилиты для отладки и бенчмаркинга производительности.
Для управления зависимостями и сборки проекта используется система Bazel. Пользователи могут установить текущую версию библиотеки через менеджер пакетов pip или собрать её из исходного кода. TorchTPU поддерживает Python версии 3.11 и выше, следуя политике совместимости PyTorch.
Для тех, кто хочет глубже изучить возможности платформы, Google подготовила подробное руководство пользователя и интерактивные ноутбуки. Они базируются на среде Marimo и позволяют тестировать код непосредственно в браузере. Рекомендуется запускать их на виртуальных машинах с TPU, например, на экземплярах GCE v6e. Для комфортной разработки C++ компонентов авторы проекта рекомендуют использовать VSCode в связке с расширением clangd, что обеспечивает полноценную поддержку навигации по коду и автодополнения.
Источник: github.com