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 chạy chương trình TensorFlow, tất cả các thao tác đều được thực thi riêng lẻ bởi trình thực thi TensorFlow. Mỗi thao tác của TensorFlow đều có một phương thức triển khai nhân GPU được biên dịch trước mà trình thực thi sẽ gửi đến.
XLA cung cấp một chế độ thay thế cho các mô hình chạy: giải pháp này biên dịch biểu đồ TensorFlow thành một trình tự các hạt nhân tính toán được tạo riêng cho mô hình nhất định. 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 cụ thể của mô hình để tối ưu hoá. Ví dụ: hãy xem một tính nă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ần XLA, biểu đồ sẽ khởi chạy 3 hạt nhân: một hạt nhân để nhân, một hạt nhân để thêm và một hạt nhân để rút gọn. 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. Thư viện thực hiện việc này bằng cách "hợp nhất" các tính năng thêm, nhân và rút gọn vào một nhân GPU duy nhất.
Hơn nữa, thao tác kết hợp 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 đó, nó "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 cho người dùng, đồng thời lưu giữ hoàn toàn chúng trong thanh ghi GPU. Fusion là 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 các trình tăng tốc phần cứng. Vì vậy, việc xoá các hoạt động của 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 bằng tf.function(jit_compile=True)
API biên dịch tường minh cung cấp quyền kiểm soát chi tiết để chọn hàm nào cần được 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 một ngoại lệ errors.InvalidArgumentError sẽ được gửi. XLA hiện không thể biên dịch các hàm mà trong đó các kích thước không thể suy luận được: tức là nếu không thể dự đoán kích thước của tất 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 được:
@tf.function
def not_compilable(x):
return tf.unique(x)
Hình dạng có thể thay đổi 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]))
Hãy xem colab hướng dẫn để tham khảo 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.
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)
Sử dụng với chiến lược được phân phối
XLA:Bạn có thể dùng GPU với chiến lược phân phối TF (MirroredStrategy hoặc MultiWorkerMirroredStrategy) bằng cách chú giải 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 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ần thay đổi nào là bật tính năng tự động phân cụm. Tính năng này sẽ tự động tìm các cụm (biểu đồ 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 phân cụm trên GPU 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ính năng tự động phân 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 colab hướng dẫn tự động phân cụm.
Biên dịch AOT (Trước khi diễn ra) cho CPU có tfcompile
Bạn cũng có thể sử dụng một 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ở kiểm tra nội tâm 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 xong, bạn có thể tìm thấy các tệp sau trong /tmp/generated:
module_XXXX.*_optimizations.txtĐã tạ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 thành phần đó khi gửi báo cáo lỗi XLA cực kỳ hữu ích!module_XXXX.ir-*.llCác tệp được tạo ở dạng trình bày trung gian LLVM, với hàm nội tại 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 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 tạo
Báo cáo lỗi sẽ dễ tái tạo hơn nếu bao gồm các tệp kết xuất cho các chương trình XLA đã tạo và tính năng nhúng tự động tạo nhóm đã sử dụng. Để tạo các đối tượng này cho chương trình TensorFlow chạy bằng tính năng tự động tạo 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 thư mục /tmp/generated (được tham chiếu ở trên).
Nếu có thể, hãy cố gắng tách riêng lỗi thành một chương trình XLA bằng cách sử dụng run_hlo_module và chạy lặp lại trên các chương trình được tạo.
Tài liệu đọc thêm
- Tài liệu OpenXLA Tài liệu 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 dành cho nhà phát triển của Google
- Hãy xem 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òn có thể được tạo bởi:
- JAX: Biến đổi thành phần kết hợp của chương trình Python+NumPy
- Julia: Ngôn ngữ Julia cho điện 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