Работает ли ваша модель на AWS Trainium? Проверьте

Опубликовано 23 октября 2026 г.

В предыдущей публикации мы представили TorchNeuron — нативный бэкенд PyTorch для AWS Trainium (пакет torch-neuronx). С ним Transformers запускает собственный код моделей на Trainium так же, как и на любом другом устройстве, поэтому на нём может работать любая модель Transformers из Hub. Но какие именно?

Space поддержки моделей torch-neuronx отвечает на этот вопрос для самых популярных моделей в Hub. Он сообщает, запускается ли модель без изменений на Trainium, и предоставляет скрипт для её самостоятельного запуска.

Найдите модель

Space открывается в представлении Selected: это наша подборка моделей, выбранных среди лидирующих моделей в Hub по числу загрузок и популярности. В каждой строке указаны результат, количество ядер Trainium, на которые распределена модель, время выполнения (включая компиляцию) и пиковый объём памяти устройства.

Чтобы найти конкретную модель, переключитесь на представление Trending или Downloads и введите часть её названия. Эти представления следуют собственным рейтингам Hub, по одной задаче за раз, независимо от того, тестировалась модель или нет. Если модель намеренно не была включена, Space сообщает причину (чаще всего ей просто не уделили приоритетного внимания, а иногда она действительно несовместима). Также можно выбрать версии компилятора Neuron (neuronx-cc) и Transformers.

На момент написания статьи тестируются только модели Transformers, а наша подборка ограничена моделями для задач text-generation и image-text-to-text. Мы планируем расширить её для diffusers, sentence transformers и других задач transformers.

Запустите модель самостоятельно

Выберите модель, чтобы узнать, как её запустить. Space указывает стек Neuron, на котором тестировалась модель, пакеты, которые нужно установить поверх него, и запустивший её скрипт run.py. Скрипт следует рецепту из анонса TorchNeuron: он загружает модель с помощью from_pretrained, при необходимости распределяя её по шардам, а затем генерирует текст со статическим кэшем и torch.compile(backend="neuron").

На экземпляре Trainium2 (trn2) с этим стеком, например в контейнере AWS Neuron Deep Learning, установите пакеты, а затем выполните:

python3 run.py

torchrun не нужен: скрипт самостоятельно запускает собственные ранги torch.distributed, когда модель распределена по шардам.

При первом запуске начинается компиляция графов модели. Чтобы пропустить этот шаг, скопируйте графы, которые тот же компилятор создал во время тестирования, из общедоступного бакета кэша NEFF. Space предоставляет команду с нужной версией компилятора:

hf buckets sync hf://buckets/aws-neuron/torch-neuronx-neff-cache/neuronxcc-<version> /tmp/neff_cache

Что означает «работает»

Каждый тест загружает модель на экземпляре trn2, распределяя её по минимальному числу ядер Trainium, способных вместить её веса. Большинство моделей распределяется с помощью параллелизма по тензорам. Для моделей со смешанными экспертами, у которых нет плана параллелизма по тензорам, используется только параллелизм экспертов. Затем тест генерирует 32 токена и проверяет их. На одном ядре токены сравниваются с токенами, полученными при запуске той же контрольной точки на CPU. Когда модель распределена по шардам, небольшая модель-судья проверяет, что текст написан на естественном английском языке.

На момент написания статьи тестируются только модели, которые помещаются на 4 устройствах Trainium2, то есть максимум на 16 ядрах. Результат PASS ничего не говорит о пропускной способности или более длинных генерациях.

Результаты — это датасет

Всё, что показывает Space, берётся из общедоступного датасета aws-neuron/torch-neuronx-model-support, содержащего по одному файлу Parquet для каждой пары версий neuronx-cc и библиотек.

В каждой строке содержатся результат для модели, ошибка при сбое, версии всего стека и запустивший тест скрипт run.py. Вы можете напрямую выполнить запрос:

import pandas as pd

df = pd.read_parquet("hf://datasets/aws-neuron/torch-neuronx-model-support/data")
df[df.model_id == "Qwen/Qwen3-0.6B"][["neuronx_cc", "library", "library_version", "neuron_cores", "result"]]

В карточке датасета приведено точное значение каждого результата.

В следующий раз, когда вы задумаетесь, работает ли модель на Trainium, проверьте это.