Amazon выпустила TorchNeuron для запуска кода PyTorch на чипах Trainium

Amazon выпустила TorchNeuron — программу с открытым кодом, с которой написанные на PyTorch¹ модели запускаются на чипах Trainium² почти без правок. В проверке помогали ИИ-агенты: они перебрали сотни тысяч моделей с Hugging Face — площадки, куда выкладывают открытые модели.

Почему это важно: большинство разработчиков пишут модели на PyTorch и запускают их на видеокартах. Чем меньше кода приходится переделывать, тем проще клиентам облака AWS попробовать вместо видеокарт чипы Amazon.

Повод: на этой неделе Amazon объявила о выходе TorchNeuron и в блоге Amazon Science рассказала, как программа устроена и как её проверяли.

Расклад: на TorchNeuron опирается и Hugging Face — компания, которая ведёт эту площадку и выпускает библиотеки для работы с моделями. 19 октября Hugging Face сообщила, что модели из её библиотек запускаются и дообучаются на Trainium тем же кодом, что и на видеокартах, без отдельной библиотеки-переходника.

Как это работает: у PyTorch два режима. В режиме немедленного выполнения операции идут по одной в порядке программы — так проще искать ошибки. В режиме компиляции компилятор заранее переводит блок кода в команды чипа и оптимизирует его: так быстрее, но подготовка занимает время.

Детали: TorchNeuron добавляет промежуточный вариант и старается не компилировать лишнего:

• Новый адаптивный режим, который включают по желанию, собирает операции в небольшие группы и компилирует их, пока чип выполняет уже готовые.
• Для многих моделей время компиляции так удаётся скрыть целиком, а разработчик по-прежнему видит, какая строка кода выполняется на чипе.
• Каждый скомпилированный фрагмент хранится в памяти и на диске, поэтому после перезапуска его не нужно собирать заново.
• Если модель разнесена по тысячам чипов, код компилирует одна машина, а остальные берут готовый.

В цифрах: в тесте с большой языковой моделью, разнесённой по многим чипам, эти приёмы сократили время до первого шага обучения примерно на треть.

Интрига: во многом новизна TorchNeuron, по словам Amazon, — в том, как программу проверяли:

• Одни ИИ-агенты прошлись по репозиториям Hugging Face, отсеяли модели с неработающими настройками и нашли сотни тысяч моделей с рабочим кодом.
• Другой агент разложил их по устройству: какие данные модель принимает, какой в ней механизм внимания³ и в каком виде она выдаёт результат.
• Оказалось, что все архитектурные особенности, которые TorchNeuron, вероятно, встретит на практике, покрывает 841 модель.
• Новую версию программы сначала гоняют по быстрым тестам из реестра десятков тысяч операций, затем по 841 модели и только потом по всему набору.

Но: скорость Trainium и видеокарт на одних и тех же моделях Amazon не сравнивает, а время до первого шага обучения сократилось на треть в одном тесте компании.

Общая картина: Amazon может упростить переезд моделей между чипами и с помощью покупки. 12 октября WSJ писала, что компания обсуждает покупку стартапа Decart за 7 млрд долларов — он помогает переносить ИИ-модели между чипами.