Эта страница была переведа с помощью Cloud Translation API.
Switch to English

XLA: Оптимизация компилятора для машинного обучения

XLA (Ускоренная линейная алгебра) является предметно-ориентированным компилятором для линейной алгебры, который может ускорять модели TensorFlow без потенциального изменения исходного кода.

Результатом являются улучшения в скорости и использовании памяти: большинство внутренних тестов работают примерно в 1,15 раза быстрее после включения XLA. Приведенный ниже набор данных оценивается на одном графическом процессоре NVidia V100:

Введение

При запуске программы TensorFlow все операции выполняются индивидуально исполнителем TensorFlow. Каждая операция TensorFlow имеет предварительно скомпилированную реализацию ядра GPU, которую отправляет исполнитель.

XLA предоставляет альтернативный режим запуска моделей: он компилирует граф TensorFlow в последовательность вычислительных ядер, созданных специально для данной модели. Поскольку эти ядра уникальны для модели, они могут использовать информацию, специфичную для модели, для оптимизации. Например, давайте посмотрим на оптимизацию, которую XLA выполняет в контексте простого вычисления TensorFlow: 

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

Запустив без XLA, график запускает три ядра: одно для умножения, одно для сложения и одно для сокращения. Однако XLA может оптимизировать график, чтобы он вычислял результат за один запуск ядра. Это делается путем «объединения» сложения, умножения и сокращения в одно ядро ​​графического процессора. Более того, эта операция слияния не записывает промежуточные значения, полученные от y*z и x+y*z в память; вместо этого он «передает» результаты этих промежуточных вычислений непосредственно своим пользователям, сохраняя их полностью в регистрах GPU. Fusion - самая важная оптимизация XLA. Пропускная способность памяти обычно является самым дефицитным ресурсом на аппаратных ускорителях, поэтому удаление операций с памятью является одним из лучших способов повышения производительности.

Включить XLA для моделей TensorFlow

Авто-кластеризация

Простейший способ начать использовать XLA в моделях TensorFlow является обеспечением автоматической кластеризации, который автоматически находит кластеры (связанные подграфы) в пределах TensorFlow графа , который может быть скомпилирован и выполнен с использованием XLA. Автоматическую кластеризацию на графическом процессоре можно включить, установив переменную среды TF_XLA_FLAGS : 

 $ TF_XLA_FLAGS=--tf_xla_auto_jit=2 path/to/your/tf/program

Автоматическая кластеризация в настоящее время оптимизирована для рабочих нагрузок графического процессора, но ее также можно включить на ЦП, дополнительно используя флаг --tf_xla_cpu_global_jit : 

 $ TF_XLA_FLAGS="--tf_xla_auto_jit=2 --tf_xla_cpu_global_jit" path/to/your/program

Подробный пример использования см. В учебнике по авто-кластеризации .

Явная компиляция с tf.function

Авто-кластеризация - отличный инструмент для ускорения модели без каких-либо изменений в коде, но может быть трудно понять, какие изменения были выполнены.

API явной компиляции предлагает более детальный контроль для выбора того, какие функции должны быть скомпилированы. Например, следующая функция TensorFlow, которая выполняет обучение MNIST, скомпилирована с 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 имеет семантику обязательной компиляции : либо вся функция компилируется с помощью XLA, либо errors.InvalidArgumentError Исключение errors.InvalidArgumentError . В настоящее время XLA не может компилировать функции, где измерения не являются выводными : то есть, если невозможно вывести размеры всех тензоров без выполнения всего вычисления. Например, следующая функция не будет компилироваться: 

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

Формы могут варьироваться в зависимости от трассы: 

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

См. Учебник Colab для более подробного примера использования.

Компиляция AOT (Ahead-of-time) для CPU с помощью tfcompile

Вы также можете использовать автономный инструмент tfcompile , который преобразует граф tfcompile в исполняемый код (только для процессора x86-64).

Проверять скомпилированные программы

XLA предоставляет средства для самоанализа, которые позволяют вам проверять сгенерированные программы. Чтобы XLA_FLAGS сгенерированные программы, используйте переменную окружения XLA_FLAGS : 

 $ XLA_FLAGS="--xla_dump_to=/tmp/generated" TF_XLA_FLAGS="--tf_xla_auto_jit=2" my/tensorflow/program

После выполнения дампа вы можете найти следующие файлы в /tmp/generated :

  • module_XXXX.*_optimizations.txt Созданные программы XLA , по одной на каждый скомпилированный кластер. Присоединение к ним при отправке отчетов об ошибках XLA чрезвычайно полезно!

  • module_XXXX.ir-*.ll Сгенерированные файлы в промежуточном представлении LLVM с внутренними свойствами NVPTX .

  • module_XXXX.ptx Сгенерированные файлы PTX .

Вы также можете вывести график, визуализирующий встраивание кластеров XLA в граф TensorFlow с помощью: 

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

Воспроизводимые сообщения об ошибках

Отчет об ошибке гораздо проще воспроизвести, если он содержит дампы для сгенерированных программ XLA и использованное встраивание авто-кластеризации. Чтобы сгенерировать их для программы TensorFlow, работающей с авто-кластеризацией, запустите: 

 $ 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"

При отправке сообщений об ошибках прикрепляйте содержимое каталога /tmp/generated (ссылка выше).

Если возможно, попытайтесь изолировать ошибку в одной XLA-программе, используя replay_computation и итеративно выполняя ее на сгенерированных программах.

XLA Frontends

Помимо TensorFlow, программы XLA могут быть сгенерированы:

  • JAX : Композиционные преобразования программ Python + NumPy
  • Julia : язык Julia для научных вычислений
  • PyTorch : фреймворк PyTorch

дальнейшее чтение