Aus Quelle erstellen

In diesem Dokument wird beschrieben, wie Sie XLA-Komponenten erstellen.

Wenn Sie das XLA-Repository nicht geklont oder Bazel nicht installiert haben, lesen Sie die ersten Abschnitte des XLA Developer Guide.

Linux

Konfigurieren

XLA-Builds werden durch die Datei .bazelrc im Stammverzeichnis des Repositorys konfiguriert. Mit dem ./configure.py-Script können Sie allgemeine Einstellungen anpassen.

Wenn Sie die Konfiguration ändern müssen, führen Sie das Skript ./configure.py aus dem Stammverzeichnis des Repositorys aus. Dieses Skript enthält Flags für den Speicherort von XLA-Abhängigkeiten und zusätzliche Build-Konfigurationsoptionen (z. B. Compiler-Flags). Weitere Informationen finden Sie im Abschnitt Beispielsitzung.

CPU-Unterstützung

Wir empfehlen, ein geeignetes Docker-Image wie ml-build zu verwenden, das auch in den CI-Workflows von XLA auf GitHub verwendet wird, um XLA zu erstellen und zu testen. Im ml-build-Image ist Clang 18 vorinstalliert.

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

Mit einem Docker-Container können Sie XLA mit CPU-Unterstützung erstellen, indem Sie die folgenden Befehle ausführen:

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

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

Wenn Sie XLA-Ziele mit CPU-Unterstützung ohne Docker erstellen möchten, müssen Sie Clang installieren. XLA wird derzeit mit Clang 18 in CI erstellt, aber frühere Versionen sollten auch funktionieren.

Führen Sie die folgenden Befehle aus, um die Ziele zu konfigurieren und zu erstellen:

./configure.py --backend=CPU

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

GPU-Unterstützung

Wir empfehlen, denselben Docker-Container wie oben zu verwenden, um XLA mit GPU-Unterstützung zu erstellen.

Führen Sie den folgenden Befehl aus, um den Docker-Container mit Zugriff auf alle GPUs zu starten:

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

Führen Sie die folgenden Befehle aus, um XLA mit GPU-Unterstützung zu erstellen:

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

docker exec xla_gpu bazel build \
  --spawn_strategy=sandboxed \
  --test_output=all \
  //xla/...
  • Verwenden Sie das Flag --gpus all nicht, wenn Sie den Docker-Container starten.
  • Geben Sie die CUDA-Rechenfunktionen manuell an, z. B.:
docker exec xla_gpu ./configure.py --backend=CUDA \
  --cuda_compute_capabilities="9.0"

Weitere Informationen zu den GPU-Docker-Images von TensorFlow finden Sie in diesem Dokument.

Sie können XLA-Ziele mit GPU-Unterstützung auch ohne Docker erstellen. Konfigurieren und erstellen Sie Ziele mit den folgenden Befehlen:

./configure.py --backend=CUDA

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

Weitere Informationen zu hermetischem CUDA

XLA mit CUDA-/cuDNN-Unterstützung mit dem JAX CI-/Release-Container erstellen

XLA ist ein Compiler, der intern von JAX verwendet wird. JAX wird über PyPI-Wheels verteilt. In der Dokumentation zur kontinuierlichen Integration von JAX wird beschrieben, wie Sie JAX-Wheels mit dem Docker-Container tensorflow/ml-build:latest erstellen.

Wir können diese Anleitung auch auf die Erstellung von XLA-Zielen im JAX-Container ausweiten. So wird sichergestellt, dass die Build-Konfiguration der XLA-Ziele mit der Build-Konfiguration von JAX/XLA übereinstimmt. Das kann nützlich sein, wenn wir Arbeitslast-Ergebnisse mit XLA-Tools reproduzieren möchten, die ursprünglich in JAX erstellt wurden.

XLA-Ziele im JAX-CI-Container erstellen

  1. Klonen Sie das JAX-Repository und wechseln Sie in das Verzeichnis „jax“.
git clone <a href="https://github.com/jax-ml/jax.git">https://github.com/jax-ml/jax.git</a>

cd jax
  1. Starten Sie den JAX CI/Release-Docker-Container mit folgendem Befehl:
./ci/utilities/run_docker_container.sh

Dadurch wird ein Docker-Container mit dem Namen „jax“ gestartet.

  1. Erstellen Sie das Ziel „jax-cuda-plugin“ im Container mit folgendem Befehl:
docker exec jax ./ci/build_artifacts.sh jax-cuda-plugin

Dadurch wird die Datei „.jax_configure.bazelrc“ mit der erforderlichen Build-Konfiguration erstellt, einschließlich CUDA-/cuDNN-Unterstützung.

  1. Greifen Sie auf eine interaktive Shell im Container zu:
docker exec -ti jax /bin/bash

Sie sollten sich jetzt im Verzeichnis /jax im Container befinden.

  1. Erstellen Sie das XLA-Ziel mit dem folgenden Befehl, z.B.:
/usr/local/bin/bazel build \
  --config=cuda_libraries_from_stubs \
  --verbose_failures=true \
  @xla//xla/tools/multihost_hlo_runner:hlo_runner_main

Optional können Sie HERMETIC-Umgebungen überschreiben, z.B.:

--repo_env=HERMETIC_CUDA_COMPUTE_CAPABILITIES="sm_90"
  1. Kopieren Sie die resultierenden Artefakte in /jax/dist, um bei Bedarf über das Hostbetriebssystem darauf zuzugreifen.
cp bazel-bin/external/xla/xla/tools/multihost_hlo_runner/hlo_runner_main \
  ./dist/
  1. Beenden Sie die interaktive Shell:
exit