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

OpenXLA — это предметно-ориентированный компилятор линейной алгебры, который может ускорять модели TensorFlow без внесения изменений в исходный код.

Введение

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

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

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

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

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

Явная компиляция с помощью tf.function(jit_compile=True)

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

API jit_compile имеет семантику обязательной компиляции : либо вся функция компилируется с помощью XLA, либо выдается исключение errors.InvalidArgumentError . XLA в настоящее время не может компилировать функции, размеры которых не являются выводимыми : то есть, если невозможно вывести размеры всех тензоров без выполнения всех вычислений. Например, следующая функция не скомпилируется:

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

Однако формы могут различаться в зависимости от пробега:

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

Более подробный пример использования см. в учебном соавторстве , а также обучающее видео по использованию jit_compile=True .

Использование с Керасом

Для моделей Keras jit_compile=True можно установить в качестве аргумента model.compile : 

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

Использование с распределенной стратегией

XLA:GPU можно использовать с распределенной стратегией TF ( MirroredStrategy или MultiWorkerMirroredStrategy ), аннотируя функцию шага с помощью 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)

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

Простой способ начать использовать 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

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

AOT (заранее время) компиляция для ЦП с помощью tfcompile

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

Проверка скомпилированных программ

XLA предоставляет средства самоанализа, которые позволяют проверять сгенерированные программы. Чтобы выгрузить сгенерированные программы, используйте переменную среды 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, используя модуль run_hlo_module и итеративно запуская его в сгенерированных программах.

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

XLA-интерфейсы

Помимо TensorFlow, программы XLA могут создаваться с помощью:

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

Переговоры

Использование XLA из TF с использованием jit_compile=True

Обзор XLA