Работи ли моделът ви на AWS Trainium? Проверете

Работи ли вашият модел на AWS Trainium? Проверете

Публикувано 23 октомври 2026 г.

В предишна публикация представихме TorchNeuron — собствения бекенд на PyTorch за AWS Trainium (пакета torch-neuronx). С него Transformers изпълнява собствения си код за моделиране върху Trainium, както на всяко друго устройство, така че всеки модел на Transformers в Hub може да работи. Но кои от тях действително работят?

Space-ът torch-neuronx model support отговаря на този въпрос за най-популярните модели в Hub. Той показва дали даден модел работи без промени на Trainium и ви дава скрипта, с който да го стартирате сами.

Потърсете модел

Space-ът се отваря в изглед Selected: нашата селекция от модели, избрани сред водещите модели в Hub според броя изтегляния и актуалния резултат. Всеки ред показва резултата, броя ядра Trainium, върху които моделът е разделен, продължителността на изпълнението (включително компилацията) и пиковата памет на устройството.

За да потърсите конкретен модел, превключете към изглед Trending или Downloads и въведете част от името му. Тези изгледи следват собствените класации на Hub, по една задача наведнъж, независимо дали даден модел е тестван или не. Когато модел умишлено е пропуснат, Space-ът ви казва защо (в повечето случаи не е бил приоритетен, а понякога е действително несъвместим). Можете също да изберете версиите на компилатора Neuron (neuronx-cc) и на Transformers, които искате да разгледате.

Към момента на писане се тестват само модели на Transformers, а селекцията ни е ограничена до модели за text-generation и image-text-to-text. Планираме да разширим обхвата до diffusers, sentence transformers и други задачи на transformers.

Стартирайте го сами

Изберете модел, за да видите как да го стартирате. Space-ът показва Neuron стека, върху който е тестван моделът, пакетите, които трябва да инсталирате върху него, и изпълнения скрипт run.py. Скриптът следва рецептата от анонса на TorchNeuron: зарежда модела с from_pretrained, като при нужда го разделя на части, след което генерира със статичен кеш и torch.compile(backend="neuron").

На инстанция Trainium2 (trn2) с този стек, например в AWS Neuron Deep Learning Container, инсталирайте пакетите и след това изпълнете:

python3 run.py

Няма нужда от torchrun: скриптът сам стартира собствените си рангове на torch.distributed, когато моделът е разделен на части.

При първото стартиране започва компилация на графите на модела. За да пропуснете тази стъпка, копирайте графите, изградени от същия компилатор по време на теста, от публичния bucket NEFF cache. Space-ът ви дава командата с правилната версия на компилатора:

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

Какво означава „работи“

Всеки тест зарежда модела на инстанция trn2, разделен на части върху минималния брой ядра Trainium, които могат да поберат теглата му. Повечето модели се разделят чрез паралелизъм по тензори. Моделите със смес от експерти, които нямат план за паралелизъм по тензори, използват само паралелизъм по експерти. След това тестът генерира 32 токена и ги проверява. При едно ядро токените се сравняват с тези от изпълнение на същия чекпойнт върху CPU. Когато моделът е разделен на части, малък съдийски модел проверява дали текстът е гладък английски.

Към момента на писане се тестват само модели, които се побират на 4 устройства Trainium2, което означава най-много 16 ядра. PASS не казва нищо за пропускателната способност или за по-дълги генерирания.

Резултатите са набор от данни

Всичко, което Space-ът показва, идва от публичния набор от данни aws-neuron/torch-neuronx-model-support, с по един Parquet файл за всяка двойка версии на neuronx-cc и библиотеките.

Всеки ред съдържа резултата за даден модел, грешката при неуспех, версиите на целия стек и изпълнения run.py. Можете да го направите заявка директно:

import pandas as pd

df = pd.read_parquet("hf://datasets/aws-neuron/torch-neuronx-model-support/data")
df[df.model_id == "Qwen/Qwen3-0.6B"][["neuronx_cc", "library", "library_version", "neuron_cores", "result"]]

Картата на набора от данни дава точното значение на всеки резултат.

Следващия път, когато се чудите дали даден модел работи на Trainium, проверете.