AWS-ове белешке о издању Neuron 2.33 наводе 19. октобар као датум јавне бета верзије TorchNeuron-а. PyTorch објава, такође датирана на 19. октобар, каже да јавна бета верзија и издање отвореног кода постају доступни 20. октобра. Пакет torch-neuronx региструје neuron уређај у PyTorch-у за Trainium2 и Trainium3, па подржане операције над моделима и тензорима могу да користе познате PyTorch API-је. За инжењере који разматрају ове чипове, корисна промена је краћи пут од постојећег PyTorch програма до пробе на уређају. Да ли ће тај програм радити исправно и ефикасно и даље зависи од његових операција, облика и софтверског окружења.

Ранији PyTorch пут за Trainium користио је библиотеку torch-neuronx засновану на XLA и код за уређај torch_xla. AWS каже да изворни бекенд користи PyTorch проширење PrivateUse1 и замењује xm.xla_device() смештањем на уређај neuron. То је значајна промена интерфејса за пројекте који већ користе Neuron, иако је назив пакета познат. Издање важи за обуку и инференцију на инстанцама Trn2 и Trn3; матрица подршке не наводи Trn1.

Шта се мења у PyTorch програму

У AWS-овим белешкама о издању приказан је модел премештен помоћу .to("neuron") и улазни тензор направљен уз device="neuron". Подржане операције torch.nn, autograd-а и оптимизатора затим се извршавају преко бекенда. Код написан за PyTorch API независан од уређаја torch.accelerator може да открије Neuron без прераде специфичне за тај уређај, према документу Преглед TorchNeuron-а.

Постоје два пута извршавања. Eager режим извршава PyTorch операције када их Python позове, уз очување уобичајеног тока контроле, исписа и прекидних тачака. Свака операција се компајлира при првом коришћењу, а наредни одговарајући позиви могу да користе кеш. Инжењери могу и да компајлирају модел или изабрани део помоћу torch.compile(model, backend="neuron", dynamic=False). Изворни бекенд компајлира подржане графове пролаза унапред и уназад; код ван обухваћеног дела наставља да ради у eager режиму. PyTorch објава описује и једнократно постављање torch.compiler.set_default_backend("neuron") пре компилације, тако да каснији позиви компилације могу да изоставе аргумент бекенда.

Издање укључује neuron бекенд за torch.distributed, са документованим колективним операцијама и операцијама од тачке до тачке. AWS међу подржаним API-јима наводи DistributedDataParallel, FSDP2, DTensor и DeviceMesh, као и цевоводни паралелизам. То су корисни путеви за разматрање у пројекту дистрибуиране обуке, али белешке о издању не доказују да ће постојећи задатак на више чворова радити без измена или достићи одређени циљ пропусне моћи.

Евалуација почиње усклађеним окружењем

AWS-ово кратко упутство за инсталацију захтева инстанцу Trn2 или Trn3 и Neuron Base DLAMI или подржани Ubuntu AMI. Матрица издања наводи Linux x86_64 на Ubuntu 24.04 или 26.04. Подршка за Ubuntu 22.04 је завршена за ово издање. Наведени су Python 3.12, 3.13 и 3.14, уз torch-neuronx 2.13.3.0 за PyTorch 2.13.0 или torch-neuronx 2.14.3.0 за PyTorch 2.14.0. Породица пакета мора да одговара верзији PyTorch-а.

Кратко упутство инсталира torch-neuronx==2.14.* заједно са neuronx-cc==2.* и nki из Neuron и PyTorch индекса пакета. Затим проверава доступност уређаја и пореди малу матричну операцију на Neuron-у са резултатом на CPU-у. Тиме се проверава пут до инсталираног уређаја у AWS-овом примеру. Не проверава се читаочев модел. Разуман следећи корак евалуације је да се изврше пролази унапред и уназад стварног модела, прегледају забележене операције пребачене на CPU помоћу torch_neuronx.get_fallback_ops(), упореде излази са тренутним окружењем и одвојено измере компилација и извршавање у устаљеном стању. Ти последњи кораци су питања за евалуацију која предлаже BIG CHANGE, а не тестови које смо ми спровели.

Где бета верзија још може да промени рад

У eager режиму, неподржана операција може да пређе на CPU и прикаже упозорење. Када се користи torch.compile, операција обухваћена графом коју Neuron бекенд не може да преведе у нижи извршни облик подразумевано доводи до неуспешне компилације. AWS упућује програмере на своју табелу компатибилности оператора за покривеност сваке операције. Зато успешно завршено извршавање у eager режиму само по себи не доказује да је све извршено на уређају или да је подржана компилација.

У издању се такође наводи да torch.compile захтева dynamic=False; нови облик улаза може да захтева нову компилацију. AWS предлаже допуњавање или груписање по облицима како би се смањиле њихове варијације. Прво извршавање укључује компилацију ако нема одговарајућег артефакта у кешу. Ови детаљи су важни при мерењу радног оптерећења: кашњење компилације и образац облика улаза могу да промене резултат евалуације.

Друга изричита ограничења сужавају тврдњу о компатибилности. Бета верзија не подржава torch.sparse, torch.func, torch.nested или torch.distributions. Само неке torch.fft трансформације раде на Neuron-у, а float64 се извршава као float32 уз упозорење. Ово издање нема интеграцију са TorchAO нити подршку за сложени тип података MXFP8. Тим који користи ове интерфејсе мора да провери сопствени пут кроз код пре него што смештање на уређај схвати као план миграције.

Велика промена

TorchNeuron у јавној бета верзији поставља Trainium иза PyTorch-ових изворних интерфејса за уређаје и компилацију. Тиме прва проба постаје непосреднија за подржане PyTorch програме, а постојећи корисници Neuron-а добијају другачији пут од свог кода заснованог на XLA. Пресудни докази за усвајање долазиће из сваког радног оптерећења: да ли његови оператори остају на предвиђеном уређају, да ли нумерички резултати остају прихватљиви и да ли измерене перформансе и трошкови оправдавају инстанцу и рад на миграцији. BIG CHANGE је за ову анализу прегледао документацију и објаве; нисмо покретали TorchNeuron нити независно проверавали AWS-ове тврдње о перформансама.

Извори и додатно читање

  • Белешке о издању AWS Neuron 2.33 TorchNeuron: датум издања, бета статус, матрица верзија и оперативних система, API-ји и позната ограничења.
  • AWS-ов преглед TorchNeuron-а: интеграција уређаја, eager и компајлирано извршавање, понашање при преласку на CPU и софтверски пут кроз Neuron.
  • AWS-ово кратко упутство за инсталацију и проверу: предуслови за инстанцу и слику система, инсталација пакета и документована провера уређаја.
  • PyTorch блог: представљање TorchNeuron-а: објашњење промене уређаја, компилације и дистрибуираних API-ја из угла PyTorch-а. Чланак је објава AWS сарадника на PyTorch блогу, а не независни тест перформанси.
  • Amazon-ова објава: AWS-ов приказ циљева и раних корисника. Његове тврдње о миграцији, перформансама и економским користима захтевају проверу за конкретно радно оптерећење.