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

Оптимизируйте свои подборки Сохраняйте и классифицируйте контент в соответствии со своими настройками.

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

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

Введение

Когда программа 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]))

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

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

Для моделей 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. Автокластеризацию на GPU можно включить, установив переменную среды 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, используя replay_computation и итеративно запуская ее в сгенерированных программах.

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

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

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

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

Переговоры

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

Обзор XLA