Есть вопрос? Присоединяйтесь к сообществу на форуме TensorFlow. Посетите форум.

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

XLA (Accelerated Linear Algebra) - это предметно-ориентированный компилятор для линейной алгебры, который может ускорять модели 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 исключение 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]))

Более подробный пример использования см. В учебном пособии .

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

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

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

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

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

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

Компиляция AOT (Ahead-of-time) для ЦП с tfcompile

Вы также можете использовать автономный инструмент tfcompile , который преобразует график tfcompile в исполняемый код (только для процессоров 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.