W poście opublikowanym na platformie Hugging Face, Amazon rozpoczyna serię poświęconą analizie wydajności inferencji, czyli wykonywania już wytrenowanego modelu. Ta pierwsza część wyjaśnia, jak zidentyfikować, gdzie tracony jest czas, przy użyciu narzędzi profilowania zastosowanych do układu Trainium od AWS, zaprojektowanego do trenowania i inferencji modeli głębokiego uczenia.
Przewodnik skierowany jest do czytelników bez wcześniejszego doświadczenia w profilowaniu. Opiera się na TorchNeuron, natywnym, open-source’owym backendzie PyTorcha dla Trainium, który obsługuje standardowe przepływy pracy PyTorcha, w tym wykonanie w trybie eager i torch.compile, przy niewielkich modyfikacjach kodu. TorchNeuron integruje się również z Neuron Kernel Interface (NKI), która umożliwia tworzenie niestandardowych operatorów.
Konfiguracja profilera
Profilowanie opiera się na dwóch obiektach. NeuronConfig definiuje strumienie danych do przechwycenia, NeuronCores do śledzenia oraz folder wyjściowy, a następnie jest przekazywany do torch.profiler.profile. NeuronProfiler dostarcza funkcję eksportu, która zapisuje ślad frameworka, w formacie trace.json, obok śladów systemowych Neuron. Parametr correlate_device_time=True pozwala wyświetlić czas spędzony na urządzeniu w tabeli key_averages(), gdzie można go porównać z czasem CPU.
Ślady można następnie otworzyć w Neuron Explorer, narzędziu wizualizacyjnym działającym w przeglądarce. Pozwala ono śledzić operację od podziału w frameworku aż do jej wykonania na sprzęcie, przez harmonogramowanie w środowisku wykonawczym. Dla NeuronCore, czyli rdzenia obliczeniowego układu, plik NEFF zawiera skompilowany plik wykonywalny, podczas gdy odpowiadający mu plik NTFF zawiera ślad instrukcji sprzętowych.
Trzy wersje tej samej operacji
Post analizuje funkcję RMSNorm, prostą operację składającą się z sekwencji obliczeń arytmetycznych i redukcji. Porównywane są trzy implementacje: wersja napisana operacja po operacji w PyTorchu, wywołanie wbudowanej funkcji torch.rms_norm() oraz wersja napisana operacja po operacji, ale skompilowana za pomocą torch.compile.
W wersji naiwnej tabela pokazuje, że trzy wywołania aten::mul stanowią największą część czasu spędzonego na urządzeniu. Całkowity czas Neuron osiąga 2 355 µs dla sześciu wywołań operatorów. Według autora, wbudowana funkcja torch.rms_norm(), która jest przekształcana w pojedyncze, zfuzowane jądro obliczeniowe, zajmuje tylko jedną operację, aten::_fused_rms_norm, z czasem Neuron wynoszącym 403,3 µs. Wersja skompilowana za pomocą torch.compile pojawia się jako pojedynczy, skompilowany region, o czasie trwania 375,5 µs.
W poście zaznaczono, że w tych testach oba podejścia zajmują około jednej szóstej czasu obliczeniowego wersji naiwnej. Fuzja eliminuje przerwy między wykonaniami sprzętowymi i unika materializowania wartości pośrednich, takich jak x_sq i mean_sq, między operacjami wysyłanymi osobno.
Autor wskazuje, że metodę można zastosować do większych modeli: zidentyfikować wąskie gardło w tabeli, zweryfikować je w śladzie, a następnie określić, czy zfuzowana operacja lub torch.compile może zmniejszyć jej koszt. Druga część serii będzie dotyczyć profilowania na poziomie urządzenia, z analizą pliku NEFF.
