Tạo từ nguồn

Tài liệu này mô tả cách tạo các thành phần XLA.

Nếu bạn chưa sao chép kho lưu trữ XLA hoặc cài đặt Bazel, hãy xem các phần ban đầu của Hướng dẫn dành cho nhà phát triển XLA.

Linux

Định cấu hình

Các bản dựng XLA được định cấu hình bằng tệp .bazelrc trong thư mục gốc của kho lưu trữ. Bạn có thể dùng tập lệnh ./configure.py để điều chỉnh các chế độ cài đặt thường dùng.

Nếu bạn cần thay đổi cấu hình, hãy chạy tập lệnh ./configure.py từ thư mục gốc của kho lưu trữ. Tập lệnh này có các cờ cho vị trí của các phần phụ thuộc XLA và các lựa chọn cấu hình bản dựng bổ sung (ví dụ: cờ trình biên dịch). Hãy tham khảo phần Phiên mẫu để biết thông tin chi tiết.

Hỗ trợ CPU

Bạn nên dùng một hình ảnh Docker phù hợp (chẳng hạn như ml-build, cũng được dùng trong quy trình CI của XLA trên GitHub) để tạo và kiểm thử XLA. Hình ảnh ml-build được cài đặt sẵn Clang 18.

docker run -itd --rm \
--name xla \
-w /xla \
-v $PWD:/xla \
us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest \
bash

Bằng cách sử dụng một vùng chứa Docker, bạn có thể tạo XLA có hỗ trợ CPU bằng cách chạy các lệnh sau:

docker exec xla ./configure.py --backend=CPU

docker exec xla bazel build \
  --spawn_strategy=sandboxed \
  --test_output=all \
  //xla/...

Nếu muốn tạo các mục tiêu XLA có hỗ trợ CPU mà không sử dụng Docker, bạn cần cài đặt Clang. XLA hiện được tạo bằng Clang 18 trong CI, nhưng các phiên bản trước đó cũng sẽ hoạt động.

Để định cấu hình và tạo các mục tiêu, hãy chạy các lệnh sau:

./configure.py --backend=CPU

bazel build \
  --spawn_strategy=sandboxed \
  --test_output=all \
  //xla/...

Hỗ trợ GPU

Bạn nên sử dụng cùng một vùng chứa Docker như đã đề cập ở trên để tạo XLA có hỗ trợ GPU.

Để khởi động vùng chứa Docker có quyền truy cập vào tất cả GPU, hãy chạy lệnh sau:

docker run -itd --rm \
  --gpus all \
  --name xla_gpu \
  -w /xla \
  -v $PWD:/xla \
  us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest \
  bash

Để tạo XLA có hỗ trợ GPU, hãy chạy các lệnh sau:

docker exec xla_gpu ./configure.py --backend=CUDA

docker exec xla_gpu bazel build \
  --spawn_strategy=sandboxed \
  --test_output=all \
  //xla/...

Không sử dụng cờ --gpus all khi khởi động vùng chứa Docker.
Chỉ định khả năng tính toán CUDA theo cách thủ công, ví dụ:

docker exec xla_gpu ./configure.py --backend=CUDA \
  --cuda_compute_capabilities="9.0"

Để biết thêm thông tin chi tiết về các hình ảnh docker GPU của TensorFlow, bạn có thể xem tài liệu này.

Bạn cũng có thể tạo các mục tiêu XLA có hỗ trợ GPU mà không cần Docker. Định cấu hình và tạo các mục tiêu bằng các lệnh sau:

./configure.py --backend=CUDA

bazel build \
  --spawn_strategy=sandboxed \
  --test_output=all \
  //xla/...

Để biết thêm thông tin chi tiết về CUDA khép kín, bạn có thể xem tài liệu này.

Tạo XLA có hỗ trợ CUDA/cuDNN bằng vùng chứa JAX CI/Release

XLA là một trình biên dịch được JAX sử dụng nội bộ. JAX được phân phối thông qua các bánh xe PyPI. Tài liệu về tính năng Tích hợp liên tục của JAX giải thích cách tạo các bánh xe JAX bằng vùng chứa Docker tensorflow/ml-build:latest.

Chúng ta cũng có thể mở rộng các hướng dẫn này để tạo mục tiêu XLA trong vùng chứa JAX. Điều này đảm bảo rằng cấu hình bản dựng của các mục tiêu XLA nhất quán với cấu hình bản dựng JAX/XLA. Điều này có thể hữu ích nếu chúng ta muốn tái tạo kết quả của khối lượng công việc bằng cách sử dụng các công cụ XLA được tạo ban đầu trong JAX.

Tạo mục tiêu XLA trong vùng chứa CI JAX

Sao chép kho lưu trữ JAX và chuyển đến thư mục "jax"

git clone <a href="https://github.com/jax-ml/jax.git">https://github.com/jax-ml/jax.git</a>

cd jax

Khởi động vùng chứa Docker JAX CI/Release bằng cách chạy:

./ci/utilities/run_docker_container.sh

Thao tác này sẽ khởi động một vùng chứa Docker có tên là "jax".

Tạo đích jax-cuda-plugin bên trong vùng chứa bằng cách sử dụng:

docker exec jax ./ci/build_artifacts.sh jax-cuda-plugin

Thao tác này sẽ tạo tệp .jax_configure.bazelrc có cấu hình bản dựng bắt buộc, bao gồm cả hỗ trợ CUDA/cuDNN

Truy cập vào một shell tương tác bên trong vùng chứa:

docker exec -ti jax /bin/bash

Giờ đây, bạn sẽ ở trong thư mục /jax trong vùng chứa

Tạo mục tiêu XLA bằng lệnh sau, ví dụ:

/usr/local/bin/bazel build \
  --config=cuda_libraries_from_stubs \
  --verbose_failures=true \
  @xla//xla/tools/multihost_hlo_runner:hlo_runner_main

Bạn có thể ghi đè các env HERMETIC (không bắt buộc), ví dụ:

--repo_env=HERMETIC_CUDA_COMPUTE_CAPABILITIES="sm_90"

Sao chép các cấu phần phần mềm thu được vào /jax/dist để truy cập vào chúng từ hệ điều hành máy chủ nếu cần

cp bazel-bin/external/xla/xla/tools/multihost_hlo_runner/hlo_runner_main \
  ./dist/

Thoát khỏi màn hình shell tương tác:

exit