{tools}[gfbf/2023a] jax v0.4.25 w/ CUDA 12.1.1 by ThomasHoffmann77 · Pull Request #20119 · easybuilders/easybuild-easyconfigs

ThomasHoffmann77 · 2024-03-14T15:17:37Z

(created using eb --new-pr)
requires:

{tools}[gfbf/2023a] ml_dtypes v0.3.2 #20707

edit: requires bug fix in framework for "cp %s %(builddir)s/archives" to work as extract command:

Fix resolved (template) values in case of failure easybuild-framework#4532

ThomasHoffmann77 · 2024-03-14T15:54:22Z

Test report by @ThomasHoffmann77
FAILED
Build succeeded for 1 out of 2 (2 easyconfigs in total)
srv-mahamid-01.embl.de - Linux AlmaLinux 8.8, x86_64, AMD EPYC 7513 32-Core Processor, 2 x NVIDIA NVIDIA GeForce RTX 3090, 535.113.01, Python 3.6.8
See https://gist.github.com/ThomasHoffmann77/43d87811306655a013126860c0bb6777 for a full test report.

ThomasHoffmann77 · 2024-03-14T16:10:36Z

Test report by @ThomasHoffmann77
FAILED
Build succeeded (with --ignore-test-failure) for 1 out of 2 (2 easyconfigs in total)
srv-mahamid-01.embl.de - Linux AlmaLinux 8.8, x86_64, AMD EPYC 7513 32-Core Processor, 2 x NVIDIA NVIDIA GeForce RTX 3090, 535.113.01, Python 3.6.8
See https://gist.github.com/ThomasHoffmann77/c51c43986eae5a7afe56f715d7c5c38c for a full test report.

ThomasHoffmann77 · 2024-03-14T18:28:48Z

Test report by @ThomasHoffmann77
SUCCESS
Build succeeded (with --ignore-test-failure) for 2 out of 2 (2 easyconfigs in total)
proline - Linux AlmaLinux 8.8, x86_64, 12th Gen Intel(R) Core(TM) i7-12700, 1 x NVIDIA NVIDIA RTX A4000, 535.113.01, Python 3.6.8
See https://gist.github.com/ThomasHoffmann77/b2b075b38d9d9d5c6fe4b0503dab7279 for a full test report.

branfosj · 2024-03-14T21:39:44Z

Test report by @branfosj
SUCCESS
Build succeeded (with --ignore-test-failure) for 2 out of 2 (2 easyconfigs in total)
bear-pg0208u15a - Linux RHEL 8.6, x86_64, Intel(R) Xeon(R) Platinum 8360Y CPU @ 2.40GHz (icelake), 1 x NVIDIA NVIDIA A100-SXM4-40GB, 535.154.05, Python 3.6.8
See https://gist.github.com/branfosj/83b07adf11f9a9eea619d5b7e45eddb5 for a full test report.

Same three failures as #19841 (comment)

branfosj · 2024-03-15T00:20:37Z

Test report by @branfosj
SUCCESS
Build succeeded (with --ignore-test-failure) for 2 out of 2 (2 easyconfigs in total)
bear-pg0208u31a - Linux RHEL 8.6, x86_64, Intel(R) Xeon(R) Platinum 8360Y CPU @ 2.40GHz (icelake), 4 x NVIDIA NVIDIA A100-SXM4-40GB, 535.154.05, Python 3.6.8
See https://gist.github.com/branfosj/bec290f9c00aa6309ee649e8ff185675 for a full test report.

Same three failures as #19841 (comment)

ThomasHoffmann77 · 2024-03-15T13:37:46Z

Test report by @ThomasHoffmann77
SUCCESS
Build succeeded (with --ignore-test-failure) for 2 out of 2 (2 easyconfigs in total)
proline - Linux AlmaLinux 8.8, x86_64, 12th Gen Intel(R) Core(TM) i7-12700, 1 x NVIDIA NVIDIA RTX A4000, 535.113.01, Python 3.6.8
See https://gist.github.com/ThomasHoffmann77/59e7a52712f520a524e93b5b5210551b for a full test report.

