Kết xuất các phép tính HLO

Kết xuất HLO là một bản trình bày bằng văn bản về các mô-đun HLO ở các giai đoạn khác nhau của quá trình tính toán. Thông tin này hữu ích cho việc gỡ lỗi và bạn thường cần đưa thông tin này vào báo cáo lỗi. Đây thường là một tệp văn bản mà con người có thể đọc được, liệt kê các chỉ dẫn HLO và thuộc tính của chúng. Đôi khi, các mô-đun HLO được kết xuất dưới dạng:

HloProto: Tệp vùng đệm giao thức, là một định dạng có cấu trúc và dễ đọc hơn đối với máy.
HloSnapshot: Mô-đun HLO cùng với các đầu vào của mô-đun đó. Để phát lại HLO, đôi khi bạn cần các đầu vào thực tế được đưa vào một phép tính nhất định thay vì dữ liệu ngẫu nhiên.

Bạn có thể sử dụng cờ XLA để chỉ định và nhận các kết xuất. Trong hầu hết các trường hợp, bạn có thể đặt giá trị này bằng một biến môi trường. JAX cũng cung cấp một cách có lập trình để in kết xuất HLO.

Thực thi cục bộ

Sử dụng biến môi trường

Bạn có thể đặt biến môi trường XLA_FLAGS bằng các cờ cần thiết để nhận được các kết xuất. Điều này áp dụng cho JAX, TensorFlow và PyTorch/XLA.

Để kết xuất các mô-đun HLO và thông tin gỡ lỗi khác vào một thư mục cụ thể, hãy chạy chương trình của bạn bằng cờ --xla_dump_to:

XLA_FLAGS="--xla_dump_to=DIRECTORY_PATH"

Ví dụ: bạn có thể sử dụng /tmp hoặc /tmp/xladump làm đường dẫn.

Theo mặc định, thao tác này sẽ kết xuất các mô-đun HLO dưới dạng văn bản, ngay từ đầu và cuối quy trình tối ưu hoá.

Bạn cũng có thể chỉ định rõ ràng định dạng:

Kết xuất văn bản

XLA_FLAGS="--xla_dump_hlo_as_text --xla_dump_to=DIRECTORY_PATH"

HLO proto

XLA_FLAGS="--xla_dump_hlo_as_proto --xla_dump_to=DIRECTORY_PATH"

Ảnh chụp nhanh HLO

XLA_FLAGS="--xla_dump_hlo_snapshots --xla_dump_to=DIRECTORY_PATH"

Kết xuất đồ thị bằng máy chủ graphviz (chỉ hoạt động tốt đối với đồ thị nhỏ)

XLA_FLAGS="--xla_dump_hlo_as_url --xla_dump_to=DIRECTORY_PATH"

Kết xuất đồ thị thành tệp HTML (chỉ hoạt động hiệu quả đối với đồ thị nhỏ)

XLA_FLAGS="--xla_dump_hlo_as_html --xla_dump_to=DIRECTORY_PATH"

Đối với các biểu đồ lớn hơn, bạn có thể sử dụng interactive_graphviz để trực quan hoá các phần của biểu đồ.

Kết xuất các lượt truyền trung gian cụ thể

Ngoài các HLO được tối ưu hoá trước / tối ưu hoá cuối cùng tiêu chuẩn, bạn cũng có thể kết xuất trạng thái của HLO sau một lần truyền trình biên dịch cụ thể.

XLA_FLAGS="--xla_dump_hlo_pass_re=regex --xla_dump_to=DIRECTORY_PATH"

Các mô-đun HLO sẽ được kết xuất cho những lượt truyền có tên khớp với biểu thức chính quy (regex). Ví dụ: bạn có thể quan sát các HLO thu được từ các lượt truyền liên quan đến việc phân vùng SPMD bằng cách:

XLA_FLAGS="--xla_dump_to=DIRECTORY_PATH --xla_dump_hlo_pass_re=spmd|propagation"

Để kết xuất kết quả sau mỗi lượt XLA (thao tác này sẽ tạo ra nhiều tệp), bạn có thể đặt:

XLA_FLAGS="--xla_dump_to=DIRECTORY_PATH --xla_dump_hlo_pass_re=.*"

Các lựa chọn dành riêng cho JAX

Theo phương thức lập trình trong JAX

Thay vì truyền cờ hoặc biến môi trường, bạn cũng có thể kết xuất HLO theo cách lập trình bằng cách sử dụng các API lower và compile của JAX.

Tìm nạp cục bộ HLO gốc chưa được tối ưu hoá bằng cách:

jax.jit(f).lower(*args).as_text('hlo')

Để kết xuất vào các tệp trong quá trình biên dịch HLO, hãy chỉ định:

compilation_args = {
    'xla_dump_to': DIRECTORY_PATH,
    'xla_dump_hlo_pass_re': 'spmd|propagation', # or some other pass filter
    ...
    }

jax.jit(f).lower(*args).compile(compilation_args)

Kết xuất jaxpr

jaxprs là biểu thị trung gian của JAX cho dấu vết chương trình. Để kết xuất thông tin này, hãy đặt các biến môi trường:

JAX_DUMP_IR_TO="DIRECTORY_PATH" JAX_DUMP_IR_MODES=jaxpr

Tìm hiểu thêm trong tài liệu JAX về Xuất và chuyển đổi các phép tính theo giai đoạn: Gỡ lỗi.

Google Colab

Biến môi trường

