Dołącz do TensorFlow na Google I/O, 11-12 maja Zarejestruj się teraz

XLA: kompilator optymalizujący pod kątem uczenia maszynowego

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

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

Rozmowy

Korzystanie XLA z TF za pomocą jit_compile=True

Przegląd XLA