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 Machine Learning

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 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 thiện về tốc độ và sử dụng bộ nhớ: hầu hết các điểm chuẩn nội bộ chạy nhanh hơn ~ 1,15 lần sau khi bật XLA. Bộ dữ liệu dưới đây được đánh giá trên một GPU NVidia V100 duy nhất:

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 bởi người thực thi TensorFlow. Mỗi hoạt động của TensorFlow có một triển khai nhân GPU được biên dịch sẵn mà người thi hành gửi đến.

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 tính toán được tạo riêng cho mô hình đã cho. 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 cụ thể của mô hình để tối ưu hóa. Ví dụ: chúng ta hãy xem xét một XLA tối ưu hóa 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, đồ 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 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" việc thêm, nhân và giảm vào một nhân GPU. Ngoài ra, thao tác hợp nhất này không ghi các giá trị trung gian được tạo bởi y*zx+y*z vào bộ nhớ; thay vào đó, nó "truyền" kết quả của các tính toán trung gian này trực tiếp đến 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 hiếm nhất trên các bộ tăng tốc phần cứng, vì vậy 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

Tự động phân cụm

Một cách đơn giản nhất để bắt đầu sử dụng XLA trong các mô hình TensorFlow là cho phép tự động phân cụm , tự động tìm các cụm (sơ đồ con được kết nối) trong biểu đồ 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 kích hoạt trên CPU bằng cách sử dụng 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 rõ ràng với tf.feft

Tự động phân cụm là một công cụ tuyệt vời để làm cho mô hình nhanh hơn mà không có bất kỳ thay đổi nào đối với mã, nhưng có thể khó hiểu những thay đổi đã được thực hiện.

API biên dịch rõ ràng cung cấp một điều khiển chi tiết hơn để 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 bằng XLA:

 @tf.function(experimental_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))
 

Các experimental_compile API có ngữ nghĩa phải-biên dịch: hoặc toàn bộ chức năng được biên soạn với XLA, hoặc một errors.InvalidArgumentError ngoại lệ được ném. 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 thang đo 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)
 

Hình dạng có thể khác nhau trên các lần chạy mặc dù:

 @tf.function(experimental_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.

Biên dịch AOT (trước thời gian) cho CPU với 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ã 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 hướng nội cho phép bạn kiểm tra các chương trình được tạo. Để kết xuất các chương trình đã tạo, 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 việc bán phá giá đượ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 Tạo các chương trình XLA , mỗi chương trình được biên dịch. Đính kèm những người khi gửi báo cáo lỗi XLA là vô cùng hữu ích!

  • module_XXXX.ir-*.ll Tạo các tệp trong biểu diễn trung gian LLVM , với nội tại NVPTX .

  • module_XXXX.ptx Tạo các tệp PTX .

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 sinh sản

Một báo cáo lỗi sẽ dễ dàng tái tạo hơn nếu nó bao gồm các bãi chứa 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 phân cụm tự động, 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 báo lỗi, đính kèm nội dung của /tmp/generated (được tham chiếu ở trên).

Nếu có thể, hãy thử cách ly một lỗi với một chương trình XLA 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 được tạo.

Mặt trận 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

đọc thêm