Чи працює ваша модель на AWS Trainium? Перевірте

Ваша модель працює на AWS Trainium? Перевірте

Опубліковано 23 жовтня 2026 року

У попередній публікації ми представили TorchNeuron — нативний бекенд PyTorch для AWS Trainium (пакет torch-neuronx). Завдяки йому Transformers запускає власний код моделей на Trainium так само, як і на будь-якому іншому пристрої, тож будь-яка модель Transformers на Hub може працювати. Але які саме?

Простір підтримки моделей torch-neuronx відповідає на це запитання для найпопулярніших моделей на Hub. Він показує, чи працює модель на Trainium без змін, і надає скрипт для її самостійного запуску.

Знайти модель

Простір відкривається у режимі Вибрані: це наша власна добірка моделей, вибраних серед найпопулярніших моделей на Hub за кількістю завантажень і трендовим показником. У кожному рядку наведено результат, кількість ядер Trainium, на які розподілено модель, час виконання (включно з компіляцією) і піковий обсяг пам’яті пристрою.

Щоб знайти певну модель, перемкніться в режим Трендові або Завантаження та введіть частину її назви. Ці режими використовують власні рейтинги Hub, по одному завданню за раз, незалежно від того, чи тестувалася модель. Якщо модель навмисно не було включено, простір пояснює чому (здебільшого її не було пріоритизовано, а іноді вона справді несумісна). Також можна вибрати версії компілятора Neuron (neuronx-cc) і Transformers для перегляду.

На момент написання статті тестуються лише моделі Transformers, а наша добірка обмежена моделями для завдань text-generation і image-text-to-text. Ми плануємо розширити її на diffusers, sentence transformers та інші завдання transformers.

Запустіть самостійно

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

На екземплярі Trainium2 (trn2) із цим стеком, наприклад у контейнері AWS Neuron для глибокого навчання, встановіть пакети, а потім запустіть:

python3 run.py

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

Під час першого запуску починається компіляція графів моделі. Щоб пропустити цей крок, скопіюйте графи, які той самий компілятор створив під час тестування, із загальнодоступного бакета кешу NEFF. Простір надає команду з правильною версією компілятора:

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

Що означає «працює»

Під час кожного тесту модель завантажується на екземплярі trn2 і розподіляється між мінімальною кількістю ядер Trainium, здатною вмістити її ваги. Більшість моделей розподіляється за допомогою паралелізму тензорів. Для моделей суміші експертів без плану паралелізму тензорів використовується лише паралелізм експертів. Потім тест генерує 32 токени та перевіряє їх. На одному ядрі токени порівнюються з токенами під час запуску тієї самої контрольної точки на CPU. Коли модель розподілена між пристроями, невелика модель-суддя перевіряє, чи є текст зв’язною англійською мовою.

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

Результати — це набір даних

Усе, що показує простір, надходить із загальнодоступного набору даних 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, перевірте це.