Tensor Parallelism

Tensor Parallelism (TP) dzieli tensory i obliczenia pojedynczych warstw modelu między urządzenia. GPU liczą części tej samej operacji, a wyniki łączą zgodnie z wybranym podziałem. Pozwala to rozłożyć pamięć wag i pracę obliczeniową; wymaga też komunikacji.

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.

Wykorzystuję treści generowane przez AI jako część mojego codziennego procesu nauki.