利用 XLA 将 GPU 性能推向极限

文 / Toby Boyd、Yanan Cao、Sanjoy Das、Thomas Joerg、Justin Lebar

XLA 是 TensorFlow 图表的编译器,只需更改极少的源代码,便可加速您的 TensorFlow ML 模型。这篇文章将介绍 XLA,并说明如何在您自己的代码中试用 XLA。

在使用 NVIDIA® Tesla® V100 GPU 训练 ResNet50 v1.0 时,相比 TF 1.11(未配备 XLA),TensorFlow 1.12(配备 XLA)的性能有显著提升:合成数据为每秒 10526 张图像,真实数据为每秒 10267 张图像(复制说明请见附录)。我们在各种内部模型上观察到速度提升(从 1.13 倍至 3.04 倍)。

image

XLA:成功编译 TensorFlow!

通常情况下,当您运行 TensorFlow 图表时,所有运算都由 TensorFlow 图表执行器单独执行。每个运算都会安装由图表执行器分派的预编译 GPU 内核(随附于 TensorFlow 二进制文件中)。

XLA 提供了另一种运行 TF 模型的模式:这种模式会将您的 TensorFlow 图表编译成专为您的模型生成的 GPU 内核序列。由于这些是您程序独有的内核,因此它们可以利用模型的特定信息进行优化。

举一个例子,我们一起看看 XLA 在简单 TensorFlow 计算环境中的优化过程:
1 def model_fn(x,y,z):
2 return tf.reduce_sum(x + y * z)

如果运行模型时不使用 XLA,图表会启动三个内核,分别用于乘法、加法和减法。

但是,XLA 可以优化图表,以便在启动单个内核时计算结果。方法是将加法、乘法和减法 “融合” 到单个 GPU 内核中。此外,这种融合运算不会将 yz 和 x+yz 生成的中间值写入内存,而是将这些中间计算的结果直接 “流式传输” 给用户,并完整保存在 GPU 寄存器中。

融合是 XLA 最重要的一种优化方式。内存带宽通常是硬件加速器上最稀缺的资源,因此删除内存运算是提升性能的最佳方法之一。

在模型中使用 XLA

XLA 使用 xla.compile API,让您可以在部分 TensorFlow 图表上显式调用 XLA 编译器。xla.compile 会接受生成 TensorFlow 计算的 Python 函数,然后连接所生成的计算以供 XLA 编译。xla.compile 还会返回张量列表,其中每个张量对应传入函数构建的计算结果,但现在会立即由 XLA 优化。

因此,通过调用 xla.compile,您可以使用 XLA 运行以上由 model_fn 生成的计算,如下所示:
1 from tensorflow.contrib.compiler import xla
2
3 def model_fn (x, y, z):
4 return tf.reduce_sum (x + y * z)
5
6 def create_and_run_graph ():
7 with tf.Session () as sess:
8 x = tf.placeholder (tf.float32, name=‘x’)
9 y = tf.placeholder (tf.float32, name=‘y’)
10 z = tf.placeholder (tf.float32, name=‘z’)
11 result = xla.compile (computation=model_fn, inputs=(x, y, z))[0]
12 # result is a normal Tensor (albeit one that is computed by an XLA
13 # compiled executable) and can be used like any other Tensor.
14 result = tf.add (result, result)
15 return sess.run (result, feed_dict={ … })

您可以使用命令行标记(或其他任意逻辑)来控制是否由 XLA 编译计算。模型通常按照以下方式调用 xla.compile 以便完成简单实验。
1 if should_use_xla ():
2 result = xla.compile (model_fn, (x, y, z))[0]
3 else:
4 result = model_fn (x, y, z)

我们已经设置 colab,让您可以在稍微复杂一些的模型上使用 xla.compile。

xla.compile 不是在 TensorFlow 子图表上调用 XLA 的唯一方法;具体来说,有一些方法可以让 TensorFlow 自动 找到与 XLA 兼容的子图表并使用 XLA 进行编译,但我们不会在这篇文章中讨论这些方法。

使用 XLA 的注意事项

第一,XLA GPU 后端目前仍处于实验阶段,虽然我们没有发现任何重大问题,但其尚未进行广泛的生产使用测试。

第二,xla.compile 仍不适用于 model.fit 等 Keras 高级 API(但您可以使用 Keras 运算),也不支持 Eager 模式。我们正在积极开发 API 以便在这些模式下启用 XLA;敬请期待。

第三,XLA 无法编译所有 TensorFlow 图表;只有具有以下属性的图表才能传递给 xla.compile。

所有运算都必须具有可推断的形状

XLA 需要能够在给定计算输入的情况下,推断出其编译的所有运算的形状。因此,如果模型函数生成的 Tensor 具有不可推断的形状,则运行时将会出现错误,进而导致运行失败。(在这个例子中,tf.expand_dims 的输出形状取决于 random_dim_size,但其无法在给定 x、y 和 z 的情况下推断出来。)

请注意,由于 XLA 是 JIT 编译器,形状会因不同的运行过程而异,前提是能够根据给定的集群输入条件推断出来。这个例子很好地体现了这一点。

XLA 必须支持所有运算

