Model kosztów LHS

tldr;

Na tej stronie opisujemy wewnętrzne działanie modelu kosztów używanego przez harmonogram ukrywania opóźnień. Jeśli chcesz dostroić model, przejdź bezpośrednio do sekcji Dostrajanie.

Harmonogram ukrywania opóźnień (LHS) to etap kompilacji, który planuje DAG HLO w sposób minimalizujący czas rzeczywisty.

Decyzje są podejmowane na podstawie ujednoliconego modelu kosztów, który wykorzystuje kombinację tabel wydajności i modeli analitycznych. W szczególności XLA zawiera tabele wydajności dla GEMM i szybkich połączeń zbiorowych oraz wykorzystuje analityczny model kosztów sieci i fuzji w innych przypadkach. W pozostałej części dokumentu opisano ogólne zasady działania tych funkcji.

Tabele skuteczności – organizacje zbiorowego zarządzania prawami autorskimi ICI

Tabela wydajności składa się z 2 głównych komponentów: modułu zbierającego i interpolatora.

Kolektor

Kolektor to narzędzie w C++ odpowiedzialne za generowanie tabel wydajności dla operacji zbiorczych. Mierzy skuteczność poszczególnych operacji HLO (np. all-gather, all-reduce) w statycznie zdefiniowanej przestrzeni parametrów.

Jak to działa

Narzędzie przeprowadza testy w zakresie operacji zbiorczych, rozmiarów transferu i schematów transferu dla danego klastra. Korzysta z infrastruktury narzędzia do uruchamiania HLO na wielu hostach i danych ExecutionProfile, aby uruchamiać wygenerowane HLO i zbierać dane o wydajności.

Parametry zbierania danych

Tabele opóźnień są zbierane dla kombinacji tych parametrów:

  • Collective Type:
    • all-reduce
    • all-gather
    • reduce-scatter
  • Rozmiar transferu:
    • Skala logarytmiczna od 1024 B do 2 GiB (np. 1024 B, 2048 B, 4096 B, ...)
  • Program przenoszenia:
    • rail-aligned
    • non-rail-aligned

Ten test jest przeprowadzany w przypadku klastrów wewnątrz węzła z 2, 4 i 8 urządzeniami.

Wyniki

Wynikiem uruchomienia zbierania danych jest tabela czasu oczekiwania w formacie .pbtxt (około 116 KB na platformę).

Interpolator

Interpolator to komponent kompilatora, który wykorzystuje wygenerowane tabele wydajności, aby podczas kompilacji podawać szacunki czasu działania.

Wewnętrzna struktura danych

Podczas inicjowania interpolator przetwarza tabelę wydajności na mapę. Ta mapa używa krotki (collective_type, transfer_scheme) jako klucza.

Wartość powiązana z każdym kluczem to 2-wymiarowa płaszczyzna euklidesowa. Ta płaszczyzna indeksuje przepustowość sieci (mierzoną przez moduł zbierający) na podstawie 2 osi:

  1. Rozmiar przesłanych danych.
  2. Liczba urządzeń, których dotyczy problem.

Wyszukiwanie i interpolacja

Gdy kompilator napotka operację zbiorową, Interpolator wykonuje te czynności:

  1. Identyfikuje prawidłową płaszczyznę przepustowości 2D, używając (collective_type, transfer_scheme) operacji jako klucza mapy.
  2. Następnie w ramach tej płaszczyzny 2D używa wyszukiwania średniej ważonej (na podstawie odległości euklidesowej), a jako punktu zapytania używa (transfer_size, num_devices) operacji.
  3. Wynikiem tego wyszukiwania jest pojedyncza, unikalna wartość przepustowości sieci.

Uzasadnienie: przepustowość i ekstrapolacja

System jest przeznaczony do przechowywania przepustowości sieci, a nie surowych danych o opóźnieniach. Ten wybór projektu znacznie upraszcza ekstrapolację wydajności w przypadku rozmiarów transferu, które nie są wyraźnie widoczne w tabeli.

Jeśli tabele opóźnień rejestrują nasycenie przepustowości sieci przy łącznym rozmiarze S, przepustowość T w tym momencie jest uważana za maksymalną. W przypadku każdej nowej grupy o rozmiarze S' > S czas działania można oszacować w ten sposób:

\[\text{EstimatedTime}(S') = \frac{S'}{T_{\text{saturated} } }\]

Dzięki temu model może szacować wydajność kolekcji o dowolnym rozmiarze, nawet tych większych niż 2 GiB, czyli maksymalny rozmiar mierzony przez moduł zbierający.

  • Zaniż maksymalną przepustowość.
  • Dlatego w przypadku dużych transferów zawyżaj czas działania.

