Tensor Parallelism

Model nie mieści się na jednej karcie obliczeniowej albo jedna operacja jest zbyt duża. Chcesz, żeby kilka urządzeń pracowało nad częściami tego samego obliczenia, zamiast każde liczyło pełną kopię.

Tensor Parallelism dzieli tabele liczb i operacje pojedynczych warstw między urządzenia. Tensor to wielowymiarowy zbiór liczb; w prostym przypadku może być wektorem lub macierzą. Karty GPU obliczają swoje części, a następnie przekazują potrzebne wyniki.

Jeśli jedna karta liczy wkład pierwszej połowy wejścia, a druga drugiej, ich wyniki mogą wymagać sumowania. Nie muszą to być dwie gotowe połowy zdania dla użytkownika.

Podział zmniejsza część obciążenia pamięci i obliczeń na urządzenie, lecz dodaje komunikację. Wolne połączenie może ograniczyć zysk. Pipeline Parallelism dzieli kolejne grupy warstw, a Tensor Parallelism wnętrze ich operacji.

Mechanizm i szczegóły

Megatron-LM, §3 opisuje taki podział dla Position-wise Feed-Forward Network oraz Multi-Head Attention. Częściowy wynik nie musi być gotowym fragmentem odpowiedzi. Czasem trzeba go zsumować z wynikiem drugiego GPU, zanim będzie można wykonać następną operację.

Dwa GPU, jedna suma

We własnym przykładzie liczymy $y=xW$. Dzielimy cztery współrzędne wejścia na dwie pary, a macierz wag na odpowiadające im wiersze:

$$x=[1,2,-1,3],\qquad W=\begin{bmatrix}2&-1\\1&2\\3&0\\-2&1\end{bmatrix}$$

GPU A liczy wkład pierwszych dwóch wierszy: $[4,3]$. GPU B liczy pozostałe: $[-9,3]$. Ich suma to $[-5,6]$, dokładnie jak przy mnożeniu całej macierzy. Sumowanie między urządzeniami można zrealizować operacją all-reduce; każde otrzymuje wtedy sumę.

Teraz dodajmy ReLU, czyli zamianę ujemnych współrzędnych na zero. Po zsumowaniu dostajemy $[0,6]$. Jeśli oba GPU zastosują ReLU przed sumą, wynik będzie błędny: $[4,6]$. Zmień ostatnią współrzędną wejścia i sprawdź obie kolejności.

Podział musi pasować do dalszego obliczenia

To przykład jednego mnożenia i nieliniowości, bez biasu, uczenia ani modelowania czasu GPU. ReLU ułatwia rachunek ręczny; występuje w oryginalnym Transformerze (Vaswani et al., §3.3). Megatron-LM używa GeLU i dobiera podział kolumn pierwszej macierzy oraz wierszy drugiej tak, by ograniczyć synchronizację. Przy podziale kolumn nieliniowość można liczyć lokalnie, bo powstają różne współrzędne wyniku.

Pipeline Parallelism rozdziela kolejne grupy warstw; TP rozdziela pracę wewnątrz warstw. Można je łączyć. Więcej GPU nie gwarantuje proporcjonalnego przyspieszenia: mniejsze operacje i komunikacja mogą zdominować zysk, co pokazuje Megatron-LM, dodatek D.2.