verdurin · 2024-03-15T15:29:52Z

I don't have a build node setup to upload test reports. Did see this test error:

tests/lax_scipy_special_functions_test.py::LaxScipySpcialFunctionsTest::testScipySpecialFun_gammainc_s_1x4_float32_float64 PASSED                                                                                                    [ 55%]
tests/lax_scipy_special_functions_test.py::LaxScipySpcialFunctionsTest::testScipySpecialFun_gammainc_s_2x1x4_float32_float32 Fatal Python error: Aborted

verdurin · 2024-03-15T15:38:32Z

I see you're all building with --ignore-test-failure - is that expected with jax?

ThomasHoffmann77 · 2024-03-15T15:43:32Z

I don't have a build node setup to upload test reports. Did see this test error:

tests/lax_scipy_special_functions_test.py::LaxScipySpcialFunctionsTest::testScipySpecialFun_gammainc_s_1x4_float32_float64 PASSED                                                                                                    [ 55%]
tests/lax_scipy_special_functions_test.py::LaxScipySpcialFunctionsTest::testScipySpecialFun_gammainc_s_2x1x4_float32_float32 Fatal Python error: Aborted

#16:38 thoffman@srv-mahamid-01#NVIDIA_TF32_OVERRIDE=0 CUDA_VISIBLE_DEVICES=0 XLA_PYTHON_CLIENT_ALLOCATOR=platform JAX_ENABLE_X64=true pytest -vv tests/lax_scipy_special_functions_test.py::LaxScipySpcialFunctionsTest::testScipySpecialFun_gammainc_s_2x1x4_float32_float32
============================= test session starts ==============================
platform linux -- Python 3.11.3, pytest-7.4.2, pluggy-1.2.0 -- /g/easybuild/x86_64/Rocky/8/rome/software/Python/3.11.3-GCCcore-12.3.0/bin/python
cachedir: .pytest_cache
hypothesis profile 'default' -> database=DirectoryBasedExampleDatabase(PosixPath('/tmp/jax-jax-v0.4.25/.hypothesis/examples'))
rootdir: /tmp/jax-jax-v0.4.25
configfile: pyproject.toml
plugins: xdist-3.3.1, hypothesis-6.88.1
collected 1 item                                                               

tests/lax_scipy_special_functions_test.py::LaxScipySpcialFunctionsTest::testScipySpecialFun_gammainc_s_2x1x4_float32_float32 PASSED [100%]
============================== 1 passed in 3.21s ===============================

ThomasHoffmann77 · 2024-03-19T13:39:11Z

Test report by @ThomasHoffmann77
SUCCESS
Build succeeded for 2 out of 2 (2 easyconfigs in total)
srv-mahamid-01.embl.de - Linux AlmaLinux 8.8, x86_64, AMD EPYC 7513 32-Core Processor, 2 x NVIDIA NVIDIA GeForce RTX 3090, 535.113.01, Python 3.6.8
See https://gist.github.com/ThomasHoffmann77/37e79d6b1006b4e8bee5438a97ef2ccd for a full test report.

ThomasHoffmann77 · 2024-03-19T14:02:28Z

Test report by @ThomasHoffmann77
SUCCESS
Build succeeded for 2 out of 2 (2 easyconfigs in total)
proline - Linux AlmaLinux 8.8, x86_64, 12th Gen Intel(R) Core(TM) i7-12700, 1 x NVIDIA NVIDIA RTX A4000, 535.113.01, Python 3.6.8
See https://gist.github.com/ThomasHoffmann77/61812c0e50d74c911c9d72e03155eac6 for a full test report.

Flamefire · 2024-03-22T13:31:16Z

Test report by @Flamefire
FAILED
Build succeeded for 6 out of 7 (2 easyconfigs in total)
n1438 - Linux RHEL 8.7 (Ootpa), x86_64, Intel(R) Xeon(R) Platinum 8470 (icelake), Python 3.8.13
See https://gist.github.com/Flamefire/ee41d9059916ce8b1f93b9267d0c847f for a full test report.

