Shardy to system partycjonowania tensorów oparty na MLIR dla wszystkich dialektów. Jest ona wynikiem współpracy zespołów GSPMD i PartIR. Łączy w sobie to, co najlepsze z obu systemów, oraz wspólne doświadczenia zespołów i użytkowników.
Zalety
- Większa kontrola i przewidywalność dla użytkowników dzięki połączeniu propagacji GSPMD z partycjonowaniem przyrostowym PartIR.
- Nowe funkcje oparte na wspólnych doświadczeniach, np. nowe funkcje obsługi zmian kształtu, które generują dodatkowe komunikaty, chyba że użytkownicy wiedzą, jak je obejść.
- Lepsza użyteczność i możliwość debugowania w celu zwiększenia szybkości działania dla użytkowników końcowych, np. przez użycie reprezentacji dzielenia na fragmenty na podstawie osi.
- Prosta, otwarta baza kodu korzystająca z MLIR, z szerszym gronem aktywnych współtwórców (wewnętrznych, zewnętrznych i z różnych stref czasowych), którzy udzielają wsparcia użytkownikom.
Komponenty
- Reprezentacja podziału: reprezentacja podziału na podstawie osi powiązana z konkretną siatką logiczną (z potencjalnie wielu siatek) i obsługująca ograniczanie podziału wymiarów i osi, podział osi na potrzeby operacji takich jak przekształcanie, priorytety dla partycjonowania cząstkowego itp.
- Interfejsy kompilatora: zestaw komponentów kompilatora, które można używać razem z reprezentacją dzielenia na części, aby wpływać na propagowanie dzielenia na części.
- Dzielenie na segmenty danych wejściowych i wyjściowych – dodaj do danych wejściowych lub wyjściowych funkcji głównej funkcję dzielenia na segmenty, aby wskazać, że tensor danych wejściowych i wyjściowych powinien być podzielony na segmenty podczas przekazywania do funkcji lub zwracania z niej.
- Ograniczenie podziału na fragmenty – dołącz podział na fragmenty do pośredniego tensora (np. wyniku matmul), aby wskazać, że tensor lub podzbiór jego zastosowań powinien zostać podzielony na fragmenty.
- Shard As/Like (Utwórz fragmenty jak/podobnie) – zgrupuj wiele tensorów według identyfikatora, aby wskazać, że powinny być podzielone w ten sam sposób.
- Obliczenie ręczne – zawiera podobliczenie, które jest dzielone ręcznie za pomocą podzbioru osi siatki, gdzie podziały na części w przypadku tych osi ręcznych są określone dla wszystkich wejść i wyjść, a wewnątrz tego podobliczenia typy tensorów są lokalne w odniesieniu do tych podziałów.
- Propagation sharding: algorytm propagacji, który łączy priorytety użytkowników i ograniczenia podziału z modelami kosztów kompilatora i heurystycznymi:
- Priorytety zdefiniowane przez użytkownika, np. do równoległego przetwarzania zbiorczego, a potem ZeRO
- priorytety oparte na operacjach, np. najpierw operacje element po elemencie, a potem mnożenia itp.
- Bardziej szczegółowe heurystyki, np. preferowanie wymiarów wsadów.
- SPMD Partitioner: komponent, który zmniejsza liczbę decyzji dotyczących propagowania podziału, dzieląc program na program SPMD, dodając niezbędne operacje przesuwania/formatowania danych i operacje zbiorcze.
- W krótkim czasie początkowa implementacja będzie używać obecnego partycjonera GSPMD SPMD.
- W dłuższej perspektywie planujemy utworzyć nowy moduł partycjonowania SPMD oparty na MLIR.
Repozytorium kodu
Projekt Shardy jest w trakcie opracowywania, a my chcemy uzyskać opinie społeczności open source. Kod Shardy jest dostępny na stronie https://github.com/openxla/shardy.