Jak profilować TorchNeuron na AWS Trainium: poradnik część 1

Amazon opublikowało pierwszą część serii poświęconej profilowaniu wydajności modeli na AWS Trainium — autorskim chipie AI firmy przeznaczonym do treningu i inferencji modeli głębokiego uczenia. Materiał, przygotowany przez Josha Longeneckera, tłumaczy od podstaw, jak systematycznie analizować wydajność inferencji w TorchNeuron, natywnym backendzie PyTorch dla Trainium, i jak zlokalizować, gdzie dokładnie model traci czas obliczeniowy.

Artykuł nie zakłada wcześniejszego doświadczenia w profilowaniu i krok po kroku pokazuje konfigurację narzędzi, odczyt wyników oraz analizę trzech różnych implementacji tej samej operacji — RMSNorm.

Kluczowe wnioski

• TorchNeuron pozwala uruchamiać standardowe przepływy PyTorch (eager execution, torch.compile, standardowe API) na AWS Trainium przy minimalnych zmianach w kodzie, a integracja z Neuron Kernel Interface (NKI) umożliwia tworzenie niestandardowych operatorów i optymalizację pod sprzęt.
• Backend Neuron wykorzystuje mechanizm rozszerzeń PyTorch PrivateUse1, dzięki czemu akcelerator integruje się z frameworkiem bez potrzeby dodawania dedykowanego klucza dyspatchu, jak w przypadku CUDA.
• Do profilowania służą dwa obiekty: NeuronConfig, który określa, jakie strumienie danych i NeuronCore’y mają być rejestrowane, oraz NeuronProfiler, który eksportuje ślad frameworka (trace.json) wraz z systemowymi śladami Neurona.
• Neuron Explorer to przeglądarkowe narzędzie do wizualizacji śladu systemowego, osi czasu urządzenia i podziału na poszczególne silniki obliczeniowe — pozwala prześledzić pojedynczą operację od dyspatchu w frameworku, przez harmonogramowanie w runtime, aż po wykonanie na sprzęcie NeuronCore.
• Autorzy porównują trzy implementacje RMSNorm — naiwną w czystym PyTorch, wbudowaną funkcję torch.rms_norm() oraz naiwną wersję skompilowaną przez torch.compile — żeby pokazać, jak zmieniają się sygnały wydajnościowe w zależności od sposobu implementacji.

Jak Neuron integruje się z PyTorch

Zanim przystąpi się do profilowania, warto zrozumieć mechanizm, który łączy Neuron z frameworkiem PyTorch. Backend Neuron opiera się na punkcie rozszerzeń PrivateUse1, który pozwala akceleratorowi zintegrować się z PyTorch bez dodawania osobnego klucza dyspatchu (jak dzieje się to w przypadku cuda). Pakiet torch_neuronx współpracuje z główną gałęzią PyTorch i rejestruje Neuron jako natywny backend dla kodu PyTorch. Dzięki temu programiści mogą pisać kod w znanym stylu, przenosząc tensory i operacje na urządzenie neuron, tak jak robiliby to z GPU.

Konfiguracja profilera i narzędzie Neuron Explorer

Do profilowania obciążenia na Neuronie potrzebne są dwa obiekty. Pierwszy, NeuronConfig, kontroluje, które strumienie danych i NeuronCore’y mają zostać przechwycone oraz gdzie zapisać wynik — przekazuje się go do torch.profiler.profile jako experimental_config. Drugi, NeuronProfiler, przyjmuje NeuronConfig i udostępnia callback export_trace, który zapisuje ślad frameworka (trace.json) obok systemowych śladów Neurona.

Parametr modes wybiera strumienie danych do przechwycenia, np. instrukcje urządzenia czy ślady runtime, a capture_enabled_for_nc wskazuje, które NeuronCore’y mają być śledzone. Ustawienie correlate_device_time=True sprawia, że czas urządzenia jest propagowany do tabeli key_averages(), dzięki czemu można bezpośrednio porównać czas CPU z czasem na urządzeniu Neuron.

