Amazon анонсировала TorchNeuron — инструментарий с открытым исходным кодом, позволяющий запускать рабочие нагрузки PyTorch на ИИ-ускорителях AWS Trainium. До настоящего момента запуск на чипах Trainium самого популярного в отрасли фреймворка PyTorch для разработки ИИ-систем требовал значительного объёма ручной работы — затраты на разработку увеличивались, создавался барьер для пользователей.
TorchNeuron позволяет запускать проекты PyTorch на чипах Trainium, не внося масштабных изменений в код. Разработчикам, подчеркнули в Amazon, достаточно лишь задать параметр, какой именно чип следует использовать для выполнения задач. Поддерживаются Trainium2 и Trainium3 — две новейшие модели в линейке этих ускорителей; они подходят как для обучения, так и для инференса моделей.
Источник изображения: aboutamazon.com
Для серии Trainium был разработан собственный язык программирования NKI; благодаря наличию специального компилятора этот язык позволяет производить низкоуровневую оптимизацию моделей, благодаря чему снижается потребление аппаратных ресурсов и повышается скорость работы ИИ. В состав TorchNeuron входит набор готовых операторов NKI — программных компонентов, которые можно использовать при создании ИИ-моделей. Одни из них упрощают реализацию MoE-архитектуры, другие помогают сократить объём используемой памяти. Ускоряется выполнение и сопутствующих задач, например, квантования.
TorchNeuron поддерживает не только PyTorch, но и другие открытые инструменты, которые разработчики часто используют в связке с этим фреймворком. Среди них есть TorchTitan — популярный механизм для обучения ИИ-моделей. Он автоматизирует часть задач, связанных с распределением процессов между ускорителями; метод полностью сегментированного параллелизма данных разделяет не только обучаемую модель, но также состояния оптимизатора и градиенты. В состав TorchNeuron вошли функции, позволяющие дополнительно оптимизировать рабочий процесс обучения с использованием ускорителей Trainium.
Источник:
