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

Zadbaj o dobrą organizację dzięki kolekcji Zapisuj i kategoryzuj treści zgodnie ze swoimi preferencjami.

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

Rezultatem jest poprawa szybkości i wykorzystania pamięci: np. w przesyłaniu BERT MLPerf przy użyciu procesorów graficznych 8 Volta V100 przy użyciu XLA osiągnięto ~7x poprawę wydajności i ~5x poprawę wielkości 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, ta połączona operacja nie zapisuje wartości pośrednich wytworzonych przez y*z i x+y*z do 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 jest zwykle najrzadszym zasobem w akceleratorach sprzętowych, więc usunięcie 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 powinny 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))

Interfejs API jit_compile ma semantykę konieczności kompilacji : albo cała funkcja jest kompilowana przy użyciu XLA, albo zgłaszany jest wyjątek errors.InvalidArgumentError . XLA nie może obecnie kompilować funkcji, w których wymiary nie są możliwe do wywnioskowania : to znaczy, jeśli nie można wywnioskować wymiarów wszystkich tensorów bez uruchamiania całego obliczenia. 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 , aby uzyskać bardziej szczegółowy przykład użycia, oraz samouczek wideo na temat jit_compile=True .

Korzystanie z Keras

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

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

Wykorzystanie ze strategią rozproszoną

XLA:GPU może być używany z rozproszoną strategią TF ( MirroredStrategy lub MultiWorkerMirroredStrategy ) poprzez adnotację funkcji step z 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 klastrowanie

Prostym sposobem na rozpoczęcie korzystania z XLA w modelach TensorFlow bez żadnych zmian jest włączenie autoklastrowania , które automatycznie wyszukuje klastry (połączone podgrafy) w ramach funkcji TensorFlow, które można skompilować i wykonać 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, 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

Aby zapoznać się ze szczegółowym przykładem użycia, zobacz colab samouczek dotyczący automatycznego klastrowania .

Kompilacja AOT (z wyprzedzeniem) dla procesora z tfcompile

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

Sprawdź skompilowane programy

XLA zapewnia narzędzia do introspekcji, które pozwalają sprawdzić wygenerowane programy. 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 możesz znaleźć następujące pliki w /tmp/generated :

  • 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 Generowane pliki w pośredniej reprezentacji LLVM , z wewnętrznymi funkcjami NVPTX .

  • module_XXXX.ptx Wygenerowane pliki PTX .

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"

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

Jeśli to możliwe, spróbuj wyizolować błąd do pojedynczego programu XLA, używając replay_computation iteracyjne uruchamianie go na wygenerowanych programach.

Dalsza lektura

Interfejsy XLA

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

  • JAX : Komponowalne przekształcenia programów Python+NumPy
  • Julia : Język Julii w obliczeniach naukowych
  • PyTorch : framework PyTorch
  • Nx : Numeryczna biblioteka obliczeniowa dla języka programowania Elixir

Rozmowy

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

Przegląd XLA