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: np. W przesyłaniu BERT MLPerf przy użyciu 8 procesorów graficznych Volta V100 używających XLA osiągnięto ~ 7-krotną poprawę wydajności i ~ 5-krotną poprawę wielkości partii:
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
Jawna kompilacja z tf.function(jit_compile=True)
Interfejs API do jawnej kompilacji oferuje precyzyjną 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(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))
jit_compile API jit_compile ma semantykę obowiązkowej kompilacji : albo cała funkcja jest skompilowana z XLA, albo wystąpił errors.InvalidArgumentError . XLA nie może obecnie kompilować funkcji, w których wymiarów nie można wywnioskować : to znaczy, jeśli nie można wywnioskować 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)
Kształty mogą się jednak różnić w zależności od biegó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]))
Bardziej szczegółowy przykład użycia znajduje się w tutorialu colab .
Automatyczne grupowanie
Prostym sposobem na rozpoczęcie korzystania z XLA w modelach TensorFlow bez żadnych zmian jest włączenie automatycznego grupowania , które automatycznie wyszukuje klastry (połączone podgrafy) w funkcjach TensorFlow, które mogą być kompilowane i wykonywane przy użyciu 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, 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 można znaleźć w samouczku dotyczącym automatycznego klastrowania - 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.txtWygenerowane 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-*.llWygenerowane pliki w reprezentacji pośredniej LLVM , z elementami wewnętrznymi NVPTX .module_XXXX.ptxWygenerowane 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.
Dalsza lektura
- Znane problemy Lista znanych problemów z XLA
- Architektura XLA : przegląd architektury XLA
- XLA - TensorFlow, skompilowany : przeczytaj na blogu Google Developers
- Sprawdź źródło XLA na Github!
Interfejsy XLA
Oprócz TensorFlow programy XLA mogą być generowane przez: