Tesla chce wprowadzić system uczenia maszynowego do Dojo

Aby zaspokoić pragnienie coraz większych modeli sztucznej inteligencji i uczenia maszynowego, Tesla ujawniła na Hot Chips 34 wiele szczegółów dotyczących ich w pełni niestandardowej architektury superkomputerowej zwanej Dojo.

System jest zasadniczo ogromnym superkomputerem, który można komponować, chociaż w przeciwieństwie do tego, co widzimy na Top 500, jest zbudowany z całkowicie niestandardowej architektury, która obejmuje układy obliczeniowe, sieciowe i wejścia/wyjścia (I/O), architekturę zestawu instrukcji (ISA), zasilanie, pakowanie i chłodzenie. Wszystko to zostało zrobione z wyraźnym celem uruchomienia dostosowanych, konkretnych algorytmów szkoleniowych uczenia maszynowego na dużą skalę.

„Przetwarzanie danych w świecie rzeczywistym jest możliwe tylko dzięki technikom uczenia maszynowego, czy to przetwarzaniu języka naturalnego, poruszaniu się po ulicach stworzonych z myślą o ludzkim wzroku, czy robotyce łączącej się z codziennym środowiskiem” – powiedział Ganesh Venkataramanan, starszy dyrektor ds. inżynierii sprzętu w Tesli. podczas jego przemówienia programowego.

Twierdził jednak, że tradycyjne metody skalowania rozproszonych obciążeń nie przyspieszyły w tempie niezbędnym do nadążania za wymaganiami uczenia maszynowego. W efekcie prawo Moore'a go nie ogranicza, podobnie jak systemy do szkolenia AI/ML na dużą skalę, a mianowicie niektóre kombinacje CPU/GPU lub, w rzadszych przypadkach, przy użyciu specjalnych akceleratorów AI. 

„Tradycyjnie budujemy chipy, umieszczamy je na opakowaniach, opakowania trafiają na płytki PCB, które trafiają do systemów. Systemy trafiają do szaf” — powiedział Venkataramanan. Problem polega na tym, że za każdym razem, gdy dane są przenoszone z chipa do pakietu i poza pakiet, wiąże się to z opóźnieniem i karą przepustowości.

Kanapka z centrum danych

Aby więc ominąć ograniczenia, Venkataramanan i jego zespół zaczęli od zera.

„Zaraz w moim wywiadzie z Elonem zapytał mnie, co możesz zrobić, co różni się od procesorów i kart graficznych dla sztucznej inteligencji. Czuję, że cały zespół wciąż odpowiada na to pytanie”.

Doprowadziło to do opracowania kafelka treningowego Dojo, samodzielnego klastra obliczeniowego zajmującego pół stopy sześciennej, zdolnego do 556 TFLOPS wydajności FP32 w pakiecie chłodzonym cieczą o mocy 15 kW.

Każda płytka jest wyposażona w 11 GB pamięci SRAM i jest połączona siecią szkieletową 9 TB/s przy użyciu niestandardowego protokołu transportowego w całym stosie.

„Ten kafelek treningowy reprezentuje niezrównane ilości integracji od komputera do pamięci, dostarczania zasilania i komunikacji, bez konieczności stosowania dodatkowych przełączników” – powiedział Venkataramanan.

Sercem kafelka treningowego jest D1 Tesli, 50-miliardowa matryca tranzystorowa, oparta na 7-nanometrowym procesie TSMC. Tesla twierdzi, że każdy D1 jest zdolny do 22 TFLOPS wydajności FP32 przy TDP 400 W. Jednak Tesla zauważa, że ​​chip jest w stanie wykonać szeroki zakres obliczeń zmiennoprzecinkowych, w tym kilka niestandardowych.

„Jeśli porównasz tranzystory do milimetra kwadratowego, to prawdopodobnie jest to krwawa krawędź wszystkiego, co tam jest”, powiedział Venkataramanan.

Tesla następnie wziął 25 modułów D1, pogrupował je w znane dobre matryce, a następnie spakował je przy użyciu technologii TSMC system-on-wafer, aby „osiągnąć ogromną integrację obliczeniową przy bardzo małych opóźnieniach i bardzo dużej przepustowości” – powiedział.

Jednak konstrukcja typu system-on-wafer i architektura stosu pionowego wprowadziły wyzwania, jeśli chodzi o dostarczanie energii.

Według Venkataramanan, większość dzisiejszych akceleratorów umieszcza moc bezpośrednio w sąsiedztwie krzemu. I choć takie podejście jest sprawdzone, oznacza to, że duża część akceleratora musi być dedykowana tym komponentom, co czyniło go niepraktycznym dla Dojo – wyjaśnił. Zamiast tego Tesla zaprojektowała swoje chipy, aby dostarczać energię bezpośrednio przez dno kości. 

Kładąc wszystko razem

„Możemy zbudować całe centrum danych lub cały budynek z tego kafelka szkoleniowego, ale kafelek szkoleniowy to tylko część obliczeniowa. Musimy go również nakarmić” – powiedział Venkataramanan.

W tym celu Tesla opracowała również procesor interfejsu Dojo (DIP), który działa jako pomost między procesorem hosta a procesorami treningowymi. DIP służy również jako źródło współużytkowanej pamięci o dużej przepustowości (HBM) i jako szybka karta sieciowa 400 Gb/s.

Każdy DIP zawiera 32 GB pamięci HBM, a do pięciu z tych kart można podłączyć do płytki szkoleniowej z prędkością 900 GB/s, co daje łącznie 4.5 TB/s do hosta, co daje łącznie 160 GB pamięci HBM na płytkę.

Pary konfiguracyjne Tesli V1 tych płytek – lub 150 matryc D1 – w macierzy obsługują cztery procesory hosta wyposażone w pięć kart DIP, aby osiągnąć deklarowaną wydajność BF16 lub CFP8.

Podsumowując, Venkataramanan mówi, że architektura – szczegółowa dogłębnie tutaj by Następna platforma – umożliwia Tesli przezwyciężenie ograniczeń związanych z tradycyjnymi akceleratorami pokroju Nvidii i AMD.

„Jak działają tradycyjne akceleratory, zazwyczaj próbuje się dopasować cały model do każdego akceleratora. Replikuj go, a następnie przesyłaj dane przez każdy z nich” – powiedział. „Co się stanie, jeśli będziemy mieli coraz większe modele? Te akceleratory mogą spaść, ponieważ zabraknie im pamięci”.

Zauważył, że to nie jest nowy problem. Na przykład przełącznik NVidii umożliwia łączenie pamięci w dużych bankach procesorów graficznych. Jednak Venkataramanan twierdzi, że to nie tylko zwiększa złożoność, ale także wprowadza opóźnienia i kompromisy w zakresie przepustowości.

„Myśleliśmy o tym od samego początku. Nasze płytki obliczeniowe i każda z matryc zostały wykonane z myślą o dopasowaniu dużych modeli” – powiedział Venkataramanan.

Tworzenie

Tak wyspecjalizowana architektura obliczeniowa wymaga specjalistycznego stosu oprogramowania. Jednak Venkataramanan i jego zespół uznali, że programowalność albo stworzy, albo zniszczy Dojo.

„Łatwość programowania dla odpowiedników oprogramowania jest najważniejsza, gdy projektujemy te systemy”, powiedział. „Naukowcy nie będą czekać, aż twoi ludzie od oprogramowania napiszą odręcznie napisane jądro w celu dostosowania się do nowego algorytmu, który chcemy uruchomić”.

Aby to zrobić, Tesla porzucił pomysł używania jądra i zaprojektował architekturę Dojo wokół kompilatorów.

„To, co zrobiliśmy, to użycie PiTorch. Stworzyliśmy warstwę pośrednią, która pomaga nam w równoległym skalowaniu sprzętu znajdującego się pod nią. Pod wszystkim jest skompilowany kod” – powiedział. „To jedyny sposób na tworzenie stosów oprogramowania, które można dostosować do wszystkich przyszłych obciążeń”.

Pomimo nacisku na elastyczność oprogramowania, Venkataramanan zauważa, że ​​platforma, która obecnie działa w ich laboratoriach, jest na razie ograniczona do użytku Tesli.

„Najpierw koncentrujemy się na naszych wewnętrznych klientach” – powiedział. „Elon podał do wiadomości publicznej, że z czasem udostępnimy to naukowcom, ale nie mamy na to ram czasowych. ®