Модели из библиотеки Transformers могут работать на чипах Trainium от AWS благодаря TorchNeuron — нативному бэкенду для PyTorch, библиотеки машинного обучения с открытым исходным кодом. Остаётся вопрос, какие именно модели реально работают. Space torch-neuronx model support, опубликованный Hugging Face, даёт ответ для самых популярных моделей на Hub.
Поиск модели
Space открывается видом Selected — подборкой моделей, выбранных среди самых скачиваемых и популярных на Hub. Каждая строка показывает результат теста, количество ядер Trainium, на которые распределена модель, общее время выполнения (включая компиляцию) и максимальную память, использованную на устройстве.
Чтобы найти конкретную модель, достаточно переключиться на виды Trending или Downloads и ввести часть её названия. Эти виды следуют рейтингу Hub по задачам, независимо от того, тестировалась модель или нет. Если модель была исключена, Space указывает причину: чаще всего она не была приоритетной, иногда — действительно несовместима. Пользователь также может выбрать версии компилятора Neuron (neuronx-cc) и Transformers для проверки.
Пока тестируются только модели Transformers, и выбор ограничен задачами генерации текста и «текст-изображение в текст». Команда планирует расширить инструмент на diffusers, sentence transformers и другие задачи Transformers.
Запуск модели самостоятельно
Выбрав модель, можно увидеть стек программного обеспечения Neuron, на котором она тестировалась, пакеты для установки и скрипт run.py, использованный при тесте. Этот скрипт загружает модель через from_pretrained, распределяя её по нескольким ядрам при необходимости, а затем генерирует текст со статическим кэшем и torch.compile(backend=«neuron»).
На экземпляре trn2 с таким же стеком, например, в AWS Neuron Deep Learning Container, достаточно установить пакеты и запустить python3 run.py. Torchrun не требуется — скрипт сам запускает процессы распределения, когда модель распределена. Первый запуск начинается с компиляции графов модели. Чтобы избежать этого, можно скопировать графы, созданные во время теста, из публичного бакета NEFF cache с помощью команды, предоставленной Space, которая учитывает версию компилятора.
Что означает «работает»
Каждый тест загружает модель на экземпляр trn2, распределяя её на минимальное количество ядер Trainium, способных вместить её веса. Большинство моделей используют тензорный параллелизм, который распределяет вычисления между несколькими ядрами. Модели со смешанными экспертами, без плана тензорного параллелизма, используют только параллелизм экспертов. Затем тест генерирует 32 токена (единицы текста) и проверяет их. На одном ядре эти токены сравниваются с токенами выполнения той же контрольной точки на процессоре. Когда модель распределена, небольшой модельный судья проверяет, является ли текст беглым английским.
Тестируются только модели, помещающиеся на четырёх устройствах Trainium2, то есть максимум на шестнадцати ядрах. Результат PASS ничего не говорит о пропускной способности или более длинных генерациях.
Все результаты взяты из публичного набора данных aws-neuron/torch-neuronx-model-support, который содержит файл Parquet для каждой пары версий neuronx-cc и библиотеки. Каждая строка содержит результат модели, возможную ошибку, версии стека и выполненный run.py. Набор данных можно запрашивать напрямую с помощью pandas, например, для модели Qwen/Qwen3-0.6B. Описание набора данных уточняет точное значение каждого результата.
