Trang này được dịch bởi Cloud Translation API.
Switch to English

XLA: Tối ưu hóa trình biên dịch cho Máy học

XLA (Accelerated Linear Algebra) là một 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 tới.

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 nhân tính toán được tạo riêng cho mô hình nhất định. 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 tối ưu hóa XLA 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ần 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 biểu đồ để 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*zx+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.

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

Để có 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.txt Cá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-*.ll Các tệp được tạo trong biểu diễn trung gian LLVM , với bản chất NVPTX .

  • module_XXXX.ptx Cá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 được 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 lại nó trên các chương trình đã tạo.

đọc thêm

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