Google открыла код TorchTPU для запуска моделей PyTorch на TPU
Изображение: OpenAI
Google открыла исходный код TorchTPU для запуска моделей и операций PyTorch на ускорителях TPU. По описанию проекта на GitHub, он поддерживает распределённое обучение и графовый режим выполнения через torch.compile().
В TorchTPU входят вычислительные ядра ATen для ряда операций PyTorch на TPU и кеш компиляции, который ускоряет выполнение вычислений XLA. Для распределённого обучения предусмотрены параллелизм по данным и тензорный параллелизм.
Разработчикам доступны инструменты отладки и измерения производительности, а также примеры распределённого обучения и инференса. В проекте есть примеры для следующих моделей: Llama, Qwen, ResNet, DLRM, minGPT и UNet.
На PyPI опубликована версия TorchTPU 0.2.0. Проект поддерживает Python 3.11 и новее. Код распространяется по лицензии Apache 2.0.
