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

XLA (Accelerated Linear Algebra) to specyficzny dla domeny kompilator algebry liniowej, który może przyspieszać modele TensorFlow bez potencjalnie zmian w kodzie źródłowym.

Rezultatem jest poprawa szybkości i wykorzystania pamięci: np. w przypadku BERT MLPerf przy użyciu procesorów graficznych 8 Volta V100 przy użyciu XLA osiągnięto ok. 7-krotną poprawę wydajności i ok. 5-krotną poprawę rozmiaru 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łany jest wykonawca.

XLA zapewnia alternatywny tryb uruchamiania modeli: kompiluje wykres TensorFlow w sekwencję jąder obliczeniowych generowanych 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ę wykonywaną przez 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 przy pojedynczym uruchomieniu jądra. Czyni to poprzez „łączenie” dodawania, mnożenia i redukcji w jednym jądrze GPU. Co więcej, ta połączona operacja nie zapisuje do pamięci wartości pośrednich utworzonych przez y*z i x+y*z ; zamiast tego „przesyła strumieniowo” wyniki tych pośrednich obliczeń bezpośrednio do ich użytkowników, zachowując je w całości w rejestrach GPU. Fusion to najważniejsza optymalizacja XLA. Przepustowość pamięci jest zwykle najrzadszym zasobem w akceleratorach sprzętowych, więc usuwanie operacji pamięciowych 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 precyzyjną kontrolę wyboru funkcji, które powinny zostać skompilowane. Na przykład następująca funkcja TensorFlow, która wykonuje szkolenie MNIST, jest kompilowana z 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ę „must-compile” : albo cała funkcja jest kompilowana z XLA, albo generowany jest wyjątek errors.InvalidArgumentError . XLA nie może obecnie kompilować funkcji, których wymiarów nie da się wywnioskować : to znaczy, jeśli nie jest możliwe wywnioskowanie wymiarów wszystkich tensorów bez wykonywania całego obliczenia. Na przykład następująca funkcja nie skompiluje się:

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

Kształty mogą się jednak różnić w zależności od serii:

@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 , aby uzyskać bardziej szczegółowy przykład użycia, oraz samouczek wideo na temat jit_compile=True use.

Użycie z Kerasem

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

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

Użycie ze strategią rozproszoną

XLA:GPU może być używany ze strategią rozproszoną TF ( MirroredStrategy lub MultiWorkerMirroredStrategy ) poprzez adnotację funkcji kroku za pomocą 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 , które automatycznie wyszukuje klastry (połączone podgrafy) w funkcjach TensorFlow, które można kompilować i wykonywać za pomocą XLA. Automatyczne klastrowanie 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

Automatyczne klastrowanie jest obecnie zoptymalizowane pod kątem obciążeń GPU, ale można je również 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

Aby zapoznać się ze szczegółowym przykładem użycia, zapoznaj się z samouczkiem dotyczącym automatycznego klastrowania .

Kompilacja AOT (z wyprzedzeniem) dla procesora za pomocą tfcompile

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

Sprawdź skompilowane programy

XLA zapewnia narzędzia do introspekcji, które umożliwiają kontrolę wygenerowanych programów. Aby zrzucić 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 w /tmp/generated można znaleźć następujące pliki:

  • module_XXXX.*_optimizations.txt Wygenerowane programy XLA , po jednym na każdy skompilowany klaster. Dołączanie ich podczas przesyłania raportów o błędach XLA jest niezwykle pomocne!

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

  • module_XXXX.ptx Wygenerowane pliki PTX .

Możesz również zrzucić wykres wizualizujący osadzenie 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żywanego 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"

Zgłaszając błędy, dołącz zawartość katalogu /tmp/generated (o którym mowa powyżej).

Jeśli to możliwe, spróbuj wyizolować błąd pojedynczego programu XLA, używając run_hlo_module i iteracyjnie uruchamiając go w wygenerowanych programach.

Dalsza lektura

Interfejsy XLA

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

  • JAX : Komponowalne transformacje programów Python+NumPy
  • Julia : język Julii do obliczeń naukowych
  • PyTorch : framework PyTorch
  • Nx : Biblioteka obliczeń numerycznych dla języka programowania Elixir

Rozmowy

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

Przegląd XLA