Zespoły XLA:GPU zwykle prowadzą tabele wydajności, ale jeśli użytkownik zdecyduje się podać własne, to on jest odpowiedzialny za to, aby były one reprezentatywne i zawierały pomiary w regionie nasyconym przepustowością dla docelowego sprzętu.

Tabele skuteczności – GEMM

Podobnie jak w przypadku zbiorów, tabele opóźnień GEMM są obsługiwane przez 2 komponenty: kolektor i interpolator.

Kolektor

Kolektor to narzędzie w C++, które oblicza tabele wydajności dla ogólnych mnożeń macierzy (GEMM). Mierzy wydajność mnożenia macierzy na poziomie operacji HLO dot.

Jak to działa

Narzędzie wykonuje skanowanie statycznej przestrzeni wymiarów GEMM (wymiar wsadowy, 2 wymiary niekurczliwe i 1 wymiar kurczliwy) oraz typów danych.

  • Domyślne typy danych: LHS = bf16,f32, RHS = bf16,f32, OUT = bf16,f32.
  • Infrastruktura: ponownie wykorzystuje profiler operacji HLO.

Parametry zbierania

Tabele opóźnień są zbierane dla kombinacji tych wymiarów:

  • batch: {1, 2, 4}
  • m (non-contracting): {256, 512, ..., 4096}
  • n (non-contracting): {256, 512, ..., 4096}
  • k (contracting): {256, 512, ..., 4096}

Wyniki i miejsce na dane

Pełne skanowanie generuje tabelę opóźnień .pbtxt, która jest gotowa do użycia przez interpolator.

Interpolator

Interpolator to komponent kompilatora, który używa wygenerowanych tabel do szacowania wydajności GEMM.

Uzasadnienie: nasycenie FLOPS

Zebrane tabele opóźnień umożliwiają interpolatorowi odtworzenie FLOPS dla każdego wpisu:

\[\text{FLOPS} = \frac{2 \times b \times m \times n \times k}{\text{runtime} }\]

Kluczową obserwacją jest to, że FLOPS osiągają punkt nasycenia, czyli sprzęt osiąga maksymalną liczbę FLOPS powyżej określonego kształtu macierzy. Ten poziom nasycenia umożliwia stosowanie tej samej metody ekstrapolacji, która jest używana w przypadku zbiorowości.

Wyszukiwanie i interpolacja

Interpolator tworzy na podstawie danych z tabeli 4-wymiarową przestrzeń euklidesową. Aby podać szacunkową skuteczność, wykonuje interpolację średniej ważonej w tej 4-wymiarowej przestrzeni. Jeśli dla danego typu danych nie ma tabeli, każdy wymiar jest normalizowany do liczby bajtów.

Analityczny model kosztów – DCN

Model kosztów zbiorczych w kształcie litery S

Model krzywej S to w pełni analityczny model limitu sieci.

Przegląd

Model ten ma na celu oszacowanie wydajności operacji zbiorowych na podstawie zestawu stałych właściwości sieci.

Dane wejściowe modelu

Model wymaga 2 kategorii danych wejściowych:

  1. Właściwości sieci stacjonarnej (zdefiniowane przez użytkownika):

    • Dodatkowe obciążenie związane z uruchamianiem zbiorczym
    • Szybkość karty sieciowej
    • RTT (czas błądzenia)

    Domyślnie XLA automatycznie wykrywa platformę i używa wartości dla najpopularniejszych architektur. Użytkownik może skonfigurować te właściwości. Szczegółowe informacje znajdziesz w sekcji dotyczącej dostrajania.

  2. Dane wejściowe dla poszczególnych grup:

    • Typ zbiorowy (np. AllGather, ReduceScatter)
    • Rozmiar przesłanych danych
    • Liczba węzłów biorących udział w komunikacji

Integracja

Model krzywej S jest zintegrowany z XLA:GPU i jest używany w przypadku architektur Hopper i Blackwell.

Analityczny model kosztów – fuzje

W przypadku innych jąder korzystamy z modelu kosztów wydajności GPU, aby oszacować odpowiednie czasy działania. Więcej informacji na ten temat znajdziesz tutaj.

Dostrajanie

Model krzywej S można dostroić, wydając odpowiednie flagi XLA. Konfiguracja domyślna powinna być wystarczająca w większości przypadków, ale w innych przypadkach kontrola modelu jest udostępniana.

export NIC_SPEED_GBPS=... # NIC speed per GPU in Gigabytes
export GPUS_PER_NODE=... # Num of GPUs per cluster interconnected with fast network (e.g. NVLINK)
export XLA_FLAGS=--xla_gpu_analytical_latency_estimator_options="nic_speed_gbps=$NIC_SPEED_GBPS,gpus_per_node=$GPUS_PER_NODE"