Amazon ogłosił TorchNeuron, natywny backend PyTorch o otwartym kodzie źródłowym dla rodziny chipów AI Trainium, co usuwa jedną z największych praktycznych barier w stosowaniu niestandardowych krzemów do obciążeń związanych z uczeniem maszynowym. Zgodnie z oficjalnym postem na blogu PyTorch opublikowanym przed konferencją PyTorch w Ameryce Północnej 20 października, programiści mogą teraz uruchamiać istniejący kod modelu PyTorch na Trainium, wprowadzając jednowierszową zmianę, zamiast przepisywać swoje przepływy pracy w oparciu o zastrzeżone narzędzia.

Ogłoszenie przykuło natychmiastową uwagę, ponieważ dotyczy tego, co analitycy od dawna określali jako najtrwalszą zaletę Nvidii: nie samej wydajności chipów, ale ekosystemu oprogramowania CUDA, który łączy programistów z procesorami graficznymi. Reuters sformułował tę publikację jako stwierdzenie Amazona, że ​​jego chipy AI będą działać z najpopularniejszym frameworkiem programistycznym w AI, podczas gdy TechCrunch nazwał TorchNeuron największym jak dotąd wyzwaniem firmy dla Nvidii. Aby uzyskać szerszy obraz zmian na rynku akceleratorów, na bieżąco śledzimy ten rozwój aktualności dotyczące sprzętu AI.

Co właściwie zmienia TorchNeuron

Podstawową obietnicą TorchNeuron jest prostota. W przykładzie opublikowanym na blogu PyTorch istniejąca pętla treningowa działa na Trainium po niczym innym jak przeniesieniu modelu na nowe urządzenie za pomocą `.to('neuron')`. Jak wynika z postu, kod napisany dla `torch.accelerator`, zalecanego przez PyTorch API niezależnego od urządzenia, działa bez modyfikacji. W przypadku skompilowanych obciążeń programiści ustawiają domyślny backend za pomocą `torch.compiler.set_default_backend("neuron")`, a następnie wywołują `model.compile()` bez dalszych zmian.

Pod maską TorchNeuron rejestruje Trainium poprzez PrivateUse1, oficjalny punkt rozszerzenia PyTorch dla niestandardowych backendów sprzętowych, co oznacza, że ​​chip zachowuje się jak standardowe urządzenie PyTorch, a nie dodatek przykręcony do szkieletu. AWS opisuje integrację jako obejmującą kilka trybów wykonania:

  • Eager wykonanie na potrzeby badań i debugowania, z eksperymentalnym „adaptacyjnym trybem chęci”, który zmniejsza obciążenie związane z uruchamianiem
  • torch.compile do wydajności produkcyjnej, przechwytywania modeli na wykresach optymalizowanych przez kompilator Neuron, z natywną obsługą FlexAttention
  • Interfejs jądra neuronu (NKI) dla badaczy potrzebujących niestandardowych jąder niskiego poziomu

Jak wynika z zapowiedzi, backend działa zarówno na układach Trainium2, jak i nowszych Trainium3.

Bezpośredni atak na fosę oprogramowania Nvidii

Kontrola Nvidii w dziedzinie obliczeń AI nigdy nie opierała się wyłącznie na szybszym sprzęcie. Na CUDA, platformie oprogramowania firmy działającej od dwudziestu lat, pisze się, dostraja i wdraża większość kodu badawczego AI, a każdy chip będący wyzwaniem – od linii AMD Instinct po procesory TPU Google i Trainium Amazona – musi zmagać się z kosztami poproszenia programistów o opuszczenie go. Relacje z premiery TorchNeuron w serwisach GeekWire, Benzinga i GuruFocus jednolicie ujęły to ogłoszenie jako atak na tę fosę.

Amazon nie próbuje wyciągać programistów z PyTorch, aby ich przekonać. Zamiast tego wprowadza Trainium do PyTorch, dzięki czemu ogromna baza użytkowników frameworka może przyjąć chip bez zmiany nawyków. W poście na blogu PyTorch zawarto poparcie ze strony firmy Meta, która utrzymuje platformę: „Natywne wsparcie dla AWS Trainium poszerza ten wybór, umożliwiając badaczom i inżynierom eksperymentowanie i iterację przy użyciu chętnie wykonywanych i rozproszonych interfejsów API, które już znają” – powiedziała Jana van Greunen, dyrektor PyTorch Engineering w Meta, w poście.

