Hugging Face представила Space torch-neuronx model support. Он отвечает на практический вопрос: запускается ли конкретная популярная модель из Hub на AWS Trainium без изменений и как воспроизвести проверку.
Какие данные доступны
Для проверенной модели каталог показывает результат, число ядер Trainium для размещения весов, пиковое потребление памяти и полное время запуска с учётом компиляции. Модель можно искать в представлениях Trending и Downloads, а также выбрать версии neuronx-cc и Transformers.
- На момент публикации проверяются только модели Transformers для генерации текста и image-text-to-text.
- Каталог охватывает модели, которые помещаются максимум на четырёх устройствах Trainium2 — до 16 ядер.
- Если модель исключили из приоритетной выборки или она несовместима, Space указывает причину.
Как устроен запуск
В карточке модели приведены стек Neuron, дополнительные пакеты и run.py. По описанию Hugging Face, сценарий загружает модель через from_pretrained, распределяет её по ядрам при необходимости и запускает генерацию со static cache и torch.compile с backend neuron.
Для распределённых моделей не требуется отдельно запускать torchrun: скрипт поднимает ранги torch.distributed сам. Первый запуск включает компиляцию графов; чтобы её избежать при повторном запуске, можно использовать кэш NEFF, созданный тем же компилятором.
Ограничение результата PASS
Проверка загружает модель на Trainium2, генерирует 32 токена и сверяет результат. Для модели на одном ядре его сравнивают с запуском того же чекпойнта на CPU; для распределённой модели связность английского текста оценивает небольшая модель-судья.
Поэтому PASS подтверждает прохождение этой конкретной проверки, но не измеряет пропускную способность и не описывает поведение при длинной генерации. Результаты, ошибки, версии компонентов и run.py собраны в публичном наборе aws-neuron/torch-neuronx-model-support.
В качестве одного из окружений для запуска упоминаются AWS Neuron Deep Learning Containers — Docker-образы для обучения и обслуживания моделей на инстансах Trainium и Inferentia с AWS Neuron SDK.
