TorchNeuronX: Как запустить модели Hugging Face на AWS Trainium

Если вы работаете с большими языковыми моделями (LLM) или мультимодальными системами в облачной инфраструктуре AWS, вопрос оптимизации затрат и производительности стоит особенно остро. Традиционно экосистема Hugging Face тесно связана с GPU NVIDIA, но с появлением специализированных чипов, таких как AWS Trainium, архитектура распределенных вычислений меняется. В этой статье мы подробно разберем, как убедиться, что ваша любимая модель с Hugging Face Hub действительно работает на AWS Trainium без необходимости переписывать код, и как использовать новые инструменты сообщества для ускорения разработки.

Ранее мы уже знакомы с пакетом torch-neuronx — нативным бэкендом PyTorch для AWS Trainium. Его главная сила заключается в том, что он позволяет библиотеке Transformers запускать собственный код модели на чипах Trainium так же естественно, как это происходит на любом другом устройстве. Теоретически, любая модель с Hub должна работать. Но на практике? Не все модели одинаково хорошо адаптируются к специфике нейронных процессоров. Здесь на сцену выходит новый ресурс: TorchNeuronX Model Support Space.

Этот Space отвечает на самый важный вопрос для инженеров ML: «Будет ли моя модель работать на Trainium без модификаций?». Мы рассмотрим, как пользоваться этим инструментом, что означают результаты тестов, как запустить модель самостоятельно и какие технические нюансы стоит учитывать при развертывании.

01Что такое TorchNeuronX и зачем он нужен?

Прежде чем углубляться в детали проверки моделей, важно понять контекст. AWS Trainium — это специализированный чип, разработанный Amazon Web Services специально для машинного обучения. Он предлагает высокую эффективность при обучении и выводе моделей, но требует специфического программного стека. Обычные бинарные файлы PyTorch, оптимизированные для CUDA (архитектура NVIDIA), не работают на Trainium напрямую.

Здесь вступает в игру torch-neuronx. Это пакет, который действует как мост между стандартным кодом PyTorch/Transformers и архитектурой Neuron. Он использует компилятор NeuronX (neuronx-cc), который преобразует графы вычислений PyTorch в формат, понятный чипам Trainium. Ключевое преимущество этого подхода — абстракция. Разработчикам не нужно писать код на низкоуровневых API Neuron. Они пишут стандартный Python-код с использованием transformers, а torch-neuronx берет на себя сложность трансляции и распределения вычислений.

Однако, поскольку компиляция графов — процесс нетривиальный, некоторые архитектуры моделей могут сталкиваться с ограничениями. Например, специфические операции внимания, нестандартные слои или сложные условия в коде модели могут не поддерживаться текущей версией компилятора. Именно поэтому возникла необходимость в централизованном репозитории совместимости.

02Исследование Space: Как найти свою модель

Space TorchNeuronX Model Support — это интерактивный интерфейс, который агрегирует результаты тестирования популярных моделей. Он позволяет быстро проверить статус совместимости, не тратя часы на попытки запуска, которые могут завершиться ошибкой.

Вкладка "Selected" (Избранное)

При открытии Space вы попадаете на вкладку Selected. Это curated-список моделей, отобранных командой Hugging Face. Критерии отбора основаны на двух метриках Hugging Face Hub: количество загрузок (downloads) и трендовый рейтинг (trending score). Это означает, что вы видите самые популярные и актуальные модели, которые с высокой вероятностью могут понадобиться в ваших проектах.

Каждая строка в таблице содержит критически важную информацию:

Результат (Result): Прошла ли модель тест или нет.
Количество ядер Trainium (Trainium cores): Сколько чипов использовалось для шардинга модели. Это важно для оценки требований к инфраструктуре.
Время выполнения (Wall time): Общее время, включая этап компиляции графов. Обратите внимание, что первый запуск всегда включает компиляцию, что может занимать значительное время.
Пиковая память устройства (Peak device memory): Сколько VRAM (или эквивалента в Neuron) потребовалось модели.

TorchNeuronX: Как запустить модели Hugging Face на AWS Trainium

Вкладки "Trending" и "Downloads"

