Google открыла исходный код TorchTPU — программного слоя для обучения и запуска моделей PyTorch на собственных тензорных ускорителях TPU. В сообщении от 21 октября компания объявила о публичной доступности проекта, разработанного вместе с командой PyTorch из Meta. Код и установочный пакет уже доступны разработчикам.

PyTorch задаёт, как разработчик описывает нейросеть и её вычисления; переход на другое оборудование становится сложнее, если приходится менять этот код и привычные средства отладки. TorchTPU сохраняет стандартные интерфейсы PyTorch: по заявлению Google, существующие задачи можно переносить на TPU с минимальными изменениями. Во время отладки операции выполняются по отдельности, а для рабочего запуска вычисления можно объединить и оптимизировать компилятором.

Саму архитектуру Google представила ещё 7 апреля. Тогда публичный репозиторий и глубокая интеграция с инструментами обслуживания моделей оставались в планах. Теперь разработчики получили код, примеры распределённого обучения и запуска нейросетей — переход от описания технологии к доступному инструменту.

Практическое дополнение — выпущенный 20 октября vLLM TorchTPU, который переносит на TPU знакомый сервер языковых моделей vLLM. Команды Inferact и Google утверждают, что большинство поддерживаемых vLLM моделей можно обслуживать без изменения кода самих моделей, сохранив обычные интерфейсы сервера и управление запросами. Это уменьшает объём работы при выборе другого ускорителя.

Поддержка оборудования пока неодинакова: для vLLM TorchTPU рекомендован TPU v7x, или Ironwood, который проверяют при каждом изменении кода; v6e, v5e и v5p имеют экспериментальный статус. Открытый программный слой делает альтернативу GPU доступнее для проверки и доработки. Выигрыш в скорости и стоимости по-прежнему придётся устанавливать на конкретных моделях и нагрузках.