Проблема совместимости решена через автоматизацию
Команда Hugging Face представила инструмент torch-neuronx model support, который позволяет разработчикам мгновенно узнать, запустится ли конкретная модель из Hub на чипах AWS Trainium (Trn1/Trn2). Ранее пользователям приходилось самостоятельно настраивать стек Neuron, что часто приводило к ошибкам компиляции. Теперь достаточно ввести имя модели, чтобы получить статус совместимости, количество необходимых ядер и ссылку на готовый скрипт запуска.
Методология тестирования и ограничения
Тестирование проводится на инстансах trn2 (Trainium2). Ключевые параметры проверки:
- Параллелизм: Используется тензорный параллелизм (tensor parallelism) для распределения весов. Для MoE-моделей без плана тензорного параллелизма применяется экспертный параллелизм.
- Объем памяти: Проверяются модели, помещающиеся на 4 устройства Trainium2 (максимум 16 ядер).
- Валидация: Генерируется 32 токена. На одном ядре результаты сравниваются с CPU-версией. При шардинге используется «судья» (judge model) для проверки связности английского текста.
- Что НЕ проверяется: Пропускная способность (throughput) и качество генерации длинных последовательностей.
Технические детали запуска
Для воспроизведения результата используется стандартный рецепт TorchNeuron:
- Загрузка модели через
from_pretrained. - Шардинг весов при необходимости.
- Генерация с использованием статического кэша и
torch.compile(backend="neuron").
Первый запуск включает этап компиляции графов. Чтобы ускорить процесс, можно синхронизировать готовые графы из публичного кэша NEFF:
hf buckets sync hf://buckets/aws-neuron/torch-neuronx-neff-cache/neuronxcc-<version> /tmp/neff_cacheДоступ к данным и API
Все результаты тестов выложены в открытый датасет aws-neuron/torch-neuronx-model-support. Каждая строка содержит ID модели, версию компилятора neuronx-cc, версию библиотеки, количество ядер, статус результата и ошибку (если была). Это позволяет программно запрашивать совместимость:
| Параметр | Значение/Пример | Примечание |
|---|---|---|
| Платформа | AWS Trainium (Trn1/Trn2) | Backend: torch-neuronx |
| Типы моделей | text-generation, image-text-to-text | Планируется расширение на diffusers и sentence-transformers |
| Макс. ядер | 16 (4 устройства trn2) | Ограничение текущей выборки тестов |
| Датасет | aws-neuron/torch-neuronx-model-support | Формат: Parquet, один файл на пару версий |
Почему это важно
Инструмент снижает порог входа для использования AWS Trainium. Разработчикам больше не нужно гадать, подойдет ли их модель. Наличие готового скрипта run.py и точных версий стека позволяет сразу перейти к оптимизации производительности, а не к отладке совместимости.
Источник: Hugging Face blog ↗