Google открыла TorchTPU для нативного запуска PyTorch на своих ускорителях

Google опубликовала TorchTPU — разработанный совместно с Meta бэкенд PyTorch для тензорных процессоров Google. Он позволяет переносить обучение и инференс моделей на TPU с минимальными изменениями кода и расширяет выбор аппаратных платформ за пределами GPU Nvidia.

Google открыла исходный код TorchTPU — нативного бэкенда PyTorch для тензорных процессоров TPU. Проект, созданный совместно с командой PyTorch в Meta, позволяет запускать существующие нагрузки на TPU с минимальными изменениями кода, использовать eager mode для отладки и torch.compile для оптимизированного исполнения.

Репозиторий TorchTPU опубликован по лицензии Apache 2.0. В него входят реализации операторов ATen для TPU, кэш компиляции XLA, поддержка распределённого обучения и инструменты для инференса, отладки и тестирования моделей.

TorchTPU интегрирован с Hugging Face Transformers и Diffusers, vLLM, SGLang, llm-d, TorchTitan, VeRL и Ray. Google также представила предварительную версию TPU Agent Suite для автоматизации переноса и настройки моделей.

Почти одновременно Amazon открыла TorchNeuron для запуска PyTorch-кода на ускорителях AWS Trainium. Два крупнейших облачных провайдера тем самым предлагают разработчикам собственные аппаратные альтернативы экосистеме Nvidia.