XLA: optymalizacja kompilatora pod kątem systemów uczących się

OpenXLA to specyficzny dla domeny kompilator do algebry liniowej, który może przyspieszać modele TensorFlow bez potencjalnych zmian w kodzie źródłowym.

Wstęp

Gdy program TensorFlow jest uruchamiany, wszystkie operacje są wykonywane indywidualnie przez wykonawcy TensorFlow. Każda operacja TensorFlow ma wstępnie skompilowaną implementację jądra GPU, do której odsyła wykonawca.

XLA zapewnia alternatywny tryb uruchamiania modeli: kompiluje wykres TensorFlow w sekwencję jąder obliczeniowych wygenerowanych specjalnie dla danego modelu. Ze względu na to, że jądra są unikalne dla modelu, mogą wykorzystywać informacje dotyczące modelu do optymalizacji. Przyjrzyjmy się optymalizacji XLA w kontekście prostych obliczeń w TensorFlow:

def model_fn(x, y, z):
  return tf.reduce_sum(x + y * z)

Uruchomienie go bez XLA spowoduje uruchomienie 3 jąder: 1 do mnożenia, 1 do dodawania, a drugiego do redukcji. XLA może jednak zoptymalizować wykres tak, aby obliczał wynik przy pojedynczym uruchomieniu jądra. Odbywa się to przez „łączenie” dodawania, mnożenia i redukcji w jedno jądro GPU. Co więcej, ta operacja scalona nie zapisuje w pamięci wartości pośrednich wygenerowanych przez y*z i x+y*z. Zamiast tego przesyła wyniki obliczeń pośrednich bezpośrednio do użytkowników, zachowując je w całości w rejestrach GPU. Fusion to najważniejsza optymalizacja XLA. Przepustowość pamięci jest zwykle ograniczonym zasobem w akceleratorach sprzętowych, więc usunięcie operacji pamięci to jeden z najlepszych sposobów na poprawę wydajności.

Włącz XLA dla modeli TensorFlow

Kompilacja dla pełnoletnich z użyciem narzędzia tf.function(jit_compile=True)

Interfejs API jawnej kompilacji zapewnia szczegółową kontrolę nad wyborem funkcji do skompilowania. Na przykład funkcja TensorFlow, która wykonuje trenowanie MNIST, została skompilowana do 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))

Interfejs API jit_compile wymaga semantyki musi kompilować: cała funkcja jest skompilowana za pomocą języka XLA lub przesyłany jest wyjątek errors.InvalidArgumentError. XLA nie może obecnie kompilować funkcji, w których nie można wywnioskować wymiarów, czyli gdy nie da się wywnioskować wymiarów wszystkich teniserów bez wykonania wszystkich obliczeń. Na przykład ta funkcja się nie skompiluje:

@tf.function
def not_compilable(x):
  return tf.unique(x)

Kształty mogą się jednak zmieniać w poszczególnych wersjach:

@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]))

Bardziej szczegółowe informacje na temat wykorzystania znajdziesz w samouczku Colab, a także w samouczku poświęconym korzystaniu z jit_compile=True.

Wykorzystanie z Keras

W przypadku modeli Keras funkcję jit_compile=True można ustawić jako argument funkcji model.compile:

model.compile(optimizer="adam", jit_compile=True)

Użycie ze strategią rozproszoną

XLA:GPU można używać ze strategią rozproszoną TF (MirroredStrategy lub MultiWorkerMirroredStrategy), dodając do funkcji kroku adnotację jit_compile=True:

@tf.function(jit_compile=True)
def step_fn():
  t = tf.ones(shape=[100], dtype=tf.float32)
  ctx = tf.distribute.get_replica_context()
  return ctx.all_reduce(tf.distribute.ReduceOp.SUM, t)

@tf.function
def run_fn():
  return strategy.run(step_fn)

Automatyczne grupowanie

Prostym sposobem na rozpoczęcie korzystania z XLA w modelach TensorFlow bez żadnych zmian jest włączenie automatycznego klastrowania, co automatycznie wyszukuje klastry (połączone podgrafy) w funkcjach TensorFlow, które można kompilować i wykonywać za pomocą XLA. Automatyczne klastrowanie w GPU można włączyć, ustawiając zmienną środowiskową TF_XLA_FLAGS:

$ TF_XLA_FLAGS=--tf_xla_auto_jit=2 path/to/your/tf/program

Automatyczne klastrowanie jest obecnie zoptymalizowane pod kątem zbiorów zadań GPU, ale można je też włączyć na procesorze, dodatkowo używając flagi --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życia znajdziesz w samouczku automatycznego grupowania w Colab.

Kompilacja AOT (z wyprzedzeniem) dla procesora tfcompile

Możesz też użyć osobnego narzędzia tfcompile, które konwertuje wykres TensorFlow na kod wykonywalny (tylko w przypadku procesorów x86–64).

Sprawdzanie skompilowanych programów

XLA zapewnia dostęp do narzędzi do refleksji, w których możesz badać wygenerowane programy. Aby zrzutować wygenerowane programy, użyj zmiennej środowiskowej XLA_FLAGS:

$ XLA_FLAGS="--xla_dump_to=/tmp/generated" TF_XLA_FLAGS="--tf_xla_auto_jit=2" my/tensorflow/program

Po wykonaniu zrzutu znajdziesz te pliki w /tmp/generated:

  • module_XXXX.*_optimizations.txt Wygenerowane programy XLA, po jednym na każdy skompilowany klaster. Załączenie ich podczas przesyłania raportów o błędach XLA będzie bardzo pomocne.

  • module_XXXX.ir-*.ll Wygenerowane pliki w średniej reprezentacji LLVM z elementami wewnętrznymi NVPTX.

  • module_XXXX.ptx Wygenerowane pliki PTX.

Możesz też wykonać zrzut wykresu wizualizującego osadzanie klastrów XLA w wykresie TensorFlow:

$ 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 on zrzuty wygenerowanych programów XLA i używanego umieszczania automatycznego klastrowania. Aby je wygenerować dla programu TensorFlow z automatycznym klastrem, 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"

Do zgłaszania błędów dołącz zawartość katalogu /tmp/generated (wspomnianego powyżej).

Jeśli to możliwe, spróbuj wyodrębnić błąd tylko do jednego programu XLA. W tym celu użyj narzędzia run_hlo_module i iteracyjnie uruchom go w wygenerowanych programach.

Więcej informacji

Frontendy XLA

Oprócz TensorFlow programy XLA mogą być generowane przez:

  • JAX: kompozycyjne przekształcenia programów Python+NumPy
  • Julia: język Julii do obliczeń naukowych
  • PyTorch: platforma PyTorch
  • Nx: numeryczna biblioteka obliczeniowa dla języka programowania Elixir

Dyskusje

Używanie XLA z TF przy użyciu jit_compile=True

Przegląd XLA