Trong ô được thực thi đầu tiên của sổ tay (vì các biến môi trường và cờ dòng lệnh thường chỉ được xử lý một lần, ví dụ: tại thời điểm nhập mô-đun hoặc thời điểm khởi tạo phụ trợ XLA), hãy thêm XLA_FLAGS được mô tả ở trên bằng os.environ, ví dụ:

import os
os.environ['XLA_FLAGS'] = "--xla_dump_to=DIRECTORY_PATH"

Thao tác này sẽ kết xuất quá trình tính toán sang DIRECTORY_PATH, ví dụ: /tmp. Trên Colab, hãy chuyển đến trình duyệt "Tệp" trong thanh bên trái để xem và truy cập vào thư mục này.

Bạn có thể sử dụng tất cả các cờ được đề cập trong phần Thực thi cục bộ.

Các lựa chọn dành riêng cho JAX

Tương tự như việc thực thi cục bộ; đối với hoạt động tự xem xét tương tác trực tiếp, bạn có thể in trực tiếp HLO được tối ưu hoá trước của một phép tính:

def f(x):
    return jax.numpy.sin(jax.numpy.cos(x))

c = jax.jit(f).lower(3.).compiler_ir('hlo')

print(c.as_hlo_text())

Bạn cũng có thể in trực tiếp HLO được tối ưu hoá của một phép tính:

def optimized_HLO(f, *args, platform=None):
    print(jax.jit(f).lower(*args).compile().as_text())

def f(x):
    return jax.numpy.sin(jax.numpy.cos(x))

optimized_HLO(f, 1.0)

Kết xuất tất cả/một số phép tính

Nếu bạn muốn xem mọi thứ trong một tệp kết xuất, bao gồm cả tất cả các bản biên dịch nhỏ, hãy đặt biến môi trường JAX:

JAX_COMPILER_DETAILED_LOGGING_MIN_OPS=0

Mosaic

Mosaic là một trình biên dịch cho phần phụ trợ TPU Pallas và phần phụ trợ GPU Pallas thử nghiệm. Để kết xuất quá trình tính toán ghép ảnh, hãy đặt cờ sau:

--xla_mosaic_dump_to=/tmp/mosaic_dumps

Hoặc đặt các đối số khởi động TPU làm biến môi trường:

export LIBTPU_INIT_ARGS="--xla_mosaic_dump_to=/tmp/mosaic_dumps"

Hãy xem tài liệu JAX về Pallas và Mosaic để tìm hiểu thêm.

More with HLO Dumps

Tìm được phép tính phù hợp

Thông thường, nhiều phép tính sẽ bị loại bỏ. Các tệp kết xuất được đặt tên rõ ràng bằng "tên tính toán" JAX, Tensorflow hoặc PyTorch/XLA được gọi trong nhật ký, giúp bạn dễ dàng xác định các tệp HLO có liên quan. Ví dụ:

1624325116260738.module_0065.pmap__unnamed_wrapped_function_.186875.before_optimizations.txt

Nếu không, bạn có thể dùng ripgrep để nhanh chóng xác định mô-đun chứa các ký hiệu hoặc phép tính cụ thể.

Lưu ý: Đưa 3 tệp được kết xuất trước/sau/chỉ định vùng đệm mà bạn quan tâm vào báo cáo lỗi.

Lượt chuyển đổi HLO

Một công cụ có tên là hlo-opt có thể dịch giữa HLOProto và các định dạng văn bản. Điều này hữu ích trong trường hợp bạn có một định dạng nhưng cần định dạng khác để gỡ lỗi.

Tìm hiểu cách sử dụng công cụ này: Tài liệu về công cụ XLA: hlo-opt.

Chơi lại

Bạn có thể chạy (phát lại) các phép tính đã kết xuất trên một chương trình phụ trợ XLA cụ thể bằng dữ liệu giả hoặc ảnh chụp nhanh đầu vào. Đây là một cách thuận tiện để tái tạo, lặp lại và gỡ lỗi các vấn đề trong XLA.

Các lệnh sau đây sử dụng dữ liệu giả. Nếu đã lưu HLO Snapshot, bạn có thể truyền các Snapshot đó vào thay thế và dữ liệu từ Snapshot sẽ được dùng. Để vẫn sử dụng dữ liệu giả trong khi chạy ảnh chụp nhanh, hãy truyền cờ --force_fake_data.

Phần phụ trợ CPU:

bazel run -c opt //xla/hlo/tools:run_hlo_module -- --platform=cpu
 /tmp/xladump/module_4561.before_optimizations.txt

Phần phụ trợ GPU:

bazel run -c opt //xla/hlo/tools:run_hlo_module -- --platform=CUDA
 /tmp/xladump/module_4561.before_optimizations.txt

Kết xuất các phép tính HLO Sử dụng bộ sưu tập để sắp xếp ngăn nắp các trang Lưu và phân loại nội dung dựa trên lựa chọn ưu tiên của bạn.

Thực thi cục bộ

Sử dụng biến môi trường

Kết xuất các lượt truyền trung gian cụ thể

Các lựa chọn dành riêng cho JAX

Theo phương thức lập trình trong JAX

Kết xuất jaxpr

Google Colab

Biến môi trường

Các lựa chọn dành riêng cho JAX

Kết xuất tất cả/một số phép tính

Mosaic

More with HLO Dumps

Tìm được phép tính phù hợp

Lượt chuyển đổi HLO

Chơi lại

Kết xuất các phép tính HLO