Du är på en ML-ingenjörsintervju på Google. Intervjuare: Vi behöver träna en LLM över 1 000 GPU:er. Hur skulle du se till att alla GPU:er delar med sig av det de lär sig? Du: Använd en central parameterserver för att aggregera och omfördela vikterna. Intervjun är över. Här är vad du missade:

En stor flaskhals vid körning i träning med flera GPU:er inträffar under GPU-synkroniseringen. Till exempel i multi-GPU-träning via dataparallellitet: - Samma modell distribueras till olika GPU:er. - Varje GPU bearbetar en annan delmängd av hela datauppsättningen. Kolla in det här 👇

Detta leder till olika gradienter mellan olika enheter. Så, innan vi uppdaterar modellparametrarna på varje GPU-enhet, måste vi kommunicera gradienterna till alla andra enheter för att synkronisera dem. Låt oss förstå två vanliga strategier härnäst!

Algoritm 1) All-reduce Ett uppenbart sätt är att skicka gradienterna från en enhet till alla andra enheter för att synkronisera dem. Men detta utnyttjar hög bandbredd. Om varje GPU har "N"-element och det finns "G"-GPU:er, resulterar det i en total överföring av G*(G-1)*N element 👇

Vi kan optimera detta genom att överföra alla element till en GPU, beräkna slutvärdet och skicka tillbaka det till alla andra GPU:er. Detta är en betydande förbättring. Men nu måste en enda GPU ta emot, beräkna och kommunicera tillbaka gradienterna, så detta skalar inte.

Algoritm 2) Ringreduktion Fas #1) Minskad andel Först delas gradienterna in i G-segment på varje GPU (G = totalt antal GPU:er). Kolla det här 👇

I en iteration skickar varje GPU ett segment till nästa GPU: - GPU1 skickar a₁ till GPU2, där det läggs till b₁ - GPU2 skickar b₂ till GPU3, där det läggs till c₂ - GPU3 skickar c₃ till GPU4, där det läggs till d₃ - GPU4 skickar d₄ till GPU1, där det läggs till a₄ Kolla det här 👇

Denna process görs igen: - GPU1 skickar (d₄+a₄) till GPU2, där det läggs till b₄. - GPU2 skickar (a₁+b₁) till GPU3, där den läggs till c₁. - GPU3 skickar (b₂+c₂) till GPU4, där den läggs till d₂. - GPU4 skickar (c₃+d₃) till GPU1, där den läggs till a₃. Kolla det här 👇

I den slutliga iterationen överförs följande segment till nästa GPU. Detta leder till ett tillstånd där varje GPU har ett helt segment, och vi kan överföra dessa kompletta segment till alla andra GPU:er. Kolla det här 👇

Fas #2) Endast delning Nu när varje GPU har ett helt segment kan vi överföra dessa kompletta segment till alla andra GPU:er. Processen genomförs på liknande sätt som vi diskuterade ovan, så vi går inte in på alla detaljer. Iteration 1 visas nedan👇

Därefter utför vi iterationerna 2 och 3. Dessa genomförs på samma sätt som vi lärde oss i fas 1. Kolla in det här 👇

Och där har du det! Modellvikter över GPU:er har synkroniserats. Även om det totala antalet överförda element fortfarande är desamma som vi hade i "single-GPU-master"-metoden, är denna ringmetod mycket mer skalbar eftersom den inte lägger hela belastningen på ett GPU. Kolla det här 👇

Som standard använder djupinlärningsmodeller endast ett enda GPU för träning, även om flera GPU:er finns tillgängliga. Ett idealiskt sätt att träna modeller är att fördela träningsbelastningen över flera GPU:er. Grafiken visar fyra strategier för multi-GPU-träning👇