Update XLA pin then fix up to make it compile (#9565)

qihqi · web-flow · commit 748ac9b1032c · 2025-08-22T20:53:36.000Z
Updates:

* Modify call to some OpenXLA C++ functions because their calling
convention changed.
* bazel is now more strict: explicit header and explicit deps are
enforced
* removed FORTIFY_SOURCE defines via a patch, it is making gcc segfault
on compile
* updated gcc version to 11
diff --git a/.circleci/common.sh b/.circleci/common.sh
@@ -112,7 +112,8 @@ function build_torch_xla() {
   # Need to uncomment the line below.
   # Currently it fails upstream XLA CI.
   # pip install plugins/cuda -v
-  pip install 'torch_xla[pallas]'
+  pip install --pre torch_xla[pallas] --index-url https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/ --find-links https://storage.googleapis.com/jax-releases/libtpu_releases.html
+
   popd
 }
 
diff --git a/.github/workflows/_test.yml b/.github/workflows/_test.yml
@@ -140,7 +140,7 @@ jobs:
           set -x
 
           pip install expecttest unittest-xml-reporting
-          pip install 'torch_xla[pallas]'
+          pip install --pre 'torch_xla[pallas]' --index-url https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/ --find-links https://storage.googleapis.com/jax-releases/libtpu_releases.html
 
           if [[ ! -z "$RUN_BENCHMARK_TESTS" ]]; then
             pip install -r pytorch/xla/benchmarks/requirements.txt
diff --git a/.github/workflows/_tpu_ci.yml b/.github/workflows/_tpu_ci.yml
@@ -52,8 +52,8 @@ jobs:
           pip install fsspec
           pip install rich
           # jax and libtpu is needed for pallas tests.
-          pip install 'torch_xla[pallas]'
-          pip install 'torch_xla[tpu]' -f https://storage.googleapis.com/libtpu-wheels/index.html -f https://storage.googleapis.com/libtpu-releases/index.html
+          pip install --pre 'torch_xla[pallas]' --index-url https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/ --find-links https://storage.googleapis.com/jax-releases/libtpu_releases.html'
+          pip install --pre 'torch_xla[tpu]' --index-url https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/ --find-links https://storage.googleapis.com/jax-releases/libtpu_releases.html'
           pip install --upgrade protobuf
       - name: Run Tests (${{ matrix.test_script }})
         if: inputs.has_code_changes == 'true'
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
@@ -162,7 +162,7 @@ commands on your Linux machine directly, outside of the container.
       -f https://storage.googleapis.com/libtpu-releases/index.html
 
     # Optional: if you're using custom kernels, install pallas dependencies
-    pip install torch_xla[pallas]
+    pip install --pre torch_xla[pallas] --index-url https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/ --find-links https://storage.googleapis.com/jax-releases/libtpu_releases.html
     ```
 
 1.  If you are running on a TPU VM, ensure `torch` and `torch_xla` were built and
diff --git a/README.md b/README.md
@@ -35,7 +35,7 @@ Note: Builds are available for Python 3.8 to 3.11; please use one of the support
 pip install torch==2.8.0 'torch_xla[tpu]==2.8.0'
 
 # Optional: if you're using custom kernels, install pallas dependencies
-pip install 'torch_xla[pallas]'
+pip install --pre torch_xla[pallas] --index-url https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/ --find-links https://storage.googleapis.com/jax-releases/libtpu_releases.html
 ```
 **As of 07/16/2025 and starting from Pytorch/XLA 2.8 release, PyTorch/XLA will 
 provide nightly and release wheels for Python 3.11 to 3.13**
diff --git a/WORKSPACE b/WORKSPACE
@@ -46,7 +46,7 @@ new_local_repository(
 
 # To build PyTorch/XLA with a new revison of OpenXLA, update the xla_hash to
 # the openxla git commit hash and note the date of the commit.
-xla_hash = '3d5ece64321630dade7ff733ae1353fc3c83d9cc'  # Committed on 2025-06-17.
+xla_hash = '92f7b5952dd585c5be17c9a5caad27407005b513'  # Committed on 2025-08-15.
 
 http_archive(
     name = "xla",
@@ -58,7 +58,7 @@ http_archive(
     patches = [
         "//openxla_patches:gpu_nvml.diff",
         "//openxla_patches:gpu_race_condition.diff",
-        "//openxla_patches:count_down.diff",
+        "//openxla_patches:no_fortify.diff",
     ],
     strip_prefix = "xla-" + xla_hash,
     urls = [
@@ -81,6 +81,19 @@ http_archive(
 #    path = "/path/to/openxla",
 # )
 
+# Initialize OpenXLA's external dependencies. There is an specific order
+# which those dependencies are initialized, because for bazel it's the
+# first definition that takes precedence.
+# We follow what openxla/xla does exactly: 
+# https://github.com/openxla/xla/blob/main/WORKSPACE#L37
+load("@xla//:workspace4.bzl", "xla_workspace4")
+
+xla_workspace4()
+
+load("@xla//:workspace3.bzl", "xla_workspace3")
+
+xla_workspace3()
+
 # Initialize hermetic Python
 load("@xla//third_party/py:python_init_rules.bzl", "python_init_rules")
 
@@ -115,14 +128,6 @@ install_deps()
 
 
 
-# Initialize OpenXLA's external dependencies.
-load("@xla//:workspace4.bzl", "xla_workspace4")
-
-xla_workspace4()
-
-load("@xla//:workspace3.bzl", "xla_workspace3")
-
-xla_workspace3()
 
 load("@xla//:workspace2.bzl", "xla_workspace2")
 
diff --git a/bazel/rules_def.bzl b/bazel/rules_def.bzl
@@ -1,10 +1,4 @@
 """Rules that simplify deps and compiler configuration for PyTorch/XLA."""
-
-load(
-    "@xla//xla:xla.default.bzl",
-    "xla_cc_test",
-)
-
 def ptxla_cc_library(
         deps = [],
         copts = [],
@@ -22,7 +16,7 @@ def ptxla_cc_test(
         deps,
         copts = [],
         **kwargs):
-    xla_cc_test(
+    native.cc_test(
         linkstatic = True,
         copts = copts + [
             "-isystemexternal/torch",  # Required for system includes.
diff --git a/openxla_patches/count_down.diff b/openxla_patches/count_down.diff
diff --git a/openxla_patches/no_fortify.diff b/openxla_patches/no_fortify.diff
@@ -0,0 +1,40 @@
+diff --git a/tools/toolchains/cross_compile/cc/BUILD b/tools/toolchains/cross_compile/cc/BUILD
+index be6e8968d5..51d15fa3c9 100644
+--- a/tools/toolchains/cross_compile/cc/BUILD
++++ b/tools/toolchains/cross_compile/cc/BUILD
+@@ -80,7 +80,6 @@ cc_toolchain_config(
+     opt_compile_flags = [
+         "-g0",
+         "-O2",
+-        "-D_FORTIFY_SOURCE=1",
+         "-DNDEBUG",
+         "-ffunction-sections",
+         "-fdata-sections",
+@@ -166,7 +165,6 @@ cc_toolchain_config(
+     opt_compile_flags = [
+         "-g0",
+         "-O2",
+-        "-D_FORTIFY_SOURCE=1",
+         "-DNDEBUG",
+         "-ffunction-sections",
+         "-fdata-sections",
+@@ -263,7 +261,6 @@ cc_toolchain_config(
+     opt_compile_flags = [
+         "-g0",
+         "-O2",
+-        "-D_FORTIFY_SOURCE=1",
+         "-DNDEBUG",
+         "-ffunction-sections",
+         "-fdata-sections",
+diff --git a/tools/toolchains/cross_compile/cc/cc_toolchain_config.bzl b/tools/toolchains/cross_compile/cc/cc_toolchain_config.bzl
+index de638c0159..f25c0cc0fe 100644
+--- a/tools/toolchains/cross_compile/cc/cc_toolchain_config.bzl
++++ b/tools/toolchains/cross_compile/cc/cc_toolchain_config.bzl
+@@ -218,7 +218,6 @@ def _impl(ctx):
+                     flag_group(
+                         # Security hardening requires optimization.
+                         # We need to undef it as some distributions now have it enabled by default.
+-                        flags = ["-U_FORTIFY_SOURCE"],
+                     ),
+                 ],
+                 with_features = [
diff --git a/scripts/build_torch_wheels.sh b/scripts/build_torch_wheels.sh
@@ -139,9 +139,9 @@ function install_llvm_clang() {
   sudo update-alternatives --install /usr/bin/clang++ clang++ $(which clang++-8) 70
 }
 
-function install_gcc10() {
-  sudo apt-get -y install gcc-10 g++-10
-  export CC=/usr/bin/gcc-10 export CXX=/usr/bin/g++-10
+function install_gcc() {
+  sudo apt-get -y install gcc-11 g++-11
+  export CC=/usr/bin/gcc-10 export CXX=/usr/bin/g++-11
   sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-10 100
   sudo update-alternatives --install /usr/bin/g++ g++ /usr/bin/g++-10 100
 }
@@ -332,7 +332,7 @@ function main() {
   if [[ $(uname -m) == "x86_64" ]]; then
     install_llvm_clang
   elif [[ $(uname -m) == "aarch64" ]]; then
-    install_gcc10
+    install_gcc
   fi
   install_and_setup_conda
   build_and_install_torch
diff --git a/scripts/update_deps.py b/scripts/update_deps.py
@@ -162,8 +162,8 @@ def get_latest_stable_jax_info() -> tuple[str, str, str] | None:
   published_at = data['published_at']  # e.g., "2024-04-26T22:58:34Z"
   release_date = published_at.split('T')[0]  # e.g., "2024-04-26"
 
-  # The XLA commit is in third_party/xla/workspace.bzl in the JAX repo.
-  workspace_bzl_url = f'https://raw.githubusercontent.com/google/jax/{tag_name}/third_party/xla/workspace.bzl'
+  # The XLA commit is in third_party/xla/revision.bzl in the JAX repo.
+  workspace_bzl_url = f'https://raw.githubusercontent.com/google/jax/{tag_name}/third_party/xla/revision.bzl'
   try:
     with urllib.request.urlopen(workspace_bzl_url) as response:
       workspace_content = response.read().decode()
diff --git a/setup.py b/setup.py
@@ -115,27 +115,27 @@
 
 USE_NIGHTLY = True  # Whether to use nightly or stable libtpu and JAX.
 
-_libtpu_version = '0.0.18'
-_libtpu_date = '20250617'
+_libtpu_version = '0.0.21'
+_libtpu_date = '20250813'
 
-_jax_version = '0.6.2'
-_jaxlib_version = '0.6.2'
-_jax_date = '20250617'  # Date for jax and jaxlib.
+_jax_version = '0.7.1'
+_jaxlib_version = '0.7.1'
+_jax_date = '20250813'  # Date for jax and jaxlib.
 
 if USE_NIGHTLY:
-  _libtpu_version += f".dev{_libtpu_date}"
+  _libtpu_version += f".dev{_libtpu_date}+nightly"
   _jax_version += f'.dev{_jax_date}'
   _jaxlib_version += f'.dev{_jax_date}'
-  _libtpu_wheel_name = f'libtpu-{_libtpu_version}.dev{_libtpu_date}+nightly-py3-none-manylinux_2_31_{platform_machine}'
+  _libtpu_wheel_name = f'libtpu-{_libtpu_version}-py3-none-manylinux_2_31_{platform_machine}'
   _libtpu_storage_directory = 'libtpu-nightly-releases'
 else:
   # The postfix can be changed when the version is updated. Check
   # https://storage.googleapis.com/libtpu-wheels/index.html for correct
   # versioning.
   _libtpu_wheel_name = f'libtpu-{_libtpu_version}-py3-none-manylinux_2_31_{platform_machine}'
   _libtpu_storage_directory = 'libtpu-lts-releases'
-
-_libtpu_storage_path = f'https://storage.googleapis.com/{_libtpu_storage_directory}/wheels/libtpu/{_libtpu_wheel_name}.whl'
+#https://us-python.pkg.dev/ml-oss-artifacts-published/jax/libtpu/libtpu-0.0.19.1-py3-none-manylinux_2_31_x86_64.whl
+_libtpu_storage_path = f'https://us-python.pkg.dev/ml-oss-artifacts-published/jax/libtpu/{_libtpu_wheel_name}.whl'
 
 
 def _get_build_mode():
@@ -423,22 +423,10 @@ def link_packages(self):
 
 
 def _get_jax_install_requirements():
-  if not USE_NIGHTLY:
-    # Stable versions of JAX can be directly installed from PyPI.
-    return [
-        f'jaxlib=={_jaxlib_version}',
-        f'jax=={_jax_version}',
-    ]
-
-  # Install nightly JAX libraries from the JAX package registries.
-  jax = f'jax @ https://us-python.pkg.dev/ml-oss-artifacts-published/jax-public-nightly-artifacts-registry/jax/jax-{_jax_version}-py3-none-any.whl'
-
-  jaxlib = []
-  for python_minor_version in [9, 10, 11, 12]:
-    jaxlib.append(
-        f'jaxlib @ https://us-python.pkg.dev/ml-oss-artifacts-published/jax-public-nightly-artifacts-registry/jaxlib/jaxlib-{_jaxlib_version}-cp3{python_minor_version}-cp3{python_minor_version}-manylinux2014_x86_64.whl ; python_version == "3.{python_minor_version}"'
-    )
-  return [jax] + jaxlib
+  return [
+      f'jaxlib=={_jaxlib_version}',
+      f'jax=={_jax_version}',
+  ]
 
 
 setup(
@@ -496,12 +484,12 @@ def _get_jax_install_requirements():
     },
     extras_require={
         # On Cloud TPU VM install with:
-        # pip install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-wheels/index.html -f https://storage.googleapis.com/libtpu-releases/index.html
+        # pip install torch_xla[tpu] --index-url https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/ --find-links https://storage.googleapis.com/jax-releases/libtpu_releases.html
         'tpu': [
             f'libtpu=={_libtpu_version}',
             'tpu-info',
         ],
-        # pip install torch_xla[pallas]
+        # pip install torch_xla[pallas] --index-url https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/ --find-links https://storage.googleapis.com/jax-releases/libtpu_releases.html
         'pallas': [*_get_jax_install_requirements(),]
     },
     cmdclass={
diff --git a/test/tpu/xla_test_job.yaml b/test/tpu/xla_test_job.yaml
@@ -43,7 +43,7 @@ spec:
     - |
       pip install expecttest==0.1.6
       pip install rich
-      pip install 'torch_xla[pallas]'
+      pip install --pre torch_xla[pallas] --index-url https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/ --find-links https://storage.googleapis.com/jax-releases/libtpu_releases.html
 
       cd /src/pytorch/xla
     volumeMounts:
diff --git a/torch_xla/_internal/jax_workarounds.py b/torch_xla/_internal/jax_workarounds.py
@@ -63,6 +63,8 @@ def maybe_get_jax():
     jax_import_guard()
     with jax_env_context():
       import jax
+      # TorchXLA still expects SPMD style sharding
+      jax.config.update('jax_use_shardy_partitioner', False)
       return jax
   except (ModuleNotFoundError, ImportError):
     logging.warn('You are trying to use a feature that requires jax/pallas.'
diff --git a/torch_xla/csrc/init_python_bindings.cpp b/torch_xla/csrc/init_python_bindings.cpp
@@ -1361,7 +1361,7 @@ class PyLoweringContext {
   std::string GetHloJsonText() {
     const xla::HloModuleProto& proto = computation.proto();
     std::string result;
-    google::protobuf::util::MessageToJsonString(proto, &result);
+    XLA_CHECK_OK(google::protobuf::util::MessageToJsonString(proto, &result));
     return result;
   }
 
diff --git a/torch_xla/csrc/runtime/ifrt_computation_client.cpp b/torch_xla/csrc/runtime/ifrt_computation_client.cpp
@@ -218,7 +218,8 @@ std::vector<ComputationClient::DataPtr> IfrtComputationClient::GetDataShards(
     std::vector<tsl::RCReference<xla::ifrt::Array>> arrays =
         ifrt_data->buffer
             ->DisassembleIntoSingleDeviceArrays(
-                xla::ifrt::ArrayCopySemantics::kAlwaysCopy)
+                xla::ifrt::ArrayCopySemantics::kAlwaysCopy,
+                xla::ifrt::SingleDeviceShardSemantics::kAddressableShards)
             .value();
 
     for (auto array : arrays) {
@@ -308,8 +309,7 @@ std::vector<ComputationClient::DataPtr> IfrtComputationClient::TransferToDevice(
                     ifrt_device, xla::ifrt::MemoryKind()),
                 xla::ifrt::Client::HostBufferSemantics::
                     kImmutableUntilTransferCompletes,
-                [tensor, timed]() { /* frees tensor and timer */ },
-                client_->CreateUserContext())
+                [tensor, timed]() { /* frees tensor and timer */ })
             .value();
 
     ComputationClient::DataPtr data =
diff --git a/torch_xla/csrc/runtime/profiler.cpp b/torch_xla/csrc/runtime/profiler.cpp
@@ -1,5 +1,7 @@
 #include "torch_xla/csrc/runtime/profiler.h"
 
+#include <variant>
+
 #include "absl/container/flat_hash_map.h"
 #include "absl/status/status.h"
 #include "torch_xla/csrc/runtime/tf_logging.h"
@@ -76,10 +78,23 @@ absl::Status Trace(
     int num_tracing_attempts,
     const absl::flat_hash_map<std::string, std::variant<int, std::string>>&
         options) {
+  // by 20250815 Upstream CaptureRemoteTrace changed signature of options to
+  // include bool option. For backward compatibility we don't change signature
+  // of Trace, but instead we make an adaptor to adapt the new function
+  // signature.
+  absl::flat_hash_map<std::string, std::variant<bool, int, std::string>>
+      updated_options;
+  for (const auto& item : options) {
+    if (std::holds_alternative<int>(item.second)) {
+      updated_options[item.first] = std::get<int>(item.second);
+    } else {
+      updated_options[item.first] = std::get<std::string>(item.second);
+    }
+  }
   return tsl::profiler::CaptureRemoteTrace(
       service_addr, logdir, /*worker_list=*/"",
       /*include_dataset_ops=*/false, duration_ms, num_tracing_attempts,
-      options);
+      updated_options);
 }
 
 void RegisterProfilerForPlugin(const PJRT_Api* c_api) {
diff --git a/torch_xla/csrc/runtime/xla_coordinator.cpp b/torch_xla/csrc/runtime/xla_coordinator.cpp
@@ -19,10 +19,8 @@ absl::Status XlaCoordinator::Initialize(int global_rank, int world_size,
     // Default value can be found in
     // https://github.com/openxla/xla/blob/4b88636002bc5834d7fe3f862997c66a490987bc/xla/pjrt/distributed/client.h#L63-L72.
     int heartbeat_interval_sec =
-        sys_util::GetEnvInt(env::kEnvDistSvcHeartbeatIntervalInSec, 10);
-    service_options.heartbeat_interval = absl::Seconds(heartbeat_interval_sec);
-    service_options.max_missing_heartbeats =
-        sys_util::GetEnvInt(env::kEnvDistSvcMaxMissingHeartbeats, 10);
+        sys_util::GetEnvInt(env::kEnvDistSvcHeartbeatIntervalInSec, 100);
+    service_options.heartbeat_timeout = absl::Seconds(heartbeat_interval_sec);
     int shutdown_timeout =
         sys_util::GetEnvInt(env::kEnvDistSvcShutdownTimeoutInMin, 5);
     service_options.shutdown_timeout = absl::Minutes(shutdown_timeout);
diff --git a/torch_xla/csrc/status.cpp b/torch_xla/csrc/status.cpp
@@ -8,6 +8,7 @@
 #include <stdexcept>
 
 #include "absl/log/absl_check.h"
+#include "absl/strings/str_cat.h"
 #include "tsl/platform/stacktrace.h"
 
 namespace torch_xla {

Original file line number	Diff line number	Diff line change
`@@ -112,7 +112,8 @@ function build_torch_xla() {`
`112`	`112`	`# Need to uncomment the line below.`
`113`	`113`	`# Currently it fails upstream XLA CI.`
`114`	`114`	`# pip install plugins/cuda -v`
`115`		`- pip install 'torch_xla[pallas]'`
	`115`	`+ pip install --pre torch_xla[pallas] --index-url https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/ --find-links https://storage.googleapis.com/jax-releases/libtpu_releases.html`
	`116`	`+`
`116`	`117`	`popd`
`117`	`118`	`}`
`118`	`119`