Flamefire · 2024-03-22T13:35:38Z

Test report by @Flamefire
FAILED
Build succeeded for 1 out of 2 (2 easyconfigs in total)
i8002 - Linux Rocky Linux 8.7 (Green Obsidian), x86_64, AMD EPYC 7352 24-Core Processor (zen2), 8 x NVIDIA NVIDIA A100-SXM4-40GB, 545.23.08, Python 3.8.13
See https://gist.github.com/Flamefire/327109d42642f3d3ed5c28565c08f20b for a full test report.

Flamefire · 2024-03-22T14:51:42Z

In both cases the failure is:

external/upb/upb/table.c: In function upb_inttable_pop:
external/upb/upb/table.c:588:10: error: val.val may be used uninitialized [-Werror=maybe-uninitialized]
  588 |   return val;
      |          ^~~
external/upb/upb/table.c:585:13: note: val.val was declared here
  585 |   upb_value val;
      |             ^~~

Due to -Werror added here

XLA comes with even more dependencies (workspace*.bzl). Can we add them as local repositories too? Maybe even auto-generate those lists via a Python script or so (similar to e.g. findPythonDeps which outputs a list of Python packages for use in an EC. That script is bundled with EasyBuild so readily available)

easybuild/easyconfigs/j/jax/jax-0.4.25-foss-2023a-CUDA-12.1.1.eb

Co-authored-by: Alexander Grund <Flamefire@users.noreply.github.com>

Flamefire · 2024-03-26T12:14:33Z

Test report by @Flamefire
FAILED
Build succeeded for 1 out of 2 (2 easyconfigs in total)
n1265 - Linux RHEL 8.7 (Ootpa), x86_64, Intel(R) Xeon(R) Platinum 8470 (icelake), Python 3.8.13
See https://gist.github.com/Flamefire/8cbb16221ab8da073cee85c97c0dd911 for a full test report.

This is caused by a crash. It isn't really clear why it fails or in which test, as when I run the crashing test file manually it works. Attaching GDB shows ~LogMessageFatal() as the cause. Need more investigation into why, i.e. what the fatal error is, but this looks serious...

easybuild/easyconfigs/j/jax/jax-0.4.25-foss-2023a-CUDA-12.1.1.eb

easybuild/easyconfigs/m/ml_dtypes/ml_dtypes-0.3.2-foss-2023a.eb

easybuild/easyconfigs/j/jax/jax-0.4.25-foss-2023a-CUDA-12.1.1.eb

casparvl · 2024-06-11T21:34:15Z

Test report by @casparvl
FAILED
Build succeeded for 0 out of 1 (1 easyconfigs in total)
gcn159.local.snellius.surf.nl - Linux RHEL 8.6, x86_64, AMD EPYC 9334 32-Core Processor, 4 x NVIDIA NVIDIA H100, 545.23.08, Python 3.6.8
See https://gist.github.com/casparvl/1f92447b1dfb9a8ef9cd7e0d72e739e0 for a full test report.

akesandgren · 2024-06-12T05:31:13Z

Nope, using single compute capability for the H100 (9.0) also fails in the same way.

ThomasHoffmann77 · 2024-06-12T08:04:47Z

Nope, using single compute capability for the H100 (9.0) also fails in the same way.

0.4.29 has been released in the meantime. It might be worth to try this version.

akesandgren · 2024-06-13T10:53:58Z

First attempt at using 0.4.29 with this toolchain failed:

tests/sparse_nm_test.py::SpmmTest::test_shapes0 Fatal Python error: Aborted

