XLA (Accelerated Linear Algebra) là trình biên dịch miền cụ thể cho đại số tuyến tính có thể tăng tốc các mô hình TensorFlow mà không có khả năng thay đổi mã nguồn.
Kết quả là cải thiện tốc độ và sử dụng bộ nhớ: ví dụ: trong trình BERT MLPerf sử dụng 8 GPU Volta V100 sử dụng XLA đã đạt được cải thiện hiệu suất ~ 7 lần và cải thiện kích thước lô ~ 5x:
Giới thiệu
Khi một chương trình TensorFlow được chạy, tất cả các hoạt động được thực thi riêng lẻ bởi người thực thi TensorFlow. Mỗi hoạt động TensorFlow có một triển khai nhân GPU được biên dịch trước mà người thực thi gửi đến.
XLA cung cấp một chế độ thay thế của các mô hình đang chạy: nó biên dịch đồ thị TensorFlow thành một chuỗi các hạt nhân tính toán được tạo riêng cho mô hình nhất định. Bởi vì những hạt nhân này là duy nhất cho mô hình, chúng có thể khai thác thông tin dành riêng cho mô hình để tối ưu hóa. Ví dụ: hãy xem xét một cách tối ưu hóa mà XLA thực hiện trong ngữ cảnh của một phép tính TensorFlow đơn giản:
def model_fn(x, y, z):
return tf.reduce_sum(x + y * z)
Chạy mà không có XLA, đồ thị khởi chạy ba hạt nhân: một cho phép nhân, một cho phép cộng và một cho phép giảm. Tuy nhiên, XLA có thể tối ưu hóa đồ thị để nó tính toán kết quả trong một lần khởi chạy hạt nhân. Nó thực hiện điều này bằng cách "hợp nhất" phép cộng, nhân và giảm thành một nhân GPU duy nhất. Hơn nữa, hoạt động hợp nhất này không ghi các giá trị trung gian được tạo ra bởi y*z và x+y*z vào bộ nhớ; thay vào đó, nó "truyền" trực tiếp kết quả của các phép tính trung gian này đến người dùng của họ trong khi vẫn giữ chúng hoàn toàn trong thanh ghi GPU. Fusion là tối ưu hóa quan trọng nhất của XLA. Băng thông bộ nhớ thường là tài nguyên khan hiếm nhất trên các bộ tăng tốc phần cứng, vì vậy việc xóa các thao tác trên bộ nhớ là một trong những cách tốt nhất để cải thiện hiệu suất.
Bật XLA cho các mô hình TensorFlow
Biên dịch rõ ràng với tf.function(jit_compile=True)
API biên dịch rõ ràng cung cấp một kiểm soát chi tiết để chọn các hàm nên được biên dịch. Ví dụ: hàm TensorFlow sau đây thực hiện đào tạo MNIST được biên dịch với 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 có ngữ nghĩa phải biên dịch : toàn bộ hàm được biên dịch bằng XLA hoặc lỗi. Ngoại errors.InvalidArgumentError được ném ra. XLA hiện không thể biên dịch các hàm trong đó kích thước không thể suy ra được: nghĩa là, nếu không thể suy ra kích thước của tất cả các tensors mà không chạy toàn bộ tính toán. Ví dụ, hàm sau sẽ không biên dịch:
@tf.function
def not_compilable(x):
return tf.unique(x)
Mặc dù vậy, hình dạng có thể khác nhau qua các lần chạy:
@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]))
Xem chuyên mục hướng dẫn để có ví dụ sử dụng chi tiết hơn và video hướng dẫn về cách sử dụng jit_compile=True .
Cách sử dụng với Keras
Đối với các mô hình Keras, jit_compile=True có thể được đặt làm đối số cho model.compile :
model.compile(optimizer="adam", jit_compile=True)
Sử dụng với chiến lược phân tán
XLA: GPU có thể được sử dụng với chiến lược phân tán TF ( MirroredStrategy hoặc MultiWorkerMirroredStrategy ) bằng cách chú thích hàm bước với 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)
Tự động phân nhóm
Một cách đơn giản để bắt đầu sử dụng XLA trong các mô hình TensorFlow mà không có bất kỳ thay đổi nào là bật tính năng tự động phân cụm , tự động tìm các cụm (đồ thị con được kết nối) trong các hàm TensorFlow có thể được biên dịch và thực thi bằng XLA. Tự động phân cụm trên GPU có thể được bật bằng cách đặt biến môi trường TF_XLA_FLAGS :
$ TF_XLA_FLAGS=--tf_xla_auto_jit=2 path/to/your/tf/program
Tự động phân cụm hiện được tối ưu hóa cho khối lượng công việc GPU, nhưng nó cũng có thể được bật trên CPU bằng cách sử dụng thêm cờ --tf_xla_cpu_global_jit :
$ TF_XLA_FLAGS="--tf_xla_auto_jit=2 --tf_xla_cpu_global_jit" path/to/your/program
Để biết ví dụ sử dụng chi tiết, hãy xem chuyên mục hướng dẫn phân cụm tự động .
Biên dịch AOT (Trước thời gian) cho CPU với trình biên dịch tfcompile
Bạn cũng có thể sử dụng công cụ tfcompile độc lập để chuyển đổi đồ thị TensorFlow thành mã thực thi (chỉ dành cho CPU x86-64).
Kiểm tra các chương trình đã biên dịch
XLA cung cấp các phương tiện kiểm tra nội quan cho phép bạn kiểm tra các chương trình đã tạo. Để kết xuất các chương trình đã tạo, hãy sử dụng biến môi trường XLA_FLAGS :
$ XLA_FLAGS="--xla_dump_to=/tmp/generated" TF_XLA_FLAGS="--tf_xla_auto_jit=2" my/tensorflow/program
Sau khi kết xuất được thực hiện, bạn có thể tìm thấy các tệp sau trong /tmp/generated :
module_XXXX.*_optimizations.txtCác chương trình XLA đã tạo , một chương trình cho mỗi cụm đã biên dịch. Đính kèm chúng khi gửi báo cáo lỗi XLA là vô cùng hữu ích!module_XXXX.ir-*.llCác tệp được tạo trong biểu diễn trung gian LLVM , với bản chất NVPTX .module_XXXX.ptxCác tệp PTX đã tạo.
Bạn cũng có thể kết xuất biểu đồ trực quan hóa việc nhúng các cụm XLA bên trong biểu đồ TensorFlow với:
$ TF_DUMP_GRAPH_PREFIX=/tmp/generated TF_XLA_FLAGS="--tf_xla_clustering_debug"
Báo cáo lỗi có thể lặp lại
Báo cáo lỗi dễ dàng tái tạo hơn nhiều nếu nó bao gồm các kết xuất cho các chương trình XLA đã tạo và nhúng tự động phân cụm đã sử dụng. Để tạo chúng cho chương trình TensorFlow chạy với tính năng tự động phân cụm, hãy khởi chạy:
$ 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"
Khi gửi lỗi, hãy đính kèm nội dung của /tmp/generated (đã tham khảo ở trên).
Nếu có thể, hãy cố gắng tách một lỗi thành một chương trình XLA duy nhất bằng cách sử dụng replay_computation và chạy lặp đi lặp lại nó trên các chương trình đã tạo.
đọc thêm
- Các sự cố đã biết Danh sách các sự cố đã biết với XLA
- Kiến trúc XLA : Tổng quan về kiến trúc XLA
- XLA - TensorFlow, Biên soạn : Đọc trên Blog nhà phát triển của Google
- Kiểm tra nguồn XLA trên Github!
Giao diện người dùng XLA
Ngoài TensorFlow, các chương trình XLA có thể được tạo bằng:
- JAX : Các phép biến đổi có thể dùng một lần của các chương trình Python + NumPy
- Julia : Ngôn ngữ Julia dành cho máy tính khoa học
- PyTorch : Khung PyTorch
- Nx : Thư viện tính toán số cho ngôn ngữ lập trình Elixir