Google открыла исходный код TorchTPU — программного стека для обучения и запуска моделей PyTorch на тензорных процессорах TPU. Разработка велась совместно с командой PyTorch, чтобы сделать TPU полноценной альтернативой CPU и GPU. Разработчики могут переносить существующие проекты на новую платформу с минимальными изменениями кода, используя привычные инструменты для отладки и распределенных вычислений.
Стек поддерживает режим пошагового выполнения и функцию torch.compile для оптимизации производительности. В экосистему Hugging Face добавлена поддержка TPU: в Transformers 5 запуск моделей настраивается через параметр 'device_map="tpu"', а библиотека Diffusers теперь позволяет распределять вычисления между несколькими ускорителями. Для контроля совместимости компания внедрила регулярные автоматические тесты.
Для работы с большими языковыми моделями подготовлены бэкенды vLLM и SGLang. SGLang применяет повторное использование кэша ключей и значений внимания для обработки запросов, а интеграция с системой llm-d позволяет эффективно распределять нагрузку между узлами с учетом сохраненного контекста. Эти решения упрощают развертывание моделей на TPU, разделяя этапы обработки запроса и генерации ответа.
TorchTPU интегрирован с TorchTitan для распределенного обучения нейросетей и поддерживает вычислительные ядра Pallas. В тестах с моделью Qwen 3 8B система достигла утилизации ресурсов до 52% и масштабировалась до 8192 ускорителей Ironwood. Хотя результаты зависят от конкретной конфигурации, открытие кода снижает порог входа для использования TPU, позволяя разработчикам гибко выбирать аппаратную платформу для своих задач.