Thread 0x000014a919150000 (most recent call first):
  File "/dev/shm/ake/build/jax/0.4.29/gfbf-2023a-CUDA-12.1.1/jax/jax-jax-v0.4.29/jax/_src/compiler.py", line 238
 in backend_compile
  File "/dev/shm/ake/build/jax/0.4.29/gfbf-2023a-CUDA-12.1.1/jax/jax-jax-v0.4.29/jax/_src/profiler.py", line 335
 in wrapper
  File "/dev/shm/ake/build/jax/0.4.29/gfbf-2023a-CUDA-12.1.1/jax/jax-jax-v0.4.29/jax/_src/compiler.py", line 608
 in _compile_and_write_cache
  File "/dev/shm/ake/build/jax/0.4.29/gfbf-2023a-CUDA-12.1.1/jax/jax-jax-v0.4.29/jax/_src/compiler.py", line 378
 in compile_or_get_cached
  File "/dev/shm/ake/build/jax/0.4.29/gfbf-2023a-CUDA-12.1.1/jax/jax-jax-v0.4.29/jax/_src/interpreters/pxla.py",
 line 2779 in _cached_compilation
...

VRehnberg · 2024-06-13T13:18:38Z

Test report by @VRehnberg
SUCCESS
Build succeeded for 4 out of 4 (1 easyconfigs in total)
alvis4-41 - Linux Rocky Linux 8.9, x86_64, Intel(R) Xeon(R) Gold 6338 CPU @ 2.00GHz, 1 x NVIDIA NVIDIA A100-SXM4-80GB, 545.23.08, Python 3.6.8
See https://gist.github.com/VRehnberg/d5fe8ca36746f898d5d1defc484c686e for a full test report.

VRehnberg · 2024-06-13T13:21:52Z

Test report by @VRehnberg
SUCCESS
Build succeeded for 4 out of 4 (1 easyconfigs in total)
alvis3-16 - Linux Rocky Linux 8.9, x86_64, Intel(R) Xeon(R) Gold 6338 CPU @ 2.00GHz, 1 x NVIDIA NVIDIA A100-SXM4-40GB, 545.23.08, Python 3.6.8
See https://gist.github.com/VRehnberg/332751b5d6100151342f5f2ccbcc7608 for a full test report.

VRehnberg · 2024-06-18T15:37:39Z

Test report by @VRehnberg
SUCCESS
Build succeeded for 4 out of 4 (1 easyconfigs in total)
alvis9-07 - Linux Rocky Linux 8.9, x86_64, Intel(R) Xeon(R) Gold 6338 CPU @ 2.00GHz, 1 x NVIDIA NVIDIA A40, 550.54.14, Python 3.6.8
See https://gist.github.com/VRehnberg/a7488643356322088e5d131a18309892 for a full test report.

akesandgren · 2024-06-26T05:12:05Z

To get this to run on H100 one needs a newer CUDA, 0.4.29 with foss/2023a and CUDA/12.5.0 passes all but a single broken test, i.e. the test itself is broken...

akesandgren

LGTM

Change added

ThomasHoffmann77 · 2024-06-26T07:53:22Z

🎉

easybuild/easyconfigs/j/jax/jax-0.4.25-gfbf-2023a-CUDA-12.1.1.eb

@akesandgren

fix local_extract_cmd according to @akesandgren 's suggestion

VRehnberg · 2024-07-09T11:40:26Z

Edit: Recent change in global pip.conf made build fail. Unrelated to this PR.

Flamefire · 2024-07-09T11:58:27Z

Edit: Recent change in global pip.conf made build fail. Unrelated to this PR.

Which one exactly and what was the error? Might be worth addressing in framework

VRehnberg · 2024-07-09T12:02:39Z

Edit: Recent change in global pip.conf made build fail. Unrelated to this PR.

Which one exactly and what was the error? Might be worth addressing in framework

ERROR: Could not find an activated virtualenv (required).

https://gist.github.com/VRehnberg/5be54199260e8a478002d18dd986725c

[install]
require-virtualenv = true

VRehnberg · 2024-07-09T15:06:27Z