Если нужной вам модели нет в избранном, вы можете воспользоваться вкладками Trending или Downloads. Эти представления следуют официальным рейтингам Hugging Face Hub, разделенным по задачам (например, только текстовая генерация или мультимодальные модели). Вы можете ввести часть имени модели в строку поиска.

Важно понимать логику отображения: если модель отсутствует в результатах, Space объясняет причину. В большинстве случаев это означает, что модель не была приоритетной для тестирования на данный момент. В редких случаях модель может быть действительно несовместима с текущей версией стека Neuron. Также вы можете выбрать конкретные версии компилятора neuronx-cc и библиотеки transformers, чтобы увидеть результаты для вашей конкретной конфигурации.

На момент написания этой статьи тестирование охватывает преимущественно модели для text-generation (генерация текста) и image-text-to-text (мультимодальные модели, такие как LLaVA). В планах команды — расширение покрытия на библиотеки diffusers (генерация изображений), sentence-transformers и другие задачи экосистемы Transformers.

💡

Совет по оптимизации. Если вы видите, что модель требует большого количества ядер Trainium, попробуйте проверить, можно ли запустить её на меньшем количестве устройств, пожертвовав скоростью генерации. Space показывает минимальное количество ядер, необходимое для успешного запуска, что помогает планировать бюджет на AWS.

03Как запустить модель самостоятельно

Самая ценная часть Space — это не просто статус «работает/не работает», а готовые инструкции для развертывания. Если модель прошла тест, Space предоставляет полный рецепт для запуска.

Стек и скрипты

Выбрав модель, вы увидите:

Версии стека Neuron: Конкретные версии neuronx-cc и библиотек, на которых проводилось тестирование. Строгое соблюдение версий критично, так как компилятор чувствителен к изменениям в коде PyTorch.
Пакеты для установки: Список зависимостей, которые нужно установить поверх базового образа.
Скрипт run.py: Файл, который использовался для тестирования. Вы можете скачать его и адаптировать под свои нужды.

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

Процесс развертывания

Для запуска вам потребуется инстанс AWS Trainium2 (тип trn2). Рекомендуется использовать AWS Neuron Deep Learning Container, так как он уже содержит базовые драйверы и библиотеки Neuron.

Порядок действий:

Установите необходимые пакеты в контейнере.
Скопируйте скрипт run.py из Space.
Запустите скрипт командой python3 run.py.

Обратите внимание: вам не нужно использовать torchrun для запуска распределенных задач. Скрипт run.py автоматически инициализирует torch.distributed ранги, если модель требует шардинга на несколько ядер. Это упрощает процесс, так как вам не нужно вручную настраивать параметры запуска MPI или Torch Distributed.

04Что на самом деле означает "Runs" (Работает)?

Результат PASS в Space — это хороший знак, но он не гарантирует, что модель будет работать идеально в продакшене. Важно понимать методологию тестирования, чтобы правильно интерпретировать данные.

Методология тестирования

Каждый тест выполняется на инстансе trn2. Модель загружается и шардится на минимально возможное количество ядер Trainium, способное уместить веса модели. Обычно используется Tensor Parallelism (параллелизм по тензорам). Для моделей типа Mixture-of-Experts (MoE), которые не имеют плана тензорного параллелизма, используется только Expert Parallelism.

После загрузки модель генерирует ровно 32 токена. Затем происходит проверка:

На одном ядре: Сгенерированные токены сравниваются с токенами, полученными при запуске той же модели на CPU. Если они совпадают, это подтверждает корректность вычислений.
При шардинге: Поскольку сравнение токенов в распределенной среде сложнее, используется небольшая "судейская" модель (judge model), которая проверяет, является ли сгенерированный текст связным английским языком. Это эвристический метод, но он эффективно выявляет грубые ошибки.

Ограничения тестов

На данный момент тестируются только модели, которые помещаются на 4 устройства Trainium2 (максимум 16 ядер). Это означает, что очень большие модели (например, 70B+ параметров) могут не быть протестированы, если они не могут быть эффективно размещены в этом лимите. Кроме того, результат PASS ничего не говорит о пропускной способности (throughput) при длительных генерациях или о поведении модели при высокой нагрузке. Это базовый функциональный тест, а не нагрузочное тестирование.

