Ngày Cộng đồng ML là ngày 9 tháng 11! Tham gia với chúng tôi để cập nhật từ TensorFlow, JAX, và nhiều hơn nữa Tìm hiểu thêm

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

XLA (Accelerated Linear Algebra) là 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à những cải tiến về tốc độ và bộ nhớ sử dụng: ví dụ trong Bert MLPerf trình sử dụng 8 Volta V100 GPU sử dụng XLA đã đạt được một ~ cải thiện hiệu suất 7x và cải thiện kích thước ~ 5x mẻ:

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 đến.

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 hạt nhân tính toán được tạo riêng cho mô hình nhất định. Bởi 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 một cách tối ưu hóa mà XLA thực hiện trong ngữ cảnh của một phép tính 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 đồ thị để 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 viết ra các giá trị trung gian sản xuất 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 soạn 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))

Các jit_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 chức năng mà kích thước không phải là inferrable: có nghĩa là, nếu nó không thể suy ra kích thước của tất cả tensors mà không cần 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 colab hướng dẫn cho một ví dụ sử dụng chi tiết hơn, và một video hướng dẫn trên jit_compile=True sử dụng.

Tự động phân nhóm

Một cách đơn giản để bắt đầu sử dụng trong các mô hình XLA TensorFlow mà không cần bất kỳ thay đổi được bật tự động clustering, mà tự động tìm kiếm cụm (đồ thị con được kết nối) trong chức năng TensorFlow có thể được biên soạn và thực hiện sử dụng XLA. Auto-clustering trên GPU có thể được kích hoạt bằng cách thiết lập TF_XLA_FLAGS biến môi trường:

$ TF_XLA_FLAGS=--tf_xla_auto_jit=2 path/to/your/tf/program

Auto-clustering 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 bổ sung bằng cách sử dụng lá cờ --tf_xla_cpu_global_jit :

$ TF_XLA_FLAGS="--tf_xla_auto_jit=2 --tf_xla_cpu_global_jit" path/to/your/program

Đối với một ví dụ sử dụng chi tiết xem các auto-clustering hướng dẫn colab .

AOT (Ahead-of-time) biên soạn cho CPU với tfcompile

Bạn cũng có thể sử dụng một độc tfcompile công cụ, chuyển đổi TensorFlow graph vào mã thực thi (chỉ 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. Để đổ các chương trình tạo ra, sử dụng các 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 bán phá giá được thực hiện, bạn có thể tìm thấy các tập tin sau đây trong /tmp/generated :

  • module_XXXX.*_optimizations.txt tạo chương trình XLA , mỗi mỗi cụm biên soạn. Đí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 tập tin được tạo trong LLVM đại diện trung gian, với NVPTX intrinsics.

  • module_XXXX.ptx tạo PTX tập tin.

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 đã 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 nộp lỗi, đính kèm các nội dung của /tmp/generated thư mục (tham chiếu ở trên).

Nếu có thể, cố gắng cô lập lỗi cho một chương trình XLA đơn bằng replay_computation và lặp đi lặp lại chạy nó trên các chương trình tạo ra.

đọ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 : composable biến đổi của các chương trình Python + NumPy
  • Julia : Các Julia ngôn ngữ cho tính toán khoa học
  • PyTorch khuôn khổ PyTorch:
  • Nx : số học máy tính thư viện cho các ngôn ngữ lập trình Elixir

Đàm thoại

Sử dụng XLA từ TF sử dụng jit_compile=True

Tổng quan về XLA