Do wizualnej analizy wyników służy Neuron Explorer — przeglądarkowe narzędzie pokazujące ślad systemowy, oś czasu urządzenia oraz podział na poszczególne silniki obliczeniowe. Dostęp do niego można uzyskać przez tunel SSH lub rozszerzenie Visual Studio Code. Po zakończeniu profilowania katalog wyjściowy zawiera ślad systemowy, a przy przechwytywaniu danych urządzenia — sparowane pliki .neff i .ntff. Plik NEFF (Neuron Executable File Format) zawiera skompilowany kod wykonywalny Neuron, a odpowiadający mu NTFF — ślad instrukcji sprzętowych. Oba pliki łączy wspólny hash numeryczny identyfikujący je jako parę. Do analizy pełnego profilu systemowego wgrywa się cały katalog wyjściowy, a dla profilu dotyczącego wyłącznie urządzenia — pojedynczo pliki .neff i .ntff.

Profilowanie RMSNorm: trzy implementacje tej samej operacji

Jako pierwszy przykład analizy autorzy wybrali RMSNorm — niewielką operację o jasnej sekwencji działań arytmetycznych i redukcji, co ułatwia prześledzenie przepływu obliczeń. Profilowane są trzy warianty:

• Naiwna implementacja RMSNorm złożona z osobnych operacji PyTorch: mnożenia, uśredniania, dodawania, odwrotności pierwiastka kwadratowego i kolejnego mnożenia.
• Wbudowana funkcja torch.rms_norm(), która jest redukowana do pojedynczego, zintegrowanego kernela Neuron.
• Ta sama naiwna implementacja co w pierwszym wariancie, ale opakowana w torch.compile.

Wszystkie skrypty korzystają z identycznych danych wejściowych, tej samej konfiguracji profilera i tego samego harmonogramu, dzięki czemu wyniki są bezpośrednio porównywalne. Pierwszą implementację autorzy pokazują na konkretnym fragmencie kodu: tensor wejściowy jest mnożony przez samego siebie, następnie liczona jest średnia kwadratów wzdłuż ostatniego wymiaru, potem odwrotność pierwiastka z sumy średniej kwadratów i epsilonu (1e-6), a na końcu wynik jest mnożony przez wagę. Każda linia tej funkcji odpowiada jednej lub kilku operacjom ATen i może tworzyć tensory pośrednie, co ma znaczenie przy analizie zużycia pamięci i czasu.

Wywołanie profilera wymaga dwóch dodatkowych argumentów specyficznych dla backendu Neuron: aktywności ProfilerActivity.PrivateUse1 do zbierania danych z urządzenia Neuron oraz przekazania NeuronConfig jako experimental_config, co włącza funkcje Kineto właściwe dla Neurona. Harmonogram profilowania wykorzystuje parametry wait=2, warmup=1 i active=1 w czterech wywołaniach prof.step(), co pozwala zarejestrować jedną „rozgrzaną” iterację bez wliczania narzutu związanego z początkową inicjalizacją. Kluczowe jest też wywołanie torch.neuron.synchronize() wewnątrz pętli profilującej — bez tej synchronizacji profiler może wyeksportować dane, zanim urządzenie zakończy obliczenia, przez co kolumna czasu Neurona pozostanie pusta.

Artykuł kończy się zapowiedzią dalszej analizy wyników dla poszczególnych implementacji, lecz tekst źródłowy zostaje w tym miejscu skrócony — pełna interpretacja tabeli key_averages() i porównanie sygnałów w Neuron Explorer dla wszystkich trzech wariantów RMSNorm ma zostać omówiona w dalszej części materiału.

Pierwsza część serii o profilowaniu TorchNeuron dostarcza więc solidnych podstaw: pokazuje, jak skonfigurować profiler na Trainium, jak czytać kluczowe metryki wydajnościowe oraz jak przygotować porównywalne eksperymenty na przykładzie RMSNorm — fundament pod bardziej zaawansowaną analizę wydajności modeli na infrastrukturze AWS.