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
- Архитектура XLA : Обзор архитектуры XLA
- XLA - TensorFlow, скомпилировано : читайте в блоге разработчиков Google
- Ознакомьтесь с исходным кодом XLA на Github!
Внешние интерфейсы XLA
Помимо TensorFlow, программы XLA могут быть созданы с помощью: