XLA (Đại số tuyến tính tăng tốc) là trình biên dịch dành riêng cho miền dành 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à những cải tiến về tốc độ và mức 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ô ~5 lần:
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 của TensorFlow đều có triển khai nhân GPU được biên dịch sẵn mà người thực thi gửi tới.
XLA cung cấp một chế độ thay thế cho các mô hình đang chạy: nó biên dịch biểu đồ 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 đã cho. Bởi vì các 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 XLA tối ưu hóa thực hiện trong bối cảnh tính toán TensorFlow đơn giản:
def model_fn(x, y, z):
return tf.reduce_sum(x + y * z)
Chạy mà không có XLA, biểu đồ 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 rút gọn. Tuy nhiên, XLA có thể tối ưu hóa biểu đồ để tính toán kết quả trong một lần khởi chạy kernel. Nó thực hiện điều này bằng cách "hợp nhất" phép cộng, phép nhân và phép rút gọn thành một nhân GPU duy nhất. Ngoài ra, thao tác 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 đó, nó "truyền trực tiếp" kết quả của các phép tính trung gian này tới người dùng của họ trong khi vẫn giữ chúng hoàn toàn trong các 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 bộ tăng tốc phần cứng, do đó, loại bỏ 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.
Kích hoạ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 khả năng kiểm soát chi tiết để chọn chức năng nào sẽ đượ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 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ộ chức năng được biên dịch bằng XLA hoặc một ngoại lệ errors.InvalidArgumentError được đưa 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 : nghĩa là, nếu không thể suy ra kích thước của tất cả các tenxơ mà không cần chạy toàn bộ tính toán. Ví dụ: chức năng 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 trong 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 colab hướng dẫn để biết 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 .
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 chức năng 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 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ó 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 (sơ đồ 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 colab hướng dẫn tự động phân cụm .
Biên dịch AOT (Ahead-of-time) cho CPU với 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ã 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ở 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 quá trình 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 được biên dịch. Đính kèm những thứ đó 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 ở dạng biểu diễn trung gian LLVM , với 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 hóa 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 tái sản xuất
Một báo cáo lỗi sẽ dễ 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 được tạo và nhúng tự động phân cụm được 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 thư mục /tmp/generated (được tham chiếu ở trên).
Nếu có thể, hãy cố gắng cô lập một lỗi với một chương trình XLA đơn lẻ bằng cách sử dụng run_hlo_module và chạy 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 Google Developers Blog
- 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ởi:
- JAX : Các phép biến đổi có thể kết hợp của các chương trình Python + NumPy
- Julia : Ngôn ngữ Julia cho máy tính khoa học
- PyTorch : Khung PyTorch
- Nx : Thư viện điện toán số cho ngôn ngữ lập trình Elixir