Ta strona została przetłumaczona przez Cloud Translation API.
Switch to English

XLA: optymalizacja kompilatora 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 zmian w kodzie źródłowym.

Rezultatem jest poprawa szybkości i wykorzystania pamięci: większość wewnętrznych testów porównawczych działa ~ 1,15x szybciej po włączeniu XLA. Poniższy zestaw danych jest oceniany na pojedynczym GPU NVidia V100:

Wprowadzenie

Po uruchomieniu programu TensorFlow wszystkie operacje są wykonywane indywidualnie przez moduł wykonawczy TensorFlow. Każda operacja TensorFlow ma prekompilowaną implementację jądra GPU, do której jest wysyłana przez moduł wykonawczy.

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ę wykonywaną przez XLA w kontekście prostego obliczenia TensorFlow: 

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

Uruchomiony 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 w pojedynczym uruchomieniu jądra. Czyni to poprzez „łączenie” dodawania, mnożenia i redukcji w jednym jądrze GPU. Co więcej, ta operacja połączona nie zapisuje do pamięci wartości pośrednich wytworzonych przez y*z i x+y*z ; zamiast tego „przesyła” wyniki tych pośrednich obliczeń bezpośrednio do ich użytkowników, zachowując je całkowicie w rejestrach GPU. Fusion to najważniejsza optymalizacja XLA. Przepustowość pamięci jest zwykle najbardziej deficytowym zasobem w akceleratorach sprzętowych, więc usunięcie operacji pamięci jest jednym z najlepszych sposobów na poprawę wydajności.

Włącz XLA dla modeli TensorFlow

Automatyczne grupowanie

Najprostszym sposobem rozpoczęcia korzystania z XLA w modelach TensorFlow jest włączenie automatycznego grupowania , które automatycznie wyszukuje klastry (połączone podgrafy) na wykresie TensorFlow, który można skompilować i wykonać za pomocą XLA. Automatyczne klastrowanie na GPU można włączyć, ustawiając TF_XLA_FLAGS ś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

Szczegółowy przykład użycia można znaleźć w samouczku dotyczącym automatycznego klastrowania - colab .

Jawna kompilacja z tf.function

Automatyczne grupowanie to świetne narzędzie do przyspieszania modelu bez wprowadzania zmian w kodzie, ale zrozumienie, jakie zmiany zostały wprowadzone, może być trudne.

Interfejs API do jawnej kompilacji oferuje bardziej szczegółową kontrolę wyboru funkcji, które mają zostać skompilowane. Na przykład następująca funkcja TensorFlow, która wykonuje szkolenie MNIST, jest kompilowana z XLA: 

 @tf.function(experimental_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))

experimental_compile API ma semantykę obowiązkowej kompilacji : albo cała funkcja jest skompilowana z XLA, albo errors.InvalidArgumentError się errors.InvalidArgumentError jest wyjątek errors.InvalidArgumentError . XLA nie może obecnie kompilować funkcji, w których wymiarów nie można wywnioskować : to znaczy, jeśli nie jest możliwe wywnioskowanie wymiarów wszystkich tensorów bez wykonania całego obliczenia. Na przykład następująca funkcja nie zostanie skompilowana: 

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

Jednak kształty mogą się różnić w zależności od biegów: 

 @tf.function(experimental_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ółowy przykład użycia znajduje się w tutorialu colab .

Kompilacja AOT (Ahead-of-time) dla procesora z 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 funkcje introspekcji, które pozwalają na inspekcję 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 można znaleźć następujące pliki w katalogu /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 Wygenerowane pliki w reprezentacji pośredniej LLVM , z elementami wewnętrznymi 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żywane osadzanie 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 wyodrębnić błąd w pojedynczym programie XLA, używając replay_computation i iteracyjnie uruchamiając go na wygenerowanych programach.

Interfejsy XLA

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

  • JAX : Komponowalne transformacje programów Python + NumPy
  • Julia : Język Julia do obliczeń naukowych
  • PyTorch : framework PyTorch

Dalsza lektura