Test report by @VRehnberg
SUCCESS
Build succeeded for 1 out of 1 (1 easyconfigs in total)
alvis1-11 - Linux Rocky Linux 8.9, x86_64, Intel(R) Xeon(R) Gold 6244 CPU @ 3.60GHz, 1 x NVIDIA Tesla V100-SXM2-32GB, 550.54.14, Python 3.6.8
See https://gist.github.com/VRehnberg/141b0e91571ee25d1ea81ddb0e05c169 for a full test report.

VRehnberg · 2024-07-09T19:29:13Z

Test report by @VRehnberg
SUCCESS
Build succeeded for 1 out of 1 (1 easyconfigs in total)
alvis2-02 - Linux Rocky Linux 8.9, x86_64, Intel(R) Xeon(R) Gold 6226R CPU @ 2.90GHz, 1 x NVIDIA Tesla T4, 550.54.14, Python 3.6.8
See https://gist.github.com/VRehnberg/a10d64ffec054c2de2cef9a0960a1fbd for a full test report.

Flamefire · 2024-07-10T07:46:03Z

ERROR: Could not find an activated virtualenv (required).

Ah I remember that. There is a fix in the easyblocks: easybuilders/easybuild-easyblocks#3374
If that isn't used in jax (custom easyblocks might skip/replace that step) we need to include that there. Haven't checked that it does yet

VRehnberg · 2024-07-10T07:49:16Z

Ah I remember that.

Thanks, wasn't using that easyblock still on 4.9.2 easyblocks.

It seems like the jax test suite is very VRAM hungry, since they only pass on A100 for some reason?

So the tests are not VRAM hungry. I never saw more than 1 GB used and T4 only have 16 GB in total so that's a strict limit even if our monitoring would miss a short spike. Does use about 27 GB of regular RAM though in case that could be an issue.

lexming · 2024-07-23T19:04:42Z

Test report by @lexming
SUCCESS
Build succeeded for 1 out of 1 (1 easyconfigs in total)
node409.hydra.os - Linux Rocky Linux 8.10, x86_64, AMD EPYC 7282 16-Core Processor (zen2), 1 x NVIDIA NVIDIA A100-PCIE-40GB, 550.90.07, Python 3.6.8
See https://gist.github.com/lexming/650ff1e0bf428372c2ea23fc00fbba4e for a full test report.

lexming

LGTM

review addressed

lexming · 2024-07-24T05:24:13Z

Merging, thanks a lot for keeping up with all the issues @ThomasHoffmann77 !

{tools}[foss/2023a] jax v0.4.25, ml_dtypes v0.3.2 w/ CUDA 12.1.1

8e4e184

ThomasHoffmann77 marked this pull request as draft March 14, 2024 15:17

ThomasHoffmann77 mentioned this pull request Mar 14, 2024

{tools}[foss/2023a] jax v0.4.24 w/ CUDA 12.1.1 #19841

Closed

fix style

94e9567

fix test_no_log_spam

1827967

update patch

021b248

boegel added the update label Mar 17, 2024

boegel added this to the 4.x milestone Mar 17, 2024

isolate some tests

a90bc88

ThomasHoffmann77 mentioned this pull request Mar 19, 2024

{bio}[foss/2023a] AlphaFold v2.3.2, dm-haiku v0.0.12, tensorstore v0.1.65 w/ CUDA v12.1.1 #19942

Merged

2 tasks

ThomasHoffmann77 marked this pull request as ready for review March 19, 2024 14:28

branfosj mentioned this pull request Mar 22, 2024

{lib} [foss/2023a] Bazel v6.1.0, grpcio v1.57.0, TensorFlow v2.15.1 #20191

Merged

Flamefire reviewed Mar 22, 2024

View reviewed changes

easybuild/easyconfigs/j/jax/jax-0.4.25-foss-2023a-CUDA-12.1.1.eb Outdated Show resolved Hide resolved

Update easybuild/easyconfigs/j/jax/jax-0.4.25-foss-2023a-CUDA-12.1.1.eb

8c5d1f3

