XLA (Accelerated Linear Algebra) to specyficzny dla domeny kompilator algebry liniowej, który może przyspieszyć modele TensorFlow bez potencjalnie żadnych zmian w kodzie źródłowym.
Wyniki są ulepszenia w szybkości i wykorzystanie pamięci: na przykład w BERT MLPerf składania wykorzystaniem 8 Volta V100 GPU wykorzystujących XLA osiągnęła ~ 7x wzrost wydajności i poprawę wielkości ~ 5x partii:

Wstęp
Po uruchomieniu programu TensorFlow wszystkie operacje są wykonywane indywidualnie przez executor TensorFlow. Każda operacja TensorFlow ma wstępnie skompilowaną implementację jądra GPU, do której wysyła executor.
XLA zapewnia alternatywny tryb uruchamiania modeli: kompiluje wykres TensorFlow w sekwencję jąder obliczeniowych wygenerowanych specjalnie dla danego modelu. Ponieważ te jądra są unikalne dla modelu, mogą wykorzystywać informacje specyficzne dla modelu do optymalizacji. Na przykład spójrzmy na optymalizację XLA w kontekście prostego obliczenia TensorFlow:
def model_fn(x, y, z):
return tf.reduce_sum(x + y * z)
Uruchom bez XLA, wykres uruchamia trzy jądra: jedno do mnożenia, jedno do dodawania i jedno do redukcji. Jednak XLA może zoptymalizować wykres tak, aby obliczał wynik podczas jednego uruchomienia jądra. Odbywa się to poprzez „łączenie” dodawania, mnożenia i redukcji w jedno jądro GPU. Co więcej, to połączone działanie nie pisze się wartości pośrednie wytworzone y*z
, a x+y*z
pamięci; zamiast tego „przesyła” wyniki tych pośrednich obliczeń bezpośrednio do użytkowników, jednocześnie przechowując je w całości w rejestrach GPU. Fusion to najważniejsza optymalizacja XLA. Przepustowość pamięci to zazwyczaj najrzadszy zasób w akceleratorach sprzętowych, więc usuwanie operacji na pamięci jest jednym z najlepszych sposobów na poprawę wydajności.
Włącz XLA dla modeli TensorFlow
Jawna kompilacja z tf.function(jit_compile=True)
Jawny interfejs API kompilacji oferuje szczegółową kontrolę wyboru funkcji, które mają być kompilowane. Na przykład poniższa funkcja TensorFlow realizująca szkolenie MNIST jest kompilowana za pomocą XLA:
@tf.function(jit_compile=True)
def train_mnist(images, labels):
images, labels = cast(images, labels)
with tf.GradientTape() as tape:
predicted_labels = layer(images)
loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(
logits=predicted_labels, labels=labels
))
layer_variables = layer.trainable_variables
grads = tape.gradient(loss, layer_variables)
optimizer.apply_gradients(zip(grads, layer_variables))
jit_compile
API ma semantykę trzeba kompilacji: albo cała funkcja jest zestawiane z XLA, albo errors.InvalidArgumentError
wyjątku. XLA nie może obecnie skompilować funkcje gdzie wymiary nie są inferrable: to znaczy, jeśli to nie jest możliwe, aby wywnioskować wymiary wszystkich tensorów bez uruchamiania całego obliczeń. Na przykład poniższa funkcja nie skompiluje się:
@tf.function
def not_compilable(x):
return tf.unique(x)
Kształty mogą się jednak różnić w poszczególnych biegach:
@tf.function(jit_compile=True)
def recompiled_on_launch(a, b):
return a + b
recompiled_on_launch(tf.ones([1, 10]), tf.ones([1, 10]))
recompiled_on_launch(tf.ones([1, 100]), tf.ones([1, 100]))
Zobacz samouczek colab bardziej szczegółowym przykładzie użytkowania i poradnik wideo na jit_compile=True
użytkowania.
Automatyczne klastrowanie
Prostym sposobem, aby rozpocząć korzystanie z XLA w modelach TensorFlow bez zmian jest umożliwienie automatycznego klastry, które automatycznie wyszukuje klastrów (podłączony subgraphs) w obrębie funkcji TensorFlow które mogą być skompilowane i wykonanych przy użyciu XLA. Automatyczne grupowanie na GPU można włączyć poprzez ustawienie TF_XLA_FLAGS
zmienną środowiskową:
$ TF_XLA_FLAGS=--tf_xla_auto_jit=2 path/to/your/tf/program
Automatyczne grupowanie jest obecnie zoptymalizowany pod kątem obciążenia GPU, ale może być również włączona przez CPU dodatkowo z flagą --tf_xla_cpu_global_jit
:
$ TF_XLA_FLAGS="--tf_xla_auto_jit=2 --tf_xla_cpu_global_jit" path/to/your/program
Szczegółowy przykład użytkowania zobaczyć auto-klastrów poradnik colab .
AOT (Ahead-of-time) kompilacja dla CPU z tfcompile
Można również skorzystać z oddzielnego tfcompile
narzędzie, które konwertuje TensorFlow wykres do kodu wykonywalnego (tylko dla x86-64 CPU).
Sprawdź skompilowane programy
XLA zapewnia narzędzia do introspekcji, które pozwalają sprawdzić wygenerowane programy. Zrzucić wygenerowanych programów, należy użyć zmiennej środowiskowej XLA_FLAGS
:
$ XLA_FLAGS="--xla_dump_to=/tmp/generated" TF_XLA_FLAGS="--tf_xla_auto_jit=2" my/tensorflow/program
Po dumpingu jest wykonywana, można znaleźć następujące pliki w /tmp/generated
:
module_XXXX.*_optimizations.txt
Wygenerowane programy XLA , po jednej w każdej skompilowany klastra. Dołączanie ich podczas przesyłania raportów o błędach XLA jest niezwykle pomocne!module_XXXX.ir-*.ll
generowanych plików w LLVM pośredniej reprezentacji, z NVPTX intrinsics.module_XXXX.ptx
Generated PTX plików.
Możesz również zrzucić wykres wizualizujący osadzanie klastrów XLA wewnątrz wykresu TensorFlow za pomocą:
$ TF_DUMP_GRAPH_PREFIX=/tmp/generated TF_XLA_FLAGS="--tf_xla_clustering_debug"
Powtarzalne raporty o błędach
Raport o błędzie jest znacznie łatwiejszy do odtworzenia, jeśli zawiera zrzuty dla wygenerowanych programów XLA i użytego osadzania automatycznego klastrowania. Aby wygenerować je dla programu TensorFlow działającego z automatycznym klastrowaniem, uruchom:
$ TF_DUMP_GRAPH_PREFIX=/tmp/generated \
TF_XLA_FLAGS="--tf_xla_clustering_debug --tf_xla_auto_jit=2" \
XLA_FLAGS="--xla_dump_hlo_as_text --xla_dump_to=/tmp/generated" \
my/tensorflow/program"
Kiedy zgłoszenia błędów, dołącz zawartość /tmp/generated
katalogu (mowa powyżej).
Jeśli to możliwe, staraj się wyizolować błąd na jednym programie XLA za pomocą replay_computation
i iteracyjnie uruchomienie go na generowanych programów.
Dalsza lektura
- Znanych problemów listę znanych problemów z XLA
- XLA Architektura : Przegląd architektury XLA
- XLA - TensorFlow, Opracowanie : Przeczytaj na blogu Google Developers
- Sprawdź źródło XLA na Github!
Interfejsy XLA
Oprócz TensorFlow, programy XLA mogą być generowane przez:
- JAX : composable przekształcenia programów Python + NumPy
- Julia : Język Julia do obliczeń naukowych
- PyTorch : PyTorch ramy
- Nx : Numeryczne obliczanie biblioteka dla języka programowania Elixir