jaxlib v0.4.32 #281
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,5 +1,5 @@ | ||
| {% set version = "0.4.32" %} | ||
| {% set build = 0 %} | ||
|
|
||
| {% if cuda_compiler_version != "None" %} | ||
| {% set build = build + 200 %} | ||
|
|
@@ -13,10 +13,13 @@ package: | |
| source: | ||
| # only pull sources after upstream PyPI release... | ||
| url: https://github.com/google/jax/archive/jaxlib-v{{ version }}.tar.gz | ||
| sha256: 3fe36d596e4d640443c0a5c533845c74fbc4341e024d9bb1cd75cb49f5f419c2 | ||
| patches: | ||
| - patches/0001-Allow-for-custom-CUDA-build.patch | ||
| - patches/0002-Consolidated-build-fixes-for-XLA.patch | ||
| - patches/0003-Simplify-logic-in-jaxlib-FFI_ASSIGN_OR_RETURN-macro-.patch | ||
| - patches/0004-Fix-XLA_FFIR_REGISTER-macros.patch | ||
| - patches/0005-Add-missing-typename.patch | ||
|
|
||
| build: | ||
| number: {{ build }} | ||
|
|
@@ -57,6 +60,7 @@ requirements: | |
| - cuda-cudart-dev # [(cuda_compiler_version or "").startswith("12")] | ||
| - cuda-nvml-dev # [(cuda_compiler_version or "").startswith("12")] | ||
| - cuda-nvtx-dev # [(cuda_compiler_version or "").startswith("12")] | ||
| - cuda-nvcc-tools # [(cuda_compiler_version or "").startswith("12")] | ||
|
There was a problem hiding this comment. Is there a way to confirm that JAX is using There was a problem hiding this comment. I have no idea. @mwittgen would you know how to do this/where to look? We sadly don't have a GPU device in the CI to test it. There was a problem hiding this comment. Ah was hoping this kind of check would not actually use the GPU. Sort of like If we are unable to figure out an appropriate config check, we could configure CI to use a GPU. Looks like some work has already been done here. So this would just be adding another runner type Another option we have made heavy use of elsewhere is caching the built packages on CI and then downloading them for local testing ( conda-forge/conda-forge.github.io#1424 ). This may be more useful in the end as it is easier to iterate There was a problem hiding this comment.
I looked at its implementation and it doesn't provide anything useful for the case we are testing on.
This seems like rather expensive to add a GPU here to do a simple test. Overall, it seems that this will need input from someone where it doesn't work. The current PR is an improvement on the status quo and I don't want to hold up the release anymore, thus I'm going to merge. |
||
| - libcublas-dev # [(cuda_compiler_version or "").startswith("12")] | ||
| - libcusolver-dev # [(cuda_compiler_version or "").startswith("12")] | ||
| - libcurand-dev # [(cuda_compiler_version or "").startswith("12")] | ||
|
|
@@ -65,6 +69,7 @@ requirements: | |
| - python | ||
| - pip | ||
| - numpy | ||
| - setuptools | ||
| - wheel | ||
| - cuda-version {{ cuda_compiler_version }} # [cuda_compiler_version != "None"] | ||
| # avoid not being able to pass `-C=--build-option=--python-tag=cp<x>` due to | ||
|
|
@@ -82,9 +87,7 @@ requirements: | |
| - scipy >=1.9 | ||
| - ml_dtypes >=0.2.0 | ||
| - __cuda # [cuda_compiler_version != "None"] | ||
| - cuda-nvcc-tools # [(cuda_compiler_version or "").startswith("12")] | ||
|
Comment on lines
-85
to
+90
There was a problem hiding this comment. Thanks Uwe! 🙏 Looks good (as does the |
||
| run_constrained: | ||
| - jax >={{ version }} | ||
|
|
||
|
|
||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,25 +1,25 @@ | ||
| From 1ec53ea591323e47c8ce53ed9b0736e98784ff68 Mon Sep 17 00:00:00 2001 | ||
| From: "Uwe L. Korn" <[email protected]> | ||
| Date: Sun, 8 Oct 2023 19:34:34 +0200 | ||
| Subject: [PATCH 1/5] Allow for custom CUDA build | ||
|
|
||
| --- | ||
| build/build.py | 5 +++++ | ||
| 1 file changed, 5 insertions(+) | ||
|
|
||
| diff --git a/build/build.py b/build/build.py | ||
| index c3a8627..5bfbdcb 100755 | ||
| --- a/build/build.py | ||
| +++ b/build/build.py | ||
| @@ -289,6 +289,11 @@ def write_bazelrc(*, remote_build, | ||
| f.write("build --config=nvcc_clang\n") | ||
| f.write(f"build --action_env=CLANG_CUDA_COMPILER_PATH={clang_path}\n") | ||
| if cuda_version: | ||
| + # set GCC_HOST_COMPILER_PATH for toolchain with conda-forge | ||
| + f.write("build --action_env GCC_HOST_COMPILER_PATH=\"{gcc_host_compiler_path}\"\n" | ||
| + .format(gcc_host_compiler_path=os.environ["GCC"])) | ||
| + f.write("build --action_env GCC_HOST_COMPILER_PREFIX=\"{gcc_host_compiler_prefix}\"\n" | ||
| + .format(gcc_host_compiler_prefix=os.path.dirname(os.environ["GCC"]))) | ||
| f.write("build --repo_env HERMETIC_CUDA_VERSION=\"{cuda_version}\"\n" | ||
| .format(cuda_version=cuda_version)) | ||
| if cudnn_version: |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,7 +1,7 @@ | ||
| From a7e732c129f51d16dd59b353e8c66ceae9b5529c Mon Sep 17 00:00:00 2001 | ||
| From: "Uwe L. Korn" <[email protected]> | ||
| Date: Thu, 14 Dec 2023 17:06:15 +0100 | ||
| Subject: [PATCH 2/5] Consolidated build fixes for XLA | ||
|
|
||
| jax vendors xla, but only populates the sources through bazel, so we cannot | ||
| patch as usual through conda, but rather need to teach the bazel build file | ||
|
|
@@ -14,24 +14,75 @@ which is also where we're patching in the list of patches to apply to xla. | |
|
|
||
| Co-Authored-By: H. Vetinari <[email protected]> | ||
| --- | ||
| .../xla/0001-Omit-usage-of-StrFormat.patch | 43 ++++ | ||
| ...pport-third-party-build-of-boringssl.patch | 51 ++++ | ||
| third_party/xla/0002-Fix-abseil-headers.patch | 73 ++++++ | ||
| .../xla/0003-Omit-usage-of-StrFormat.patch | 43 ++++ | ||
| ...0004-Add-missing-bits-absl-systemlib.patch | 226 ++++++++++++++++++ | ||
| third_party/xla/workspace.bzl | 6 + | ||
| 6 files changed, 442 insertions(+) | ||
| create mode 100644 third_party/xla/0001-Omit-usage-of-StrFormat.patch | ||
| create mode 100644 third_party/xla/0001-Support-third-party-build-of-boringssl.patch | ||
| create mode 100644 third_party/xla/0002-Fix-abseil-headers.patch | ||
| create mode 100644 third_party/xla/0003-Omit-usage-of-StrFormat.patch | ||
| create mode 100644 third_party/xla/0004-Add-missing-bits-absl-systemlib.patch | ||
|
|
||
| diff --git a/third_party/xla/0001-Omit-usage-of-StrFormat.patch b/third_party/xla/0001-Omit-usage-of-StrFormat.patch | ||
| new file mode 100644 | ||
| index 0000000..d1b4765 | ||
| --- /dev/null | ||
| +++ b/third_party/xla/0001-Omit-usage-of-StrFormat.patch | ||
| @@ -0,0 +1,43 @@ | ||
| +From b7d3f685ea9f58f0054af0f34d0bc3ccac43fa5c Mon Sep 17 00:00:00 2001 | ||
| +From: "Uwe L. Korn" <[email protected]> | ||
| +Date: Thu, 4 Jul 2024 10:36:03 +0200 | ||
| +Subject: [PATCH] Omit usage of StrFormat | ||
| + | ||
| +--- | ||
| + xla/stream_executor/gpu/gpu_executor.h | 9 ++++++--- | ||
| + 1 file changed, 6 insertions(+), 3 deletions(-) | ||
| + | ||
| +diff --git a/xla/stream_executor/gpu/gpu_executor.h b/xla/stream_executor/gpu/gpu_executor.h | ||
| +index 3a5945e884..9bdc2acd47 100644 | ||
| +--- a/xla/stream_executor/gpu/gpu_executor.h | ||
| ++++ b/xla/stream_executor/gpu/gpu_executor.h | ||
| +@@ -29,6 +29,7 @@ limitations under the License. | ||
| + #include <optional> | ||
| + #include <set> | ||
| + #include <string> | ||
| ++#include <sstream> | ||
| + #include <unordered_map> | ||
| + #include <utility> | ||
| + | ||
| +@@ -38,7 +39,6 @@ limitations under the License. | ||
| + #include "absl/numeric/int128.h" | ||
| + #include "absl/status/status.h" | ||
| + #include "absl/status/statusor.h" | ||
| +-#include "absl/strings/str_format.h" | ||
| + #include "absl/synchronization/mutex.h" | ||
| + #include "absl/types/span.h" | ||
| + #include "xla/stream_executor/blas.h" | ||
| +@@ -187,8 +187,11 @@ class GpuExecutor : public StreamExecutor { | ||
| + uint64_t size) override { | ||
| + auto* buffer = GpuDriver::HostAllocate(context_, size); | ||
| + if (buffer == nullptr && size > 0) { | ||
| +- return absl::InternalError( | ||
| +- absl::StrFormat("Failed to allocate HostMemory of size %d", size)); | ||
| ++ std::ostringstream stringStream; | ||
| ++ stringStream << "Failed to allocate HostMemory of size "; | ||
| ++ stringStream << size; | ||
| ++ std::string res = stringStream.str(); | ||
| ++ return absl::InternalError(res); | ||
| + } | ||
| + return std::make_unique<HostMemoryAllocation>(buffer, size, this); | ||
| + } | ||
| diff --git a/third_party/xla/0001-Support-third-party-build-of-boringssl.patch b/third_party/xla/0001-Support-third-party-build-of-boringssl.patch | ||
| new file mode 100644 | ||
| index 0000000..26a9904 | ||
| --- /dev/null | ||
| +++ b/third_party/xla/0001-Support-third-party-build-of-boringssl.patch | ||
| @@ -0,0 +1,51 @@ | ||
| +From 9a5932bb8a363f777ce39ff75a52eda2bba9c21f Mon Sep 17 00:00:00 2001 | ||
| +From: "Uwe L. Korn" <[email protected]> | ||
| +Date: Thu, 14 Dec 2023 15:04:51 +0100 | ||
| +Subject: [PATCH 1/4] Support third-party build of boringssl | ||
|
|
@@ -70,10 +121,10 @@ index 000000000..e24a45e1f | |
| ++ ], | ||
| ++) | ||
| +diff --git a/workspace2.bzl b/workspace2.bzl | ||
| +index 1809702d8b..6fc538d3a2 100644 | ||
| +--- a/workspace2.bzl | ||
| ++++ b/workspace2.bzl | ||
| +@@ -69,7 +69,7 @@ def _tf_repositories(): | ||
| + name = "boringssl", | ||
| + sha256 = "9dc53f851107eaf87b391136d13b815df97ec8f76dadb487b58b2fc45e624d2c", | ||
| + strip_prefix = "boringssl-c00d7ca810e93780bd0c8ee4eea28f4f2ea4bcdc", | ||
|
|
@@ -84,11 +135,11 @@ index 000000000..e24a45e1f | |
| + | ||
| diff --git a/third_party/xla/0002-Fix-abseil-headers.patch b/third_party/xla/0002-Fix-abseil-headers.patch | ||
| new file mode 100644 | ||
| index 0000000..96a6fec | ||
| --- /dev/null | ||
| +++ b/third_party/xla/0002-Fix-abseil-headers.patch | ||
| @@ -0,0 +1,73 @@ | ||
| +From 97ad75d0bf891be488fb223ac95ff6572b4ecd88 Mon Sep 17 00:00:00 2001 | ||
| +From: "Uwe L. Korn" <[email protected]> | ||
| +Date: Thu, 23 May 2024 15:45:52 +0200 | ||
| +Subject: [PATCH 2/4] Fix abseil headers | ||
|
|
@@ -101,10 +152,10 @@ index 000000000..7a58075e1 | |
| + 4 files changed, 10 insertions(+) | ||
| + | ||
| +diff --git a/third_party/tsl/tsl/platform/default/BUILD b/third_party/tsl/tsl/platform/default/BUILD | ||
| +index b3ce4301fb..9b72c2eb42 100644 | ||
| +--- a/third_party/tsl/tsl/platform/default/BUILD | ||
| ++++ b/third_party/tsl/tsl/platform/default/BUILD | ||
| +@@ -225,6 +225,8 @@ cc_library( | ||
| + deps = [ | ||
| + "//tsl/platform:logging", | ||
| + "@com_google_absl//absl/log:check", | ||
|
|
@@ -114,7 +165,7 @@ index 000000000..7a58075e1 | |
| + ) | ||
| + | ||
| +diff --git a/third_party/tsl/tsl/profiler/rpc/client/BUILD b/third_party/tsl/tsl/profiler/rpc/client/BUILD | ||
| +index 4b8ece7403..a2772846b8 100644 | ||
| +--- a/third_party/tsl/tsl/profiler/rpc/client/BUILD | ||
| ++++ b/third_party/tsl/tsl/profiler/rpc/client/BUILD | ||
| +@@ -101,6 +101,8 @@ cc_library( | ||
|
|
@@ -150,7 +201,7 @@ index 000000000..7a58075e1 | |
| + ], | ||
| + alwayslink = True, | ||
| +diff --git a/xla/tsl/distributed_runtime/rpc/BUILD b/xla/tsl/distributed_runtime/rpc/BUILD | ||
| +index 817c4dc5a4..d6f27deb5c 100644 | ||
| +--- a/xla/tsl/distributed_runtime/rpc/BUILD | ||
| ++++ b/xla/tsl/distributed_runtime/rpc/BUILD | ||
| +@@ -37,6 +37,7 @@ cc_library( | ||
|
|
@@ -163,11 +214,11 @@ index 000000000..7a58075e1 | |
| + "@tsl//tsl/platform:status", | ||
| diff --git a/third_party/xla/0003-Omit-usage-of-StrFormat.patch b/third_party/xla/0003-Omit-usage-of-StrFormat.patch | ||
| new file mode 100644 | ||
| index 0000000..67d2275 | ||
| --- /dev/null | ||
| +++ b/third_party/xla/0003-Omit-usage-of-StrFormat.patch | ||
| @@ -0,0 +1,43 @@ | ||
| +From a360cd33b748c4f6b1ab00e386ac8031112c5b2f Mon Sep 17 00:00:00 2001 | ||
| +From: "Uwe L. Korn" <[email protected]> | ||
| +Date: Thu, 4 Jul 2024 10:36:03 +0200 | ||
| +Subject: [PATCH 3/4] Omit usage of StrFormat | ||
|
|
@@ -177,17 +228,17 @@ index 000000000..541c06f40 | |
| + 1 file changed, 6 insertions(+), 3 deletions(-) | ||
| + | ||
| +diff --git a/xla/stream_executor/gpu/gpu_executor.h b/xla/stream_executor/gpu/gpu_executor.h | ||
| +index 8e9a8352e2..36d42493c6 100644 | ||
| +--- a/xla/stream_executor/gpu/gpu_executor.h | ||
| ++++ b/xla/stream_executor/gpu/gpu_executor.h | ||
| +@@ -27,6 +27,7 @@ limitations under the License. | ||
| + #include <memory> | ||
| + #include <optional> | ||
| + #include <string> | ||
| ++#include <sstream> | ||
| + #include <unordered_map> | ||
| + #include <utility> | ||
| + #include <variant> | ||
| +@@ -37,7 +38,6 @@ limitations under the License. | ||
| + #include "absl/numeric/int128.h" | ||
| + #include "absl/status/status.h" | ||
|
|
@@ -196,7 +247,7 @@ index 000000000..541c06f40 | |
| + #include "absl/synchronization/mutex.h" | ||
| + #include "absl/types/span.h" | ||
| + #include "xla/stream_executor/blas.h" | ||
| +@@ -166,8 +166,11 @@ class GpuExecutor : public StreamExecutorCommon { | ||
| + uint64_t size) override { | ||
| + auto* buffer = GpuDriver::HostAllocate(context_, size); | ||
| + if (buffer == nullptr && size > 0) { | ||
|
|
@@ -212,16 +263,15 @@ index 000000000..541c06f40 | |
| + } | ||
| diff --git a/third_party/xla/0004-Add-missing-bits-absl-systemlib.patch b/third_party/xla/0004-Add-missing-bits-absl-systemlib.patch | ||
| new file mode 100644 | ||
| index 0000000..1941f79 | ||
| --- /dev/null | ||
| +++ b/third_party/xla/0004-Add-missing-bits-absl-systemlib.patch | ||
| @@ -0,0 +1,226 @@ | ||
| +From fc6d67a2f5fce78eb91477fa4bca5c47b6fc31fd Mon Sep 17 00:00:00 2001 | ||
| +From: "Uwe L. Korn" <[email protected]> | ||
| +Date: Thu, 4 Jul 2024 15:58:32 +0200 | ||
| +Subject: [PATCH 4/4] Add missing bits absl systemlib | ||
| + | ||
| +--- | ||
| + .../third_party/absl/system.absl.base.BUILD | 16 +++++ | ||
| + .../third_party/absl/system.absl.crc.BUILD | 70 +++++++++++++++++++ | ||
|
|
@@ -425,14 +475,15 @@ index 000000000..e151c23c8 | |
| + name = "strings", | ||
| + linkopts = ["-labsl_strings"], | ||
| +diff --git a/third_party/tsl/third_party/absl/workspace.bzl b/third_party/tsl/third_party/absl/workspace.bzl | ||
| +index 9565a82c33..e71aa16726 100644 | ||
| +--- a/third_party/tsl/third_party/absl/workspace.bzl | ||
| ++++ b/third_party/tsl/third_party/absl/workspace.bzl | ||
| +@@ -14,12 +14,14 @@ def repo(): | ||
| + SYS_DIRS = [ | ||
| + "algorithm", | ||
| + "base", | ||
| ++ "crc", | ||
| + "cleanup", | ||
| + "container", | ||
| + "debugging", | ||
| + "flags", | ||
|
|
@@ -443,7 +494,7 @@ index 000000000..e151c23c8 | |
| + "meta", | ||
| + "numeric", | ||
| diff --git a/third_party/xla/workspace.bzl b/third_party/xla/workspace.bzl | ||
| index 8f4accc..3b7afaf 100644 | ||
| --- a/third_party/xla/workspace.bzl | ||
| +++ b/third_party/xla/workspace.bzl | ||
| @@ -30,6 +30,12 @@ def repo(): | ||
|
|
||
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Am sorry to see this is needed 😢
Is it able to parse
CUDA_HOMEwith multiple paths. IOW would this be possible?There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
No, in the end the bazel scripts do some globbing over
$LOCAL_CUDA_PATH/bin*to determine the available CUDA binaries. You can do a bit of trickery and use one ofLOCAL_CUDNN_PATH/LOCAL_NCCL_PATHto add a third directory but this then sadly led to the wrong protobuf headers being picked up.