Co-authored-by: Alexander Grund <Flamefire@users.noreply.github.com>

akesandgren requested changes Mar 27, 2024

View reviewed changes

pavelToman mentioned this pull request Jun 14, 2024

{bio}[foss/2023a] scArches v0.6.1, scib-metrics v0.5.1, scvi-tools v1.1.2, ... #20814

Merged

akesandgren approved these changes Jun 26, 2024

View reviewed changes

akesandgren modified the milestones: 4.x, release after 4.9.2 Jun 26, 2024

akesandgren previously requested changes Jun 27, 2024

View reviewed changes

easybuild/easyconfigs/j/jax/jax-0.4.25-gfbf-2023a-CUDA-12.1.1.eb Outdated Show resolved Hide resolved

Update jax-0.4.25-gfbf-2023a-CUDA-12.1.1.eb

c6518ee

fix local_extract_cmd according to @akesandgren 's suggestion

lexming approved these changes Jul 24, 2024

View reviewed changes

lexming merged commit 3ea4e45 into easybuilders:develop Jul 24, 2024

lexming mentioned this pull request Jul 24, 2024

reorganize Flax/JAX stack in 2023a: move jax + Optax to gfbf/2023a toolchain + use standalone Flax + absl-py as dependencies #21038

Merged

pavelToman mentioned this pull request Sep 11, 2025

Update jaxlib easyblock for jax/jaxlib v0.6.2 and v0.7.0 easybuilders/easybuild-easyblocks#3852

Merged

Conversation

ThomasHoffmann77 commented Mar 14, 2024 • edited by casparvl Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ThomasHoffmann77 commented Mar 14, 2024

Uh oh!

ThomasHoffmann77 commented Mar 14, 2024

Uh oh!

ThomasHoffmann77 commented Mar 14, 2024

Uh oh!

branfosj commented Mar 14, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

branfosj commented Mar 15, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ThomasHoffmann77 commented Mar 15, 2024

Uh oh!

verdurin commented Mar 15, 2024

Uh oh!

verdurin commented Mar 15, 2024

Uh oh!

ThomasHoffmann77 commented Mar 15, 2024

Uh oh!

ThomasHoffmann77 commented Mar 19, 2024

Uh oh!

ThomasHoffmann77 commented Mar 19, 2024

Uh oh!

Flamefire commented Mar 22, 2024

Uh oh!

Flamefire commented Mar 22, 2024

Uh oh!

Flamefire commented Mar 22, 2024

Uh oh!

Uh oh!

Flamefire commented Mar 26, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

casparvl commented Jun 11, 2024

Uh oh!

akesandgren commented Jun 12, 2024

Uh oh!

ThomasHoffmann77 commented Jun 12, 2024

Uh oh!

akesandgren commented Jun 13, 2024

Uh oh!

VRehnberg commented Jun 13, 2024

Uh oh!

VRehnberg commented Jun 13, 2024

Uh oh!

VRehnberg commented Jun 18, 2024

Uh oh!

akesandgren commented Jun 26, 2024

Uh oh!

akesandgren left a comment

Choose a reason for hiding this comment

Uh oh!

ThomasHoffmann77 commented Jun 26, 2024

Uh oh!

Uh oh!

VRehnberg commented Jul 9, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Flamefire commented Jul 9, 2024

Uh oh!

VRehnberg commented Jul 9, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

VRehnberg commented Jul 9, 2024

Uh oh!

VRehnberg commented Jul 9, 2024

Uh oh!

Flamefire commented Jul 10, 2024

Uh oh!

VRehnberg commented Jul 10, 2024

ThomasHoffmann77 commented Mar 14, 2024 •

edited by casparvl

Loading

branfosj commented Mar 14, 2024 •

edited

Loading

branfosj commented Mar 15, 2024 •

edited

Loading

Flamefire commented Mar 26, 2024 •

edited

Loading

VRehnberg commented Jul 9, 2024 •

edited

Loading

VRehnberg commented Jul 9, 2024 •

edited

Loading