В предыдущей публикации мы представили TorchNeuron — родной бэкенд PyTorch для AWS Trainium (пакет torch-neuronx). Благодаря ему библиотека Transformers запускает собственный код моделей на Trainium так же, как и на любом другом устройстве, поэтому любая модель Transformers из Hub может быть запущена. Какие из них работают на практике?
Пространство (Space) поддержки моделей torch-neuronx отвечает на этот вопрос для самых популярных моделей на Hub. Оно показывает, запускается ли модель на Trainium без изменений, и предоставляет скрипт для ее самостоятельного запуска.
Поиск модели
Пространство открывается на вкладке Selected (Выбранное): это наша собственная подборка моделей, отобранных среди лучших моделей на Hub по количеству загрузок и трендовому рейтингу. Каждая строка содержит результат, количество ядер Trainium, по которым шардирована модель, время выполнения (включая компиляцию) и пиковое использование памяти устройства.
Чтобы найти конкретную модель, перейдите на вкладку Trending (В тренде) или Downloads (Загрузки) и введите часть ее названия. Эти вкладки следуют собственным рейтингам Hub по одной задаче за раз, независимо от того, тестировалась модель или нет. Если модель была намеренно пропущена, Пространство объясняет почему (в большинстве случаев она не была приоритетной, иногда она действительно несовместима). Вы также можете выбрать версии компилятора Neuron (neuronx-cc) и Transformers для просмотра.
На момент написания этой статьи тестируются только модели Transformers, а наша выборка ограничена моделями генерации текста (text-generation) и «изображение-текст в текст» (image-text-to-text). Мы планируем расширить поддержку на диффузоры, трансформеры предложений и другие задачи Transformers.
На момент написания этой статьи тестируются только модели Transformers, а наша выборка ограничена моделями генерации текста (text-generation) и «изображение-текст в текст» (image-text-to-text). Мы планируем расширить поддержку на диффузоры, трансформеры предложений и другие задачи Transformers.
Запустите ее самостоятельно
Выберите модель, чтобы узнать, как ее запустить. В Пространстве указан стек Neuron, на котором тестировалась модель, пакеты, которые нужно установить поверх него, и выполненный скрипт run.py. Скрипт следует рецепту из анонса TorchNeuron: он загружает модель с помощью from_pretrained, при необходимости разделяя ее на шарды, а затем генерирует текст со статическим кэшем и torch.compile(backend="neuron").
На инстансе Trainium2 (trn2) с этим стеком, например в AWS Neuron Deep Learning Container, установите пакеты и выполните:
Нет необходимости использовать torchrun: скрипт сам запускает процессы torch.distributed при шардировании модели.
Первый запуск начинается с компиляции графов модели. Чтобы пропустить этот шаг, скопируйте графы, созданные тем же компилятором во время теста, из общедоступного бакета кэша NEFF. Пространство предоставляет команду с правильной версией компилятора:
Что означает «запускается»
Каждый тест загружает модель на инстанс trn2, шардированный по наименьшему числу ядер Trainium, способному вместить ее веса. Большинство моделей шардируются с использованием тензорного параллелизма. Модели Mixture-of-experts без плана тензорного параллелизма используют только параллелизм экспертов. Затем тест генерирует 32 токена и проверяет их. На одном ядре токены сравниваются с результатами выполнения того же чекпоинта на CPU. Когда модель шардирована, небольшая модель-судья проверяет, является ли текст беглым английским языком.
На момент написания статьи тестируются только модели, помещающиеся на 4 устройствах Trainium2, то есть максимум на 16 ядрах. Статус PASS ничего не говорит о пропускной способности или о более длинных генерациях.
Результаты — это датасет
Все, что отображается в Пространстве, получено из общедоступного датасета aws-neuron/torch-neuronx-model-support, содержащего по одному файлу Parquet для каждой пары версий neuronx-cc и библиотеки.
Каждая строка содержит результат модели, ошибку в случае сбоя, версии всего стека и выполненный скрипт run.py. Вы можете делать запросы к нему напрямую:
Карточка датасета дает точное определение каждого результата.
В следующий раз, когда вы зададитесь вопросом, запускается ли модель на Trainium, просто проверьте ее.
