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

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

Результаты улучшения в скорости и памяти: например , в BERT MLPerf представления с использованием 8 Вольта V100 графических процессоров с использованием XLA достигла ~ повышение производительности 7x и улучшение размер ~ 5x партия:

Вступление

При запуске программы 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))

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

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

Смотрите учебник colab для более подробного примера использования, а также видео - jit_compile=True учебник по jit_compile=True использование.

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

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

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

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

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

Для подробного примера использования см авто-кластеризация учебник colab .

АОТ (вперед-оф-времени) компиляция для CPU с tfcompile

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

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

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 с помощью replay_computation и итеративно запустить его на сгенерированных программ.

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

Внешние интерфейсы XLA

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

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

Переговоры

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

Обзор XLA