Bring back PyTorch/XLA GPU tests/builds

[tengyifei]

🐛 Bug

PyTorch/XLA on GPUs builds have been failing since Oct 21, 2024.

In order to bring back GPU builds and tests, the first challenge is to build PyTorch/XLA with clang and hermetic CUDA 1. After that, there may be more tests to fix.

This bug tracks work/discussions needed to bring back GPU builds.

[tengyifei]

Chengji has a wip branch for hermetic CUDA: https://github.com/yaochengji/xla/tree/chengji/clang-herm

[tengyifei]

cc @ysiraichi

[miladm]

cc @tengyifei to work with @ysiraichi. Ideally, we would like to use this bug to bring back GPU whl for 2.6 release.

[ysiraichi]

Update: I have tried Chengji's branch, but the build kept failing with:

failed: undeclared inclusion(s) in rule '@zlib//:zlib':
this rule is missing dependency declarations for the following files included by 'zutil.c':
  '/usr/lib/clang/17/include/stddef.h'
  '/usr/lib/clang/17/include/__stddef_max_align_t.h'
  '/usr/lib/clang/17/include/limits.h'
  '/usr/lib/clang/17/include/stdarg.h'

Still investigating it.

[ysiraichi]

Here's a more verbose update on transitioning PyTorch/XLA to use OpenXLA hermetic CUDA. In summary, these are the things I have added to the build system (branch diff):

• Followed the documentation, and added the necessary code to WORKSPACE and .bazelrc
• Added a few more things that I could find in the JAX repo
• Explicitly set clang-17 as CC and CXX

Even after all these steps, I am still hitting the error above.

Reproducing the Error

• Get my hermetic-cuda branch
• Go to the following directory: plugins/cuda
• Run the following command:

$ bazel build @xla//xla/pjrt/c:pjrt_c_api_gpu_plugin.so --symlink_prefix=$(pwd)/bazel- --config=cuda

My Thoughts

• As far as I understand it, we wouldn't need to specify those dependencies, since they are system dependencies (from clang itself)
• These errors are occurring outside PyTorch/XLA (I believe they are OpenXLA dependencies)
• I believe I am building an OpenXLA target similarly to how JAX builds it. I just don't know what I'm doing differently...

Question

How to fix this error?

[miladm]

Thanks @ysiraichi - I've pined openxla partners to share their input on this issue.

[beckerhe]

Update: I have tried Chengji's branch, but the build kept failing with:

failed: undeclared inclusion(s) in rule '@zlib//:zlib':
this rule is missing dependency declarations for the following files included by 'zutil.c':
  '/usr/lib/clang/17/include/stddef.h'
  '/usr/lib/clang/17/include/__stddef_max_align_t.h'
  '/usr/lib/clang/17/include/limits.h'
  '/usr/lib/clang/17/include/stdarg.h'

Still investigating it.

Can you try setting CC and CXX to the full absolute path of clang?

CC=/usr/lib/llvm-17/bin/clang

[ybaturina]

The fix openxla/xla#22165 is merged. Can you try building Pytorch again please? (without changing Clang symlink to absolute path)

[ysiraichi]

Thanks, @ybaturina. Using the absolute path did work! I still haven't tested your patch.

[ysiraichi]

@tengyifei @will-cromar What's the recommended process for installing something (e.g. clang-17) inside the current dev VM that runs on CI? I saw that the install_deps role is being skipped by default.

[tengyifei]

@ysiraichi I believe you need to install clang-17 into the development image by sending a PR to

xla/infra/ansible/development.Dockerfile

Line 15 in 3578940

ENV TAGS="bazel,configure_env,install_deps"

.

That builds a dev docker image that will be accessible at https://console.cloud.google.com/artifacts/docker/tpu-pytorch-releases/us-central1/docker/development