OpenXLA là một trình biên dịch dành riêng cho miền cho đại số tuyến tính có thể tăng tốc các mô hình TensorFlow mà có thể không cần thay đổi mã nguồn.
Giới thiệu
Khi một chương trình TensorFlow chạy, tất cả các thao tác sẽ được thực thi riêng lẻ bởi trình thực thi TensorFlow. Mỗi toán tử TensorFlow đều có một quá trình triển khai nhân GPU được biên dịch trước mà trình thực thi sẽ điều phối đến.
XLA cung cấp một chế độ chạy mô hình thay thế: biên dịch biểu đồ TensorFlow thành một trình tự hạt nhân tính toán được tạo riêng cho mô hình đã cho. Vì các hạt nhân này là dành riêng cho mô hình, nên chúng có thể khai thác thông tin dành riêng cho mô hình để tối ưu hoá. Ví dụ: hãy xem một hoạt động tối ưu hoá mà XLA thực hiện trong bối cảnh một phép tính TensorFlow đơn giản:
def model_fn(x, y, z):
return tf.reduce_sum(x + y * z)
Khi chạy mà không có XLA, biểu đồ sẽ khởi chạy 3 hạt nhân: một hạt nhân cho phép nhân, một hạt nhân cho phép cộng và một hạt nhân cho phép trừ. Tuy nhiên, XLA có thể tối ưu hoá biểu đồ để tính toán kết quả trong một lần khởi chạy hạt nhân. Phương thức này thực hiện việc này bằng cách "hợp nhất" phép cộng, phép nhân và phép trừ vào một hạt nhân GPU duy nhất.
Hơn nữa, toán tử hợp nhất này không ghi các giá trị trung gian do y*z và x+y*z tạo ra vào bộ nhớ; thay vào đó, toán tử này "truyền trực tuyến" kết quả của các phép tính trung gian này trực tiếp đến người dùng trong khi vẫn giữ hoàn toàn các giá trị đó trong các thanh ghi GPU. Fusion là tính năng tối ưu hoá 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 trình tăng tốc phần cứng, vì vậy, việc xoá các thao tác 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 tường minh bằng tf.function(jit_compile=True)
API biên dịch rõ ràng cung cấp khả năng kiểm soát chi tiết để chọn hàm cần biên dịch. Ví dụ: hàm TensorFlow sau đây thực hiện việc huấn luyện MNIST được biên dịch bằng 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 ngoại lệ errors.InvalidArgumentError sẽ được gửi. XLA hiện không thể biên dịch các hàm mà kích thước không có thể suy luận: tức là không thể suy luận kích thước của tất cả các tensor mà không chạy toàn bộ phép tính. Ví dụ: hàm sau đây sẽ không biên dịch:
@tf.function
def not_compilable(x):
return tf.unique(x)
Tuy nhiên, hình dạng có thể khác nhau giữa 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]))
Hãy xem tài liệu hướng dẫn về Colab để biết ví dụ chi tiết hơn về cách sử dụng 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, bạn có thể đặt jit_compile=True làm đối số cho model.compile:
model.compile(optimizer="adam", jit_compile=True)
Cách sử dụng với chiến lược phân phối
XLA:GPU có thể được sử dụng với chiến lược phân phối TF
(MirroredStrategy
hoặc
MultiWorkerMirroredStrategy)
bằng cách chú thích hàm bước bằng 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 cụ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ần thay đổi gì là bật tính năng tự động tạo cụm. Tính năng này sẽ tự động tìm các cụm (các đồ 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. Bạn có thể bật tính năng tự động tạo cụm trên GPU bằng cách thiết lập biến môi trường TF_XLA_FLAGS:
$ TF_XLA_FLAGS=--tf_xla_auto_jit=2 path/to/your/tf/program
Tính năng tự động tạo cụm hiện được tối ưu hoá cho khối lượng công việc của GPU, nhưng bạn cũng có thể bật tính năng này 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ụ chi tiết về cách sử dụng, hãy xem hướng dẫn tự động tạo cụm trên Colab.
Biên dịch AOT (trước khi thực thi) cho CPU bằng tfcompile
Bạn cũng có thể sử dụng công cụ tfcompile độc lập để chuyển đổi biểu đồ TensorFlow thành mã có thể 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 cơ sở tự kiểm tra 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 thực hiện việc kết xuất, bạn có thể tìm thấy các tệp sau trong /tmp/generated:
module_XXXX.*_optimizations.txtTạo chương trình XLA, mỗi chương trình cho một cụm được biên dịch. Việc đính kèm những tệp đó khi gửi báo cáo lỗi XLA sẽ rất hữu ích!module_XXXX.ir-*.llTạo tệp ở dạng trình bày trung gian LLVM, với các hàm nội tại NVPTX.module_XXXX.ptxTạo tệp PTX.
Bạn cũng có thể kết xuất biểu đồ trực quan hoá việc nhúng các cụm XLA bên trong biểu đồ TensorFlow bằng:
$ TF_DUMP_GRAPH_PREFIX=/tmp/generated TF_XLA_FLAGS="--tf_xla_clustering_debug"
Báo cáo lỗi có thể tái hiện
Báo cáo lỗi dễ tái hiện hơn nhiều nếu báo cáo đó bao gồm các tệp báo lỗi cho các chương trình XLA đã tạo và tính năng nhúng tự động phân cụm được sử dụng. Để tạo các tệp này cho một chương trình TensorFlow chạy bằng tính năng tự động tạo cụm, hãy 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 thư mục /tmp/generated
(được tham chiếu ở trên).
Nếu có thể, hãy cố gắng tách riêng một lỗi cho một chương trình XLA bằng cách sử dụng run_hlo_module và chạy lặp lại lỗi đó trên các chương trình đã tạo.
Tài liệu đọc thêm
- Tài liệu về OpenXLA Tài liệu về OpenXLA
- Vấn đề đã biết Danh sách các vấn đề đã biết với XLA+TF
- XLA – TensorFlow, Compiled: Đọc trên Blog Google Developers
- Hãy xem nguồn XLA trên GitHub!
Giao diện người dùng XLA
Ngoài TensorFlow, bạn có thể tạo chương trình XLA bằng:
- JAX: Các phép biến đổi có thể kết hợp của chương trình Python+NumPy
- Julia: Ngôn ngữ Julia dùng cho tính toán khoa học
- PyTorch: Khung PyTorch
- Nx: Thư viện điện toán số cho ngôn ngữ lập trình Elixir