ТехноСводка

ИИ

Google открыла код TorchTPU для запуска моделей PyTorch на TPU

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

Разработка создана вместе с командой PyTorch. Существующие программы можно переносить с минимальными изменениями кода, сохраняя привычные средства отладки, компиляции (torch.compile) и распределённых вычислений. В Transformers 5 запуск на TPU задаётся параметром device_map="tpu", библиотека Diffusers поддерживает распределение вычислений между ускорителями. Для vLLM и SGLang подготовлены бэкенды, ведётся интеграция с llm-d.

В связке с TorchTitan стек поддерживает распределённое обучение и ядра Pallas. По данным Google, при обучении Qwen 3 8B утилизация вычислений достигла до 52%, а масштабирование дошло до 8192 TPU Ironwood. Эти цифры относятся к конкретной конфигурации и не гарантируют такого же результата для других моделей.

Открытый код не обещает одинаковой производительности на разных ускорителях, но снижает барьер для переноса проектов на TPU и даёт разработчикам больше свободы в выборе оборудования.

  • До 52% утилизации на Qwen 3 8B
  • Масштабирование до 8192 TPU Ironwood

Источники

Текст подготовлен автоматически по материалам источников и может содержать неточности.