В этом документе описывается, как создавать компоненты XLA.
Если вы не клонировали репозиторий XLA или не установили Bazel, ознакомьтесь с начальными разделами Руководства разработчика XLA .
Linux
Настройка
Настройка сборок XLA осуществляется с помощью файла .bazelrc , расположенного в корневом каталоге репозитория. Для изменения общих параметров можно использовать скрипт ./configure.py .
Если вам необходимо изменить конфигурацию, запустите скрипт ./configure.py из корневого каталога репозитория. Этот скрипт содержит флаги для указания расположения зависимостей XLA и дополнительные параметры конфигурации сборки (например, флаги компилятора). Подробности см. в разделе «Пример сессии» .
поддержка ЦП
Для сборки и тестирования XLA мы рекомендуем использовать подходящий образ Docker, например ml-build , который также используется в рабочих процессах CI XLA на GitHub. Образ 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, вы можете собрать XLA с поддержкой ЦП, выполнив следующие команды:
docker exec xla ./configure.py --backend=CPU
docker exec xla bazel build \
--spawn_strategy=sandboxed \
--test_output=all \
//xla/...
Если вы хотите собрать целевые объекты XLA с поддержкой ЦП без использования Docker , вам потребуется установить Clang. В настоящее время XLA собирается с помощью Clang 18 в CI, но более ранние версии также должны работать.
Для настройки и сборки целевых объектов выполните следующие команды:
./configure.py --backend=CPU
bazel build \
--spawn_strategy=sandboxed \
--test_output=all \
//xla/...
поддержка графического процессора
Мы рекомендуем использовать тот же контейнер Docker, что и выше, для сборки XLA с поддержкой графического процессора.
Чтобы запустить контейнер Docker с доступом ко всем графическим процессорам, выполните следующую команду:
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
Для сборки XLA с поддержкой графического процессора выполните следующие команды:
docker exec xla_gpu ./configure.py --backend=CUDA
docker exec xla_gpu bazel build \
--spawn_strategy=sandboxed \
--test_output=all \
//xla/...
- Не используйте флаг
--gpus allпри запуске контейнера Docker. - Укажите вычислительные возможности CUDA вручную, например:
docker exec xla_gpu ./configure.py --backend=CUDA \
--cuda_compute_capabilities="9.0"
Более подробную информацию об образах Docker для работы с GPU в TensorFlow можно найти в этом документе.
Вы можете создавать целевые объекты XLA с поддержкой GPU и без Docker. Настройте и соберите целевые объекты, используя следующие команды:
./configure.py --backend=CUDA
bazel build \
--spawn_strategy=sandboxed \
--test_output=all \
//xla/...
Для получения более подробной информации о герметичной архитектуре CUDA вы можете ознакомиться с этим документом.
Сборка XLA с поддержкой CUDA/cuDNN с использованием контейнера JAX CI/Release
XLA — это компилятор, используемый внутри JAX. JAX распространяется через пакеты PyPI. В документации по непрерывной интеграции JAX объясняется, как собирать пакеты JAX с помощью контейнера Docker tensorflow/ml-build:latest .
Мы можем расширить эти инструкции, чтобы собирать цели XLA и внутри контейнера JAX. Это гарантирует, что конфигурация сборки целей XLA будет соответствовать конфигурации сборки JAX/XLA, что может быть полезно, если мы хотим воспроизвести результаты рабочей нагрузки с использованием инструментов XLA, которые изначально были созданы в JAX.
Создание целей XLA в контейнере JAX CI
- Клонируйте репозиторий JAX и перейдите в каталог 'jax'.
git clone https://github.com/jax-ml/jax.git
cd jax
- Запустите контейнер Docker JAX CI/Release, выполнив команду:
./ci/utilities/run_docker_container.sh
Это запустит контейнер Docker с именем 'jax'.
- Создайте целевой объект 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 среды, например:
--repo_env=HERMETIC_CUDA_COMPUTE_CAPABILITIES="sm_90"
- Скопируйте полученные артефакты в
/jax/dist, чтобы при необходимости получить к ним доступ из операционной системы хоста.
cp bazel-bin/external/xla/xla/tools/multihost_hlo_runner/hlo_runner_main \
./dist/
- Выйти из интерактивной оболочки:
exit
Windows
Сборка XLA в Windows нативно — это процесс, задействующий только процессор, поскольку CUDA напрямую не поддерживается в Windows; для поддержки CUDA необходимо использовать WSL2. Также требуются определенные настройки среды, включая оболочку Bash, компилятор Clang и Visual Studio.
Предварительные требования
- Visual Studio: Для настройки набора инструментов C++ (который предоставляет необходимые системные заголовочные файлы и библиотеки) необходимо установить Visual Studio 2019 версии 16.5 или новее.
- Среда Bash: Для сборки XLA в Windows необходимо использовать оболочку Bash (например, MSYS2 или Git Bash).
- Компилятор Clang: В сборках XLA для Windows используется
clang-clвместо стандартного компилятора MSVC. Убедитесь, что LLVM/Clang установлен. - Python: Python 3 должен быть установлен и доступен в
PATHвашей системы. - Режим разработчика: Необходимо включить режим разработчика в Windows или запустить оболочку Bash от имени администратора. Это необходимо, поскольку Bazel использует создание символических ссылок для дерева исполняемых файлов в процессе сборки.
Создание с нуля
Для компиляции бэкенда XLA для ЦП в Windows выполните следующую команду в терминале Bash. Мы используем конфигурацию xla_windows_x86_cpu_2022 для автоматической настройки инструментария clang-cl .
Обратите внимание, что некоторые целевые платформы и теги необходимо явно исключить из сборки Windows для ЦП, например, целевые платформы для графических процессоров и функции, которые еще не поддерживаются в Windows:
# Note: Ensure your PATH includes Python and Bazel before running this
bazel build \
--config=xla_windows_x86_cpu_2022 \
--keep_going \
--build_tag_filters=-no_oss,-oss_excluded,-gpu,-no_windows,-windows_excluded \
-- //xla/... -//xla/hlo/experimental/... -//xla/python_api/... -//xla/python/...