Hugging Face: Space проверяет, какие модели Hub заводятся на Trainium

С пакетом torch-neuronx Transformers исполняет свой же modeling-код на AWS Trainium, как на GPU — но по конкретному чекпоинту с Hub всё равно неясно, поедет ли он без правок. Hugging Face ответила это Space torch-neuronx model support и датасетом aws-neuron/torch-neuronx-model-support: статус прогона, метрики и скрипт, которым гоняли модель.

На 23 октября 2026 года в выборке только задачи text-generation и image-text-to-text; diffusers, sentence transformers и остальное обещают позже. Автор материала — David Corvoysier.

Как смотреть каталог

По умолчанию открывается Selected — подборка из топа Hub по скачиваниям и trending. Вкладки Trending и Downloads повторяют рейтинги Hub, по одной задаче за раз; можно ввести часть имени модели. В строке — итог теста, число ядер Trainium, wall time вместе с компиляцией и пик памяти на устройстве. В интерфейсе переключаются версии компилятора neuronx-cc и Transformers. Если модель не тестировали, Space пишет почему: чаще «не в приоритете», реже — несовместимость.

Запуск на своём Trainium2

Выбрав модель, видишь стек Neuron, пакеты для установки и run.py из прогона: загрузка через from_pretrained, шардирование при нехватке памяти, генерация со static cache и torch.compile(backend="neuron"). На инстансе Trainium2 (trn2), в том числе в AWS Neuron Deep Learning Container, хватает python3 run.py — без torchrun: при шардировании скрипт сам поднимает ранги torch.distributed.

Первый прогон уходит на компиляцию графов. Чтобы не ждать, можно подтянуть NEFF cache, который собрали при тесте — Space даёт команду с версией компилятора:

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

Что значит «runs»

• Модель кладут на минимальное число ядер Trainium, которое вмещает веса; у большинства — tensor parallelism, у MoE без TP-плана — только expert parallelism.
• Генерируют 32 токена и сверяют: на одном ядре — с CPU-прогоном того же чекпоинта; при шардировании связность проверяет небольшая judge-модель.
• Тестируют только модели до 4 устройств Trainium2 — не больше 16 ядер. PASS не обещает throughput и длинные генерации.

Те же строки лежат в Parquet по паре версий neuronx-cc и библиотеки; датасет читается напрямую, например через pandas с путём hf://datasets/aws-neuron/torch-neuronx-model-support/data — в каждой записи результат, ошибка при падении, версии стека и скрипт прогона.

Источник: Does Your Model Run on AWS Trainium? Look It Up.