Hugging Face и AWS показали, какие модели работают на чипах Trainium

Hugging Face и AWS открыли справочник, где видно, какие популярные модели из каталога Hugging Face запускаются на чипах Amazon Trainium¹ без правок кода.

Детали: у каждой модели указано, на сколько ядер Trainium её разбили, сколько длился запуск вместе с компиляцией и сколько памяти он занял. Там же лежит скрипт, которым модель запускали, а все результаты, включая тексты ошибок, собраны в открытый набор данных.

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

Но: пока в справочнике только модели библиотеки Transformers², которые пишут текст, в том числе по картинкам, и только те, что помещаются на четыре чипа Trainium2. Скорость работы и длинные ответы тест не оценивает.

Модели для справочника запускали через TorchNeuron — открытую программу, которую Amazon выпустила на этой неделе. Она позволяет почти без правок запускать на Trainium программы, написанные для PyTorch — популярной среды для создания ИИ-моделей. С TorchNeuron библиотека Transformers от Hugging Face выполняет на Trainium тот же код моделей, что и на видеокартах, поэтому переписывать его не приходится.