Wskazówki dotyczące flag XLA

Ten przewodnik zawiera wyselekcjonowane najważniejsze flagi XLA, które pomogą użytkownikom skutecznie korzystać z funkcji XLA. W sekcjach poniżej znajdziesz szczegółowe informacje o flagach, które mogą mieć znaczący wpływ na wydajność w czasie działania i wykorzystanie pamięci. Jeśli po włączeniu flagi wystąpią problemy, np. awarie, zalecamy przywrócenie ustawień domyślnych i utworzenie zgłoszenia w GitHub.

Oznaczenia wydajności

Te flagi mają kluczowe znaczenie dla zwiększenia wydajności w czasie działania. Eksperymentowanie z tymi ustawieniami może przynieść znaczne korzyści w zakresie skuteczności.

Flaga Opis Wartości domyślne Sugerowane wartości Wartości kandydatów
Potokowanie
1. xla_should_allow_loop_variant_parameter_in_chain
2. xla_should_add_loop_invariant_op_in_chain
3. xla_tpu_enable_ici_ag_pipelining
Te 3 flagi należy stosować łącznie, aby włączyć zbiorcze przetwarzanie potokowe operacji all-gather ICI(Interchip-Interconnect), co stwarza więcej możliwości nakładania się wykonywania. 1. xla_should_allow_loop_variant_parameter_in_chain=kDisabled
2. xla_should_add_loop_invariant_op_in_chain=kDisabled
3. xla_tpu_enable_ici_ag_pipelining=false
1. xla_should_allow_loop_variant_parameter_in_chain=kEnabled
2. xla_should_add_loop_invariant_op_in_chain=kEnabled
3. xla_tpu_enable_ici_ag_pipelining=true
1. xla_should_allow_loop_variant_parameter_in_chain=kDisabled/kEnabled/kAuto
2. xla_should_add_loop_invariant_op_in_chain=kDisabled/kEnabled/kAuto
3. xla_tpu_enable_ici_ag_pipelining=true/false
v5e/Async
xla_enable_async_all_gather
xla_tpu_enable_async_collective_fusion
xla_tpu_enable_async_collective_fusion_fuse_all_gather
Te 3 flagi należy stosować łącznie, aby aktywować asynchroniczne operacje all-gather na v5e. xla_enable_async_all_gather=kAuto
xla_tpu_enable_async_collective_fusion=true
xla_tpu_enable_async_collective_fusion_fuse_all_gather=true
xla_enable_async_all_gather=kAuto
xla_tpu_enable_async_collective_fusion=true
xla_tpu_enable_async_collective_fusion_fuse_all_gather=true
xla_enable_async_all_gather=kDisabled/kEnabled/kAuto
xla_tpu_enable_async_collective_fusion=true/false
xla_tpu_enable_async_collective_fusion_fuse_all_gather=true/false
v5e/Async
xla_tpu_enable_async_collective_fusion
xla_tpu_enable_async_collective_fusion_fuse_all_reduce
Te 2 flagi należy stosować razem, aby aktywować asynchroniczne operacje all-reduce na v5e. xla_tpu_enable_async_collective_fusion=true
xla_tpu_enable_async_collective_fusion_fuse_all_reduce=false
xla_tpu_enable_async_collective_fusion=true
xla_tpu_enable_async_collective_fusion_fuse_all_reduce=true
xla_tpu_enable_async_collective_fusion=true/false
xla_tpu_enable_async_collective_fusion_fuse_all_reduce=true/false
Async
xla_tpu_enable_async_all_to_all
Ten flag włącza asynchroniczną komunikację typu „każdy z każdym”. xla_tpu_enable_async_all_to_all=false xla_tpu_enable_async_all_to_all=true xla_tpu_enable_async_all_to_all=true/false
Ograniczone przez czas oczekiwania
xla_all_gather_latency_bound_threshold_in_bytes
Ten flag jest przeznaczony dla operacji all-gather o małych rozmiarach, w przypadku których ważny jest czas oczekiwania. Włączenie tej opcji powoduje uruchomienie określonych optymalizacji, które mogą skrócić czas wykonywania operacji all-gather ograniczonych przez opóźnienie. Zwykle jest używana w przypadku zadań związanych z wnioskowaniem. xla_all_gather_latency_bound_threshold_in_bytes=-1
(które nie jest włączone)
4~16Mb(i.e. 4~16 * 1024 * 1024) [0, 9223372036854775807]
Ograniczone przez czas oczekiwania
xla_all_reduce_latency_bound_threshold_in_bytes
Ten flag jest przeznaczony dla operacji all-gather o małych rozmiarach, w przypadku których ważny jest czas oczekiwania. Włączenie tej opcji powoduje uruchomienie określonych optymalizacji, które mogą skrócić czas wykonywania operacji all-reduce ograniczonych przez opóźnienie. Zwykle jest używana w przypadku zadań związanych z wnioskowaniem. xla_all_reduce_latency_bound_threshold_in_bytes=-1
(które nie jest włączone)
4~16Mb(i.e. 4~16 * 1024 * 1024) [0, 9223372036854775807]
Ograniczone przez czas oczekiwania
xla_collective_permute_latency_bound_threshold_in_bytes
Ten flag jest przeznaczony dla operacji all-gather o małych rozmiarach, w przypadku których ważny jest czas oczekiwania. Włączenie tej opcji powoduje uruchomienie określonych optymalizacji, które mogą skrócić czas wykonywania operacji zbiorczych permutacji ograniczonych przez opóźnienie. Zwykle jest używana w przypadku zadań związanych z wnioskowaniem. xla_collective_permute_latency_bound_threshold_in_bytes=-1
(które nie jest włączone)
4~16Mb(i.e. 4~16 * 1024 * 1024) [0, 9223372036854775807]
Ograniczone przez czas oczekiwania
xla_all_to_all_latency_bound_threshold_in_bytes
Ten flag jest przeznaczony dla operacji all-gather o małych rozmiarach, w przypadku których ważny jest czas oczekiwania. Włączenie tej opcji powoduje uruchomienie konkretnych optymalizacji, które mogą skrócić czas wykonywania operacji typu all-to-all ograniczonych przez opóźnienie. Zwykle jest używana w przypadku zadań związanych z wnioskowaniem. xla_all_to_all_latency_bound_threshold_in_bytes=-1
(które nie jest włączone)
4~16Mb(i.e. 4~16 * 1024 * 1024) [0, 9223372036854775807]
xla_enable_async_collective_permute Przekształca wszystkie operacje collective-permute na ich asynchroniczne warianty. Gdy ta opcja jest ustawiona na auto, XLA może automatycznie włączać asynchroniczne operacje zbiorowe na podstawie innych konfiguracji lub warunków. xla_enable_async_collective_permute=kAuto xla_enable_async_collective_permute=kAuto xla_enable_async_collective_permute=kAuto/kEnabled/kDisabled