To ramy mają znaczenie dla działalności chmurowej AWS. Instancje Trainium są reklamowane jako tańsza alternatywa dla procesorów graficznych Nvidia do celów szkoleniowych i wnioskowania, a chipy przyciągnęły już klientów zajmujących się sztuczną inteligencją. Sprawienie, że chipy będą sprawiać wrażenie natywnych dla PyTorch, bezpośrednio rozwiązuje zarzut, który spowalnia szersze przyjęcie: wysiłek związany z przenoszeniem.

Zakres techniczny, od notebooków po szkolenia w skali klastra

Możliwości wymienione na blogu PyTorch sugerują, że Amazon celuje daleko poza eksperymenty hobbystyczne. TorchNeuron obejmuje zestawy operatorów widoczne w modelach tekstowych, multimodalnych i wizjonerskich, a także generowanie obrazów i wideo, modele przestrzeni stanów i modele świata z precyzją BF16, FP32 i FP16; Obsługa TorchAO i FP8 znajduje się w prywatnej wersji beta przed ogólną dostępnością w następnej wersji.

Szkolenie rozproszone, dyscyplina, która przekształca kod pojedynczego procesora graficznego w uruchomienia w skali klastra, jest obsługiwane przez standardowe interfejsy API PyTorch — DDP, FSDP2, SimpleFSDP, równoległość tensorów i potoków oraz sharding w stylu ZeRO, a także DTensor i DeviceMesh dla wielowymiarowych grup procesów. Jak napisano w poście, strumienie i zdarzenia odwzorowują semantykę PyTorch na kolejki sprzętowe Trainium.

Podsumowując, stosy Neuron firmy Amazon opierają się na tym samym fundamencie: vLLM Neuron obsługuje duże modele językowe, podczas gdy vLLM Omni Neuron obsługuje modele dyfuzyjne do generowania obrazów i wideo. Interfejs jądra Neuron umożliwia badaczom pisanie niestandardowych operatorów do wykonywania zarówno w trybie pilnym, jak i skompilowanym, a AWS to NKI-lib o otwartym kodzie źródłowym, zestaw jąder zoptymalizowanych pod kątem Trainium, obsługujący obliczenia blokowe FlashAttention i MoE dla głównego nurtu scenariuszy LLM i multimodalnych.

W zestawie znajdują się również narzędzia zwiększające wydajność. Neuron Explorer rozszerza standardowy plik `torch.profiler' o śledzenie wykonania od platformy po sprzęt, łącząc każde zdarzenie z powrotem z jego linią źródłową i może eksportować metryki w formacie Parquet, które notatki pocztowe mogą wykorzystać agenci do automatycznej analizy wąskich gardeł.

Dostępność i co dalej

Jak wynika z wpisu na blogu, premiera TorchNeuron jako oprogramowania typu open source odbędzie się 20 października na konferencji PyTorch w Ameryce Północnej, a kompatybilność zostanie rozszerzona na Hugging Face Transformers v5. AWS opisuje Trainium jako zestaw chipów zaprojektowanych od podstaw do szkolenia i wnioskowania AI w skali chmury.

Kontekst strategiczny to rynek, na którym każdy liczący się dostawca rozwiązań chmurowych jest jednocześnie największym nabywcą akceleratorów Nvidii i potencjalnym konkurentem. Microsoft, Google i Meta opracowują krzem na zamówienie, a chińskie firmy, w tym Huawei, budują własne stosy oprogramowania, aby uciec przed grawitacją CUDA. Spotykając się z programistami w PyTorch, zamiast prosić ich o migrację, Amazon wybrał najmniej uciążliwą ścieżkę udziału w tym konkursie – i tę, która najprawdopodobniej, jeśli doświadczenie programistów się utrzyma, przeniesie rzeczywiste obciążenia szkoleniowe ze sprzętu Nvidia.

To, czy TorchNeuron rzeczywiście osłabi przewagę CUDA, będzie zależeć od wydajności w skali i długiego ogona modeli w świecie rzeczywistym, na pytania, na które materiały premierowe nie odpowiadają. Jednak bariera, która powstrzymywała większość programistów PyTorch przed wypróbowaniem Trainium — przepisanie — została teraz obniżona do jednej linii kodu.

---

Wyprzedź sztuczną inteligencję

Otrzymuj najnowsze wiadomości, analizy i przełomowe informacje dotyczące sztucznej inteligencji — wszystko w jednym miejscu.

Przeczytaj więcej aktualności o sztucznej inteligencji →