本文說明如何建構 XLA 元件。
如果您未複製 XLA 存放區或安裝 Bazel,請參閱 XLA 開發人員指南的初始章節。
Linux
設定
XLA 建構作業是透過存放區根目錄中的 .bazelrc 檔案設定。您可以使用 ./configure.py 指令碼調整常見設定。
如要變更設定,請從存放區的根目錄執行 ./configure.py 指令碼。這個指令碼包含 XLA 依附元件位置的旗標,以及其他建構設定選項 (例如編譯器旗標)。詳情請參閱「範例工作階段」一節。
CPU 支援
建議使用合適的 Docker 映像檔 (例如 ml-build),這類映像檔也用於 GitHub 上的 XLA CI 工作流程,可建構及測試 XLA。ml-build 映像檔已預先安裝 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
使用 Docker 容器時,您可以執行下列指令,建構支援 CPU 的 XLA:
docker exec xla ./configure.py --backend=CPU
docker exec xla bazel build \
--spawn_strategy=sandboxed \
--test_output=all \
//xla/...
如要不使用 Docker 建構支援 CPU 的 XLA 目標,您需要安裝 Clang。XLA 目前在 CI 中以 Clang 18 建構,但舊版也應可運作。
如要設定及建構目標,請執行下列指令:
./configure.py --backend=CPU
bazel build \
--spawn_strategy=sandboxed \
--test_output=all \
//xla/...
GPU 支援
建議您使用上述 Docker 容器,建構支援 GPU 的 XLA。
如要啟動 Docker 容器並存取所有 GPU,請執行下列指令:
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
如要建構支援 GPU 的 XLA,請執行下列指令:
docker exec xla_gpu ./configure.py --backend=CUDA
docker exec xla_gpu bazel build \
--spawn_strategy=sandboxed \
--test_output=all \
//xla/...
- 啟動 Docker 容器時,請不要使用
--gpus all旗標。 - 手動指定 CUDA 運算功能,例如:
docker exec xla_gpu ./configure.py --backend=CUDA \
--cuda_compute_capabilities="9.0"
如要進一步瞭解 TensorFlow 的 GPU Docker 映像檔,請參閱這份文件。
您也可以在不使用 Docker 的情況下,建構支援 GPU 的 XLA 目標。使用下列指令設定及建構目標:
./configure.py --backend=CUDA
bazel build \
--spawn_strategy=sandboxed \
--test_output=all \
//xla/...
如要進一步瞭解密封式 CUDA,請參閱這份文件。
使用 JAX CI/發布容器建構支援 CUDA/cuDNN 的 XLA
XLA 是 JAX 內部使用的編譯器。JAX 是透過 PyPI Wheel 發布。 JAX 持續整合說明文件說明如何使用 tensorflow/ml-build:latest Docker 容器建構 JAX 輪子。
我們也可以擴充這些指令,在 JAX 容器中建構 XLA 目標。這可確保 XLA 目標的建構設定與 JAX/XLA 建構設定一致。如果我們想使用最初在 JAX 中建立的 XLA 工具重現工作負載結果,這就非常實用。
在 JAX CI 容器中建構 XLA 目標
- 複製 JAX 存放區,然後前往「jax」目錄
git clone <a href="https://github.com/jax-ml/jax.git">https://github.com/jax-ml/jax.git</a>
cd jax
- 執行下列指令,啟動 JAX CI/Release Docker 容器:
./ci/utilities/run_docker_container.sh
這會啟動名為「jax」的 Docker 容器。
- 在容器內使用下列指令建構 jax-cuda-plugin 目標:
docker exec jax ./ci/build_artifacts.sh jax-cuda-plugin
這會建立 .jax_configure.bazelrc 檔案,其中包含必要的建構設定,包括 CUDA/cuDNN 支援
- 存取容器內的互動式殼層:
docker exec -ti jax /bin/bash
您現在應該位於容器內的 /jax 目錄
- 使用下列指令建構 XLA 目標,例如:
/usr/local/bin/bazel build \
--config=cuda_libraries_from_stubs \
--verbose_failures=true \
@xla//xla/tools/multihost_hlo_runner:hlo_runner_main
您也可以視需要覆寫 HERMETIC envs,例如:
--repo_env=HERMETIC_CUDA_COMPUTE_CAPABILITIES="sm_90"
- 將產生的構件複製到
/jax/dist,以便從主機 OS 存取 (如有需要)
cp bazel-bin/external/xla/xla/tools/multihost_hlo_runner/hlo_runner_main \
./dist/
- 結束互動式殼層:
exit