Amazon анонсировала TorchNeuron — инструментарий с открытым исходным кодом, позволяющий запускать рабочие нагрузки PyTorch на ИИ-ускорителях AWS Trainium. До настоящего момента запуск на чипах Trainium самого популярного в отрасли фреймворка PyTorch для разработки систем искусственного интеллекта требовал значительного объёма ручной работы — затраты на разработку увеличивались, создавался барьер для пользователей.
TorchNeuron позволяет запускать проекты PyTorch на чипах Trainium, не внося масштабных изменений в код. Разработчикам, подчеркнули в Amazon, достаточно лишь задать параметр, какой именно чип следует использовать для выполнения задач. Поддерживаются Trainium2 и Trainium3 — две новейшие модели в линейке этих ускорителей; они подходят как для обучения, так и для инференса моделей.
Более мощный из двух ускоритель Trainium3 оснащён восемью вычислительными ядрами, каждое из которых состоит из четырёх модулей, оптимизированных для определённых типов вычислений, необходимых ИИ-моделям при обработке данных. Объём встроенной памяти составляет 144 Гбайт. На сервере могут работать до 144 ускорителей Trainium3 — одна такая система способна выполнять до 362 квадриллионов операций в секунду при работе с данными в формате MXFP8, который широко применяется в моделях ИИ.
Источник изображения: aboutamazon.com
Для серии Trainium был разработан собственный язык программирования NKI; благодаря наличию специального компилятора этот язык позволяет производить низкоуровневую оптимизацию моделей, за счёт чего снижается потребление аппаратных ресурсов и повышается скорость работы ИИ.
В состав TorchNeuron входит набор готовых операторов NKI — программных компонентов, которые можно использовать при создании ИИ-моделей. Одни из них упрощают реализацию архитектуры Mixture-of-Experts, другие помогают сократить объём памяти, который занимают нейросети. Ускоряется выполнение и сопутствующих задач, например, квантования. TorchNeuron поддерживает не только PyTorch, но и другие открытые инструменты, которые разработчики часто используют в связке с этим фреймворком. Среди них поддерживается TorchTitan — популярный механизм для обучения ИИ-моделей. Он автоматизирует часть задач, связанных с распределением процессов между ускорителями; метод полностью сегментированного параллелизма данных разделяет не только обучаемую модель, но также состояния оптимизатора и градиенты. Оптимизатор — программа совершенствующая способности модели; градиенты — данные, которые играют в этом процессе важную роль. В состав TorchNeuron вошли функции, позволяющие дополнительно оптимизировать рабочий процесс обучения с использованием ускорителей Trainium.
Ранее о поддержке PyTorch на своих ускорителях заявили Intel, Meta✴ и корейский стартап HyperAccel.
Источник: servernews.ru