并非所有 TensorFlow 运算都可由 XLA 编译,如果模型中有 XLA 不支持的运算,XLA 编译就会失败。例如,XLA 不支持 tf.where 运算,因此如果您的模型函数包含此运算,使用 xla.compile 运行模型时便会失败。

XLA 支持的每项 TensorFlow 运算都可在 tensorflow/compiler/tf2xla/kernels/ 中调用 REGISTER_XLA_OP,因此您可以使用 grep 来搜索 REGISTER_XLA_OP 宏实例,以查找支持的 TensorFlow 运算列表。

附录

在 Google 基准下的表现
下方是所有 XLA 团队基准模型(在 V100 GPU 上运行)的 TensorFlow 在使用和不使用 XLA 时的相对加速 / 减速图。我们没有任何保留;这是我们今天在评估编译器时使用的全部基准。

image

图中每个竖条代表一个完整的模型,例如 “resnet50 训练图像 / 秒” 或 “基于 Google 内部模型的推理吞吐量”。X 轴按加速情况排序。

您的情况可能有所不同,特别是我们已经根据其中很多基准专门对 XLA 进行了优化!但是,其中很多基准可以立即使用,而且效果非常出色,我们还将持续作出改进。

复制 ResNet50 v1.0 基准
以下部分将介绍如何设置 Google Cloud 实例和执行 ResNet50 基准。

准备数据
只有进行真实数据测试时需要完成这个步骤,并可能需时几小时。我们建议只在 CPU 实例中执行此步骤,以降低计算成本。根据 imagenet_to_gcs.py 说明创建 TFRecord 格式的 imagenet 数据,并将其发送至 Google Cloud Storage 存储分区。

创建 GCE 实例
下方代码片段会在使用八个 Tesla® V100 GPU 的 Google Cloud Platform 上创建 Google 深度学习 VM。
1 export INSTANCE_NAME=“xla-benchmark-8xV100”
2 export IMAGE_FAMILY=“tf-1-12-cu100”
3 export PROJECT_NAME=""
4
5 gcloud beta compute instances create $INSTANCE_NAME \
6 --project=$PROJECT_NAME \
7 --machine-type=n1-standard-64 \
8 --maintenance-policy=TERMINATE \
9 --accelerator=type=nvidia-tesla-v100,count=8 \
10 --tags=http-server,https-server \
11 --image-family=$IMAGE_FAMILY \
12 --image-project=deeplearning-platform-release \
13 --boot-disk-size=100GB \
14 --boot-disk-type=pd-ssd \
15 --local-ssd interface=nvme \
16 --local-ssd interface=nvme \
17 --local-ssd interface=nvme \
18 --local-ssd interface=nvme \
19 --metadata install-nvidia-driver=True
20
21 ## Combines the 4 local nvme SSD drives into a single RAID 0 drive.
22 # Install raid management tool.
23 sudo apt-get update && sudo apt-get install mdadm --no-install-recommends
24
25 # Creates RAID 0 array.
26 sudo mdadm --create /dev/md0 --level=0 --raid-devices=4
27 /dev/nvme0n1 /dev/nvme0n2 /dev/nvme0n3 /dev/nvme0n4
28
29 # Formats and mounts the array.
30 sudo mkfs.ext4 -F /dev/md0
31 sudo mkdir -p /data/imagenet
32 sudo mount /dev/md0 /data
33 sudo chmod a+w /data
34
35 # Installs custom TensorFlow 1.12 binary with AVX2. Binary included on
36 # the image already has XLA but the custom binary is compiled with AVX2.
37 sudo pip install --force-reinstall https://storage.googleapis.com/t … mu-linux_x86_64.whl

执行基准测试
1 gcloud compute SSH $INSTANCE_NAME
2
3 # Clone TensorFlow benchmark repository.
4 git clone https://github.com/tensorflow/benchmarks.git && cd benchmarks
5 git reset --hard 1e7d788042dfc6d5e5cd87410c57d5eccee5c664
6 cd scripts/tf_cnn_benchmarks
7
8 ## Synthetic data test
9
10 # 8 GPUs
11 python tf_cnn_benchmarks.py \
12 --batch_size=364 \
13 --num_batches=100 \
14 --model=resnet50 \
15 --optimizer=momentum \
16 --variable_update=replicated \
17 --all_reduce_spec=nccl \
18 --use_fp16=True \
19 --nodistortions \
20 --gradient_repacking=2 \
21 --compute_lr_on_cpu=True \
22 --single_l2_loss_op=True \
23 --xla_compile=True \
24 --num_gpus=8 \
25 --loss_type_to_report=base_loss
26
27 # 1 GPU
28 python tf_cnn_benchmarks.py \
29 --batch_size=364 \
30 --num_batches=100 \
31 --model=resnet50 \
32 --optimizer=momentum \
33 --use_fp16=True \
34 --nodistortions \
35 --compute_lr_on_cpu=True \
36 --single_l2_loss_op=True \
37 --xla_compile=True \
38 --loss_type_to_report=base_loss
39
40 ## Real data test
41 # add --data_dir=/data/imagenet to the 1 or 8 GPU command.


大白白 2018-11-22 15:22:03