⚠️

Важно: Компиляция. Первый запуск модели всегда включает этап компиляции графов Neuron. Это может занять от нескольких минут до десятков минут в зависимости от сложности модели. Не пугайтесь долгого ожидания при первом запуске — это нормальный процесс.

05Ускорение запуска: Кэш NEFF

Поскольку компиляция графов — ресурсоемкий процесс, повторять её каждый раз при запуске скрипта неэффективно. AWS предоставляет публичный кэш скомпилированных графов, известный как NEFF cache.

Space предоставляет команду для синхронизации этого кэша. Если вы знаете, какую версию компилятора (neuronx-cc) использовал тест, вы можете скопировать готовые графы из публичного бакета S3:

terminalbash

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

После этого, запустив скрипт с указанием пути к кэшу, вы сможете пропустить этап компиляции и сразу перейти к генерации. Это значительно ускоряет итерации разработки и тестирование.

06Данные как актив: Датасет TorchNeuronX

Все результаты, отображаемые в Space, основаны на публичном датасете aws-neuron/torch-neuronx-model-support. Это означает, что вы не ограничены интерфейсом веб-приложения. Вы можете получить прямой доступ к данным для анализа, построения собственных дашбордов или интеграции в CI/CD пайплайны.

Датасет структурирован так, что для каждой пары версий neuronx-cc и библиотеки существует отдельный файл Parquet. Каждая строка содержит:

ID модели.
Результат теста (PASS/FAIL).
Сообщение об ошибке (если тест провалился).
Версии всего стека (компилятор, библиотеки).
Скрипт run.py, который использовался.

Вы можете запросить данные напрямую с помощью Python и библиотеки Pandas:

terminalpython

import pandas as pd

df = pd.read_parquet(
    "hf://datasets/aws-neuron/torch-neuronx-model-support/data"
)

# Пример запроса для конкретной модели
model_data = df[df.model_id == "Qwen/Qwen3-0.6B"][[
    "neuronx_cc",
    "library",
    "library_version",
    "neuron_cores",
    "result"
]]

print(model_data)

Карточка датасета (dataset card) на Hugging Face Hub содержит подробное описание каждого поля и значения, что делает его удобным для программистов, предпочитающих работать с данными, а не с UI.

07Что это значит на практике

Для инженеров, работающих с AWS Trainium, появление этого Space и датасета снимает главный барьер входа — неопределенность. Раньше проверка совместимости модели требовала ручного развертывания инстансов, установки зависимостей и отладки ошибок компиляции. Теперь этот процесс стандартизирован и документирован.

Стратегия работы должна выглядеть следующим образом:

Проверка совместимости: Перед началом разработки проверьте модель в Space. Если она есть и имеет статус PASS, вы можете быть уверены в базовой функциональности.
Планирование ресурсов: Используйте данные о количестве ядер и памяти для выбора правильного типа инстанса AWS. Это помогает избежать ситуаций, когда модель не помещается в память или требует избыточных ресурсов.
Использование кэша: Всегда синхронизируйте NEFF кэш для ускорения разработки. Это особенно важно при частых перезапусках скриптов.
Анализ данных: Если вы работаете с кастомными моделями или редкими архитектурами, используйте датасет для анализа паттернов ошибок. Возможно, ваша модель падает из-за известной проблемы с конкретной версией компилятора, которую уже можно обойти обновлением.

Экосистема AWS Trainium продолжает развиваться. Команда Hugging Face планирует расширить покрытие на генерацию изображений и другие задачи, что сделает этот инструмент еще более универсальным. Для тех, кто стремится к оптимизации затрат на ML-инфраструктуру, TorchNeuronX и связанные с ним ресурсы становятся обязательным инструментом в арсенале.

Помните, что "работает" на Trainium — это не просто технический факт, это возможность запустить мощные модели LLM с высокой эффективностью. Используйте эти инструменты, чтобы сократить время от идеи до продакшена.

Источник: Hugging Face ↗