Kompilacja ze źródła

Z tego dokumentu dowiesz się, jak tworzyć komponenty XLA.

Jeśli nie sklonujesz repozytorium XLA ani nie zainstalujesz Bazela, zapoznaj się z początkowymi sekcjami przewodnika dla programistów XLA.

Linux

Skonfiguruj

Kompilacje XLA są konfigurowane przez plik .bazelrc w katalogu głównym repozytorium. Skrypt ./configure.py może służyć do dostosowywania typowych ustawień.

Jeśli chcesz zmienić konfigurację, uruchom skrypt ./configure.py z katalogu głównego repozytorium. Ten skrypt zawiera flagi lokalizacji zależności XLA i dodatkowe opcje konfiguracji kompilacji (np. flagi kompilatora). Więcej informacji znajdziesz w sekcji Przykładowa sesja.

Obsługa procesora

Do kompilowania i testowania XLA zalecamy używanie odpowiedniego obrazu Dockera, takiego jak ml-build, który jest też używany w przepływach pracy CI XLA w GitHubie. Obraz ml-build ma fabrycznie zainstalowany kompilator 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

Za pomocą kontenera Dockera możesz utworzyć XLA z obsługą procesora, uruchamiając te polecenia:

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

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

Jeśli chcesz tworzyć platformy XLA z obsługą procesora bez użycia Dockera, musisz zainstalować Clang. XLA jest obecnie kompilowany w CI za pomocą Clang 18, ale starsze wersje też powinny działać.

Aby skonfigurować i skompilować środowiska docelowe, uruchom te polecenia:

./configure.py --backend=CPU

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

Obsługa GPU

Do utworzenia XLA z obsługą GPU zalecamy użycie tego samego kontenera Dockera, o którym wspomnieliśmy powyżej.

Aby uruchomić kontener Dockera z dostępem do wszystkich procesorów GPU, uruchom to polecenie:

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

Aby skompilować XLA z obsługą procesora GPU, uruchom te polecenia:

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

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

Nie używaj flagi --gpus all podczas uruchamiania kontenera Dockera.
Ręcznie określ możliwości obliczeniowe CUDA. Przykład:

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

Więcej informacji o obrazach Dockera TensorFlow na GPU znajdziesz w tym dokumencie.

Możesz też tworzyć cele XLA z obsługą GPU bez użycia Dockera. Skonfiguruj i utwórz cele za pomocą tych poleceń:

./configure.py --backend=CUDA

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

Więcej informacji o hermetycznym CUDA znajdziesz w tym dokumencie.

Kompilowanie XLA z obsługą CUDA/cuDNN za pomocą kontenera CI/Release JAX

XLA to kompilator używany wewnętrznie przez JAX. JAX jest rozpowszechniany za pomocą pakietów PyPI. W dokumentacji dotyczącej ciągłej integracji JAX znajdziesz informacje o tym, jak tworzyć pakiety JAX za pomocą kontenera Dockera tensorflow/ml-build:latest.

Możemy też rozszerzyć te instrukcje, aby tworzyć cele XLA w kontenerze JAX. Dzięki temu konfiguracja kompilacji docelowych XLA jest zgodna z konfiguracją kompilacji JAX/XLA, co może być przydatne, jeśli chcemy odtworzyć wyniki obciążenia za pomocą narzędzi XLA, które zostały pierwotnie utworzone w JAX.

Kompilowanie celów XLA w kontenerze CI JAX

Sklonuj repozytorium JAX i przejdź do katalogu „jax”:

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

cd jax

Uruchom kontener Dockera CI/Release JAX, wpisując:

./ci/utilities/run_docker_container.sh

Spowoduje to uruchomienie kontenera Dockera o nazwie „jax”.

Zbuduj element docelowy jax-cuda-plugin w kontenerze za pomocą tego polecenia:

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

Spowoduje to utworzenie pliku .jax_configure.bazelrc z wymaganą konfiguracją kompilacji, w tym obsługą CUDA/cuDNN.

Uzyskaj dostęp do interaktywnej powłoki w kontenerze:

docker exec -ti jax /bin/bash

Powinien teraz być widoczny katalog /jax w kontenerze.

Skompiluj środowisko docelowe XLA za pomocą tego polecenia, np.:

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

Opcjonalnie możesz zastąpić zmienne środowiskowe HERMETIC, np.:

--repo_env=HERMETIC_CUDA_COMPUTE_CAPABILITIES="sm_90"

W razie potrzeby skopiuj powstałe artefakty do /jax/dist, aby uzyskać do nich dostęp z systemu operacyjnego hosta.

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

Wyjdź z interaktywnej powłoki:

exit