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

OpenXLA to kompilator do algebry liniowej, który może przyspieszyć działanie modeli TensorFlow bez konieczności wprowadzania zmian w kodzie źródłowym.

Wprowadzenie

Gdy program TensorFlow jest uruchamiany, wszystkie operacje są wykonywane indywidualnie przez wykonawcę TensorFlow. Każda operacja TensorFlow ma wstępnie skompilowane jądro GPU, które wykonawca przekazuje do wykonania.

XLA zapewnia alternatywny tryb uruchamiania modeli: kompiluje graf TensorFlow w sekwencję rdzeni obliczeniowych wygenerowanych specjalnie dla danego modelu. Kernele są unikalne dla danego modelu, więc mogą wykorzystywać informacje o nim do optymalizacji. Spójrzmy na przykład na optymalizację, którą XLA wykonuje w ramach prostego obliczenia w TensorFlow:

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

W przypadku uruchomienia bez XLA wykres uruchamia 3 jądra: jedno do mnożenia, jedno do dodawania i jedno do redukcji. XLA może jednak zoptymalizować graf, aby obliczenia były wykonywane w ramach jednego wywołania jądra. Dokonuje tego, „łącząc” dodawanie, mnożenie i redukcję w jednym rdzeniu GPU. Ponadto ta złączona operacja nie zapisuje do pamięci wartości pośrednich wygenerowanych przez funkcje y*z i x+y*z. Zamiast tego „przesyła” wyniki tych obliczeń pośrednich bezpośrednio do użytkowników, zachowując je całkowicie w rejestrach GPU. Fusion to najważniejsza optymalizacja XLA. Przepustowość pamięci jest zwykle najrzadszym zasobem w przyspieszaczach sprzętowych, dlatego usuwanie operacji na pamięci jest jednym z najlepszych sposobów na poprawę wydajności.

Włączanie XLA w przypadku modeli TensorFlow

Kompilacja jawna z tf.function(jit_compile=True)

Interfejs API do kompilacji jawnej zapewnia szczegółową kontrolę nad tym, które funkcje mają być kompilowane. Na przykład ta funkcja TensorFlow, która wykonuje trenowanie 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))

Interfejs API jit_compile ma semantykę musi być skompilowany: albo cała funkcja jest kompilowana za pomocą XLA, albo wyrzucana jest wyjątek errors.InvalidArgumentError. XLA nie może obecnie kompilować funkcji, w których wymiary nie są wywnioskowalne, czyli gdy nie można wywnioskować wymiarów wszystkich tensorów bez wykonania całego obliczenia. Na przykład ta funkcja nie zostanie skompilowana:

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

Kształty mogą się jednak różnić w trakcie poszczególnych przebiegów:

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

Aby zobaczyć bardziej szczegółowy przykład użycia i film instruktażowy dotyczący jit_compile=True, otwórz samouczek Colab.

Korzystanie z Keras

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

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

Użycie z strategią rozproszoną

XLA:GPU można używać w ramach strategii rozproszonej w TF (MirroredStrategy lub MultiWorkerMirroredStrategy), dodając adnotację do funkcji kroku za pomocą funkcji 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 używania XLA w modelach TensorFlow bez wprowadzania żadnych zmian jest włączenie auto-clusteringu, który automatycznie znajduje grupy (połączone podgrafy) w funkcjach TensorFlow, które można skompilować i wykonywać za pomocą XLA. Automatyczne grupowanie na 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

Autoklasterowanie jest obecnie zoptymalizowane pod kątem zbiorów zadań korzystających z GPU, ale można je włączyć również w przypadku procesora, używając dodatkowo 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 samouczkielu w colab na temat automatycznego grupowania.

Kompilacja AOT (z wyprzedzeniem) dla procesora z tfcompile

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

Sprawdzanie skompilowanych programów

XLA udostępnia funkcje inspekcji, które umożliwiają sprawdzanie wygenerowanych programów. Aby wygenerować 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 zakończeniu generowania kopii zapasowej w folderze /tmp/generated znajdziesz te pliki:

  • module_XXXX.*_optimizations.txt Wygenerowane programy XLA, po jednym dla każdego skompilowanego klastra. Dołączenie ich podczas przesyłania raportów o błędach XLA jest bardzo pomocne.

  • module_XXXX.ir-*.ll Wygenerowane pliki w reprezentacji pośredniej LLVM z elementami charakterystycznymi dla NVPTX.

  • module_XXXX.ptx Wygenerowane pliki PTX.

Możesz też wygenerować zapis wykresu z wizualizacją zagłębiania klastrów XLA w grafie TensorFlow za pomocą:

$ TF_DUMP_GRAPH_PREFIX=/tmp/generated TF_XLA_FLAGS="--tf_xla_clustering_debug"

Raporty o błędach, które można odtworzyć

Raport o błędzie jest znacznie łatwiej odtworzyć, jeśli zawiera zrzuty wygenerowanych programów XLA i wykorzystywanego wbudowania autoklasteryzacji. Aby wygenerować je dla programu TensorFlow działającego z automatycznym grupowaniem, 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"

Podczas zgłaszania błędów załącz zawartość katalogu /tmp/generated (o którym mowa powyżej).

Jeśli to możliwe, spróbuj wyizolować błąd w pojedynczym programie XLA, używając narzędzia run_hlo_module i cyklicznie uruchamiając je na wygenerowanych programach.

Więcej informacji

Frontendy XLA

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

  • JAX: składane przekształcenia programów Python+NumPy
  • Julia: język Julia do obliczeń naukowych.
  • PyTorch platforma PyTorch,
  • Nx: biblioteka do obliczeń numerycznych dla języka programowania Elixir.

Dyskusje

Korzystanie z XLA z TF przy użyciu jit_compile=True

Omówienie XLA