TorchNeuron. PyTorch działa na chipach AWS Trainium bez przepisywania kodu

Co zmienia TorchNeuron

Do tej pory uruchomienie modelu na Trainium wymagało pakietu Neuron SDK opartego na PyTorch/XLA, czyli warstwie pośredniej, która tłumaczyła operacje na format zrozumiały dla chipu. Kod trzeba było dostosować, a część bibliotek nie działała od razu. Hugging Face przyznaje, że jego modele wymagały osobnej konfiguracji pod Neuron.

TorchNeuron korzysta z oficjalnego mechanizmu PyTorch do podłączania nowego sprzętu (PrivateUse1). Dzięki temu standardowe moduły, takie jak torch.nn, torch.autocast, torch.distributed i torch.profiler, działają na urządzeniu neuron bez zmian. Kod napisany z użyciem uniwersalnego API torch.accelerator ma działać bez modyfikacji. Według dokumentacji AWS Neuron PyTorch 2.9 to ostatnia wersja obsługiwana przez starsze rozwiązanie oparte na XLA, a od PyTorch 2.10 wsparcie zapewnia TorchNeuron.

Jak uruchomić model PyTorch na Trainium

Według dokumentacji AWS przejście sprowadza się do kilku kroków. Pakiet będzie dostępny od 20 października w repozytorium pip AWS Neuron i na GitHubie w projekcie aws-neuron/torch-neuronx.

1. instalacja pakietu poleceniem pip install torch-neuronx z repozytorium AWS Neuron,
2. przeniesienie modelu na Trainium zapisem model.to('neuron') lub device_map='neuron' w Transformers,
3. dla trybu produkcyjnego jednorazowe ustawienie backendu torch.compiler.set_default_backend('neuron') i wywołanie torch.compile,
4. do treningu rozproszonego standardowe moduły DDP, FSDP2 oraz równoległość tensorowa i potokowa,
5. do analizy wydajności torch.profiler z konfiguracją NeuronConfig i narzędziem Neuron Explorer.

Czego TorchNeuron jeszcze nie obsługuje

Zapowiedź „bez przepisywania kodu” dotyczy typowych modeli. AWS sam publikuje listę luk. W sierpniowym przebiegu testów PyTorch z ok. 154 tys. wyników ok. 39 tys. pominięto celowo z powodów specyficznych dla Neuron. Największa grupa, ok. 22 tys. testów, dotyczy operatorów bez natywnej obsługi, m.in. części operacji na indeksach, zaawansowanej algebry liniowej i komórek LSTM oraz GRU.

W pierwszej wersji świadomie pominięto tensory rzadkie, NestedTensor, różniczkowanie w przód, dowolne maski w natywnej uwadze wielogłowicowej, kwantyzację per kanał oraz dynamiczne i symboliczne kształty. Brak dynamicznych kształtów może utrudnić obsługę zapytań o różnej długości bez dodatkowej pracy. Tryb eager kompiluje każdą operację osobno, więc jest wolniejszy od torch.compile. Szybszy tryb adaptive eager ma status eksperymentalny.

Kto już korzysta i co deklaruje

Amazon wymienia wśród użytkowników Cartesię, Deepgram i Hugging Face. Cartesia, twórca modeli głosowych opartych na architekturze Mamba, podaje, że uruchomiła swój proces inferencji przy minimalnych zmianach w kodzie. Deepgram po kilkumiesięcznym projekcie z AWS deklaruje ok. 60 proc. więcej przetworzonych sekund nagrań na dolara w porównaniu z dotychczasowymi GPU, nie podaje jednak, z jakimi kartami porównywał wyniki.

Hugging Face dodał Trainium jako obsługiwane urządzenie w bibliotekach Transformers, Diffusers i TRL. Na liście firm, które według AWS przekazywały uwagi do projektu, są m.in. Anthropic, Databricks, Meta i Salesforce. Komunikat nie podaje cen instancji ani nowych regionów, więc koszt trzeba liczyć według obecnego cennika AWS dla instancji Trn2 i Trn3.

Co to oznacza dla zespołów w Polsce

Dla polskich firm i uczelni, które trenują lub dostrajają modele w AWS, TorchNeuron obniża koszt próby. Zamiast przenosić projekt na nowy zestaw narzędzi, można uruchomić istniejący kod na Trainium i porównać koszt z instancjami GPU. W naszej ocenie najwięcej mogą zyskać zespoły pracujące na popularnych modelach z Hugging Face i standardowych pętlach treningowych.

Przed decyzją warto sprawdzić trzy rzeczy. Pierwsza to dostępność instancji Trainium w regionie, w którym muszą zostać dane, co ma znaczenie przy RODO. Druga to lista nieobsługiwanych funkcji, zwłaszcza dynamiczne kształty. Trzecia to status projektu, bo 20 października startuje dopiero publiczna beta. Na rynku chipów AI to kolejny ruch po rundzie Tenstorrent i premierze akceleratora Intela, które również celują w koszt inferencji.