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
- XLA Архитектура : Обзор архитектуры XLA
- XLA - TensorFlow, Составитель : Чтение на Google Developers Blog
- Проверьте источник XLA на Github!
Внешние интерфейсы XLA
Помимо TensorFlow, программы XLA могут быть созданы с помощью:
- JAX : наборный преобразования программ Python + NumPy
- Джулия : Джулия язык для научных вычислений
- PyTorch рамки PyTorch:
- Nx : Численное вычисление библиотеки для языка программирования Elixir