Amazon и Google облегчили побег от Nvidia: запускать модели под PyTorch на ИИ-чипах Trainium и TPU стало проще

Amazon и Google представили на мероприятии PyTorch Conference открытые средства для запуска оптимизированного для Nvidia фреймворка PyTorch на своих ускорителях искусственного интеллекта. Инструментарий Amazon TorchNeuron позволяет запускать проекты PyTorch на ускорителях Trainium; Google представила инструмент с открытым исходным кодом TorchTPU для работы PyTorch на ускорителях TPU.

Amazon TorchNeuron решил проблему совместимости PyTorch с чипами Trainium — ранее для работы фреймворка на этих ускорителях требовались значительные усилия. Теперь проекты запускаются на ускорителях Trainium2 и Trainium3 без внесения масштабных изменений в код — достаточно лишь указать параметр, какой ускоритель будет использоваться для выполнения задач. TorchNeuron включает набор готовых операторов языка NKI, созданного для ускорителей Trainium. За счёт специального компилятора язык позволяет оптимизировать модели на низком уровне, обеспечивая прирост скорости и сокращение потребления аппаратных ресурсов.

Помимо поддержки самого PyTorch, Amazon предложила совместимость и с другими открытыми инструментами, которые используются в связке с фреймворком. Одним из них является используемый в обучении моделей механизм TorchTitan, который автоматизирует часть задач по распределению процессов между ускорителями.

Разработкой инструмента TorchTPU компания Google занималась совместно с командой PyTorch в Meta✴. Этот инструмент позволяет запускать существующие рабочие нагрузки PyTorch на ускорителях TPU с минимальными изменениями в коде. Поддерживаются стандартные средства API PyTorch, предусмотрена расширенные инструменты диагностики, есть встроенные примитивы для распределённых вычислений — всё это обеспечивает высокую производительность и эффективность работы на старте, отметили в Google Cloud.

В Google также заручились поддержкой Hugging Face, благодаря чему модели с открытыми весами из библиотеки популярной ИИ-платформы поддерживают работу с TPU «из коробки» — достаточно указать соответствующий параметр. Отмечена поддержка востребованных в профессиональной среде средств разработки ИИ: vLLM, SGLang, llm-d, TorchTitan, Ray, TorchRec и LiteRT. Другими словами, ускорители Google TPU теперь включены в привычную для разработчиков среду PyTorch с поддержкой набора уже существующих инструментов — это упрощает переход на оборудование Google и не требует перестраивать всю программную систему.

Источники:

Материалы по теме

Материалы по теме