Flagi pamięci

Poniższe flagi służą do rozwiązywania problemów związanych z HBM. Te ustawienia należy zmieniać tylko wtedy, gdy podczas kompilacji modelu wystąpią błędy „out of memory” (brak pamięci) związane z pamięcią HBM. We wszystkich pozostałych przypadkach zalecamy wartości domyślne, ponieważ ich zmiana może negatywnie wpłynąć na skuteczność.

Flaga Opis Wartości domyślne Sugerowane wartości Wartości kandydatów
Harmonogram
xla_latency_hiding_scheduler_rerun
To ustawienie dostosowuje działanie harmonogramu ukrywającego opóźnienie. Działa on poprzez stopniowe zmniejszanie limitu pamięci przydzielonego do planowania przy każdym ponownym uruchomieniu procesu. xla_latency_hiding_scheduler_rerun=1 xla_latency_hiding_scheduler_rerun=5 0~10(it doesn’t make much sense beyond 10 reruns)
Fusion
xla_tpu_rwb_fusion
Ta flaga włącza fuzje typu reduce+broadcast i może zmniejszyć wykorzystanie pamięci. xla_tpu_rwb_fusion=true xla_tpu_rwb_fusion=false xla_tpu_rwb_fusion=true/false
Harmonogram
xla_memory_scheduler
Ta flaga określa algorytm, którego harmonogram pamięci będzie używać do minimalizowania zużycia pamięci. Użycie bardziej zaawansowanego algorytmu może dać mniej pamięciożerny harmonogram, ale kosztem dłuższego czasu kompilacji. xla_memory_scheduler=kDefault xla_memory_scheduler=kBrkga xla_memory_scheduler=kDefault/kList/kDfs/kPostOrder/kBrkga
Harmonogram
xla_tpu_enable_latency_hiding_scheduler
Ten flag włącza harmonogram ukrywający opóźnienia, który umożliwia wykonywanie asynchronicznych operacji zbiorczych zamiast synchronicznych. Wyłączenie tej funkcji zmniejsza zużycie pamięci, ale powoduje utratę wzrostu wydajności wynikającego z tych operacji asynchronicznych. xla_tpu_enable_latency_hiding_scheduler=true xla_tpu_enable_latency_hiding_scheduler=false xla_tpu_enable_latency_hiding_scheduler=true/false
SPMD
xla_jf_spmd_threshold_for_windowed_einsum_mib
Ten flag określa dolny próg minimalnego rozmiaru kropki, który wywołuje zbiorcze mnożenie macierzy. Ustawienie wyższej wartości pozwoli zaoszczędzić pamięć, ale kosztem utraty możliwości wykonywania zbiorczego mnożenia macierzy. xla_jf_spmd_threshold_for_windowed_einsum_mib=-1 10Mb~1Gb (i.e. 10*1024*1024 ~ 1024*1024*1024) [0, 9223372036854775807]

Inne często używane flagi

Flaga Typ Uwagi
xla_dump_to Ciąg znaków (ścieżka do pliku) Folder, w którym będą umieszczane pliki HLO przed optymalizacją i inne artefakty (patrz Narzędzia XLA).

Flagi XLA TPU

Flaga Typ Uwagi
xla_tpu_enable_data_parallel_all_reduce_opt Wartość logiczna (prawda/fałsz) Optymalizacja w celu zwiększenia możliwości nakładania się operacji all-reduce w sieciach centrów danych (DCN) używanych do równoległego dzielenia danych.
xla_tpu_data_parallel_opt_different_sized_ops Wartość logiczna (prawda/fałsz) Umożliwia potokowe przetwarzanie operacji równoległych na danych w wielu iteracjach, nawet jeśli rozmiary ich danych wyjściowych nie pasują do rozmiarów danych, które można zapisać w miejscu w zmiennych ułożonych w stos. Może zwiększyć obciążenie pamięci.
xla_tpu_spmd_rng_bit_generator_unsafe Wartość logiczna (prawda/fałsz) Określa, czy HLO RngBitGenerator ma być uruchamiany w sposób podzielony na partycje. Jest to niebezpieczne, jeśli oczekiwane są deterministyczne wyniki z różnymi podziałami na fragmenty w różnych częściach obliczeń.
xla_tpu_megacore_fusion_allow_ags Wartość logiczna (prawda/fałsz) Umożliwia łączenie operacji all-gather z operacjami splotu lub all-reduce.
xla_tpu_enable_ag_backward_pipelining Wartość logiczna (prawda/fałsz) Potoki zbierają wszystkie dane (obecnie megaskalowe zbieranie wszystkich danych) wstecz przez pętle skanowania.

Flagi XLA GPU

Flaga Typ Uwagi
xla_gpu_enable_latency_hiding_scheduler Wartość logiczna (prawda/fałsz) Ten flag umożliwia harmonogramom ukrywającym opóźnienia efektywne nakładanie się komunikacji asynchronicznej z obliczeniami. Wartość domyślna to False.
xla_gpu_enable_triton_gemm Wartość logiczna (prawda/fałsz) Używaj mnożenia macierzy opartego na Triton.
xla_gpu_graph_level Flaga (0–3) Starsza flaga do ustawiania poziomu wykresu GPU. Używaj xla_gpu_enable_command_buffer w nowych przypadkach użycia. 0 = wyłączone; 1 = przechwytywanie fuzji i funkcji memcpys; 2 = przechwytywanie funkcji gemm; 3 = przechwytywanie konwolucji.
xla_gpu_all_reduce_combine_threshold_bytes Liczba całkowita (bajty) Te flagi określają, kiedy połączyć wiele małych operacji AllGather, ReduceScatter lub AllReduce w jedną dużą operację AllGather, ReduceScatter lub AllReduce, aby skrócić czas poświęcony na komunikację między urządzeniami. Na przykład w przypadku progów AllGather / ReduceScatter w przypadku zbioru zadań opartego na architekturze Transformer rozważ dostrojenie ich tak, aby łączyły co najmniej wagę warstwy Transformer w operacjach AllGather / ReduceScatter. Domyślnie wartość combine_threshold_bytes wynosi 256.
xla_gpu_all_gather_combine_threshold_bytes Liczba całkowita (bajty) Patrz xla_gpu_all_reduce_combine_threshold_bytes powyżej.
xla_gpu_reduce_scatter_combine_threshold_bytes Liczba całkowita (bajty) Patrz xla_gpu_all_reduce_combine_threshold_bytes powyżej.
xla_gpu_enable_pipelined_all_gather Wartość logiczna (prawda/fałsz) Włącz potokowe przetwarzanie instrukcji all-gather.
xla_gpu_enable_pipelined_reduce_scatter Wartość logiczna (prawda/fałsz) Włącz potokowe przetwarzanie instrukcji reduce-scatter.
xla_gpu_enable_pipelined_all_reduce Wartość logiczna (prawda/fałsz) Włącz potokowe przetwarzanie instrukcji all-reduce.
xla_gpu_enable_while_loop_double_buffering Wartość logiczna (prawda/fałsz) Włącz podwójne buforowanie dla pętli while.
xla_gpu_enable_all_gather_combine_by_dim Wartość logiczna (prawda/fałsz) Łączenie operacji all-gather z tym samym wymiarem zbierania lub niezależnie od wymiaru.
xla_gpu_enable_reduce_scatter_combine_by_dim Wartość logiczna (prawda/fałsz) Łączenie operacji reduce-scatter z tym samym wymiarem lub niezależnie od wymiaru.