| 11 Jul 2024 |
 SomeoneSerge (matrix works sometimes) | openai-triton broken with cuda+python3.12 😩 | 00:52:55 |
|  myrkskog joined the room. | 03:15:20 |
 connor (burnt/out) (UTC-8) | [2/2137/2141 built (157 failed), 11826 copied (948537.2 MiB), 27313.5 MiB DL] building xyce-7.8.0 (checkPhase): MEASURE/PrecisionTest............................................passed[sh] (Time: 1s = 0.
💀
| 07:06:43 |
| SomeoneSerge (matrix works sometimes) | ROCm stuff cache-missing again | 12:39:44 |
| SomeoneSerge (matrix works sometimes) | https://hydra.nixos.org/job/nixpkgs/trunk/rocmPackages.rocsolver.x86_64-linux yeah getting a timeout locally as well | 12:40:02 |
|  yorik.sar joined the room. | 15:05:02 |
| SomeoneSerge (matrix works sometimes) | Madoura I'm trying to build rocsolver again but I already suspect it's going to time out another time | 16:30:42 |
| SomeoneSerge (matrix works sometimes) | Madoura https://discourse.nixos.org/t/testing-gpu-compute-on-amd-apu-nixos/47060/2 this is falling apart 😹 | 23:16:45 |
| 12 Jul 2024 |
|  @valconius:matrix.org left the room. | 01:16:15 |
| 13 Jul 2024 |
 mcwitt | This might just be me being dumb, but am surprised that I'm unable to build jax with doCheck = false (use case is I want to override jaxlib = jaxlibWithCuda and don't want to run the tests). Repro:
nix build --impure --expr 'let nixpkgs = builtins.getFlake("github:nixos/nixpkgs/nixpkgs-unstable"); pkgs = import nixpkgs { system = "x86_64-linux"; config.allowUnfree = true; }; in pkgs.python3Packages.jax.overridePythonAttrs { doCheck = false; }'
fails with ModuleNotFoundError: jax requires jaxlib to be installed
| 01:58:54 |
| mcwitt | Something else that's puzzling me: overriding with jaxlib = jaxlibWithCuda doesn't seem to work for numpyro (which has a pretty simple derivation):
Sanity check: a Python env with just jax and jaxlibWithCuda is GPU-enabled, as expected:
nix shell --refresh --impure --expr 'let nixpkgs = builtins.getFlake "github:nixos/nixpkgs/nixpkgs-unstable"; pkgs = import nixpkgs { config.allowUnfree = true; }; in (pkgs.python3.withPackages (ps: [ ps.jax ps.jaxlibWithCuda ]))' --command python -c "import jax; print(jax.devices())"
yields [cuda(id=0), cuda(id=1)].
But when numpyro is overridden to use jaxlibWithCuda, for some reason the propagated jaxlib is still the CPU version:
nix shell --refresh --impure --expr 'let nixpkgs = builtins.getFlake "github:nixos/nixpkgs/nixpkgs-unstable"; pkgs = import nixpkgs { config.allowUnfree = true; }; in (pkgs.python3.withPackages (ps: [ ((ps.numpyro.overridePythonAttrs (_: { doCheck = false; })).override { jaxlib = ps.jaxlibWithCuda; }) ]))' --command python -c "import jax; print(jax.devices())"
yields [CpuDevice(id=0)]. (Furthermore, if we try to add jaxlibWithCuda to the withPackages call, we get a collision error, so clearly something is propagating the CPU jaxlib 🤔
Has anyone seen this, or have a better way to use the GPU-enabled jaxlib as a dependency?
| 04:30:05 |
| mcwitt | * Something else that's puzzling me: overriding with jaxlib = jaxlibWithCuda doesn't seem to work for numpyro (which has a pretty simple derivation):
Sanity check: a Python env with just jax and jaxlibWithCuda is GPU-enabled, as expected:
nix shell --refresh --impure --expr 'let nixpkgs = builtins.getFlake "github:nixos/nixpkgs/nixpkgs-unstable"; pkgs = import nixpkgs { config.allowUnfree = true; }; in (pkgs.python3.withPackages (ps: [ ps.jax ps.jaxlibWithCuda ]))' --command python -c "import jax; print(jax.devices())"
yields [cuda(id=0), cuda(id=1)].
But when numpyro is overridden to use jaxlibWithCuda, for some reason the propagated jaxlib is still the CPU version:
nix shell --refresh --impure --expr 'let nixpkgs = builtins.getFlake "github:nixos/nixpkgs/nixpkgs-unstable"; pkgs = import nixpkgs { config.allowUnfree = true; }; in (pkgs.python3.withPackages (ps: [ ((ps.numpyro.overridePythonAttrs (_: { doCheck = false; })).override { jaxlib = ps.jaxlibWithCuda; }) ]))' --command python -c "import jax; print(jax.devices())"
yields [CpuDevice(id=0)]. (Furthermore, if we try to add jaxlibWithCuda to the withPackages call, we get a collision error, so clearly something is propagating the CPU jaxlib 🤔)
Has anyone seen this, or have a better way to use the GPU-enabled jaxlib as a dependency?
| 04:30:27 |
| mcwitt | Ack, the second one was just me being dumb. I'd mixed up the proper ordering of `override` and `overridePythonAttrs` 🤦 sorry for the noise | 05:03:40 |
| SomeoneSerge (matrix works sometimes) | In reply to @mcwitt:matrix.org
Something else that's puzzling me: overriding with jaxlib = jaxlibWithCuda doesn't seem to work for numpyro (which has a pretty simple derivation):
Sanity check: a Python env with just jax and jaxlibWithCuda is GPU-enabled, as expected:
nix shell --refresh --impure --expr 'let nixpkgs = builtins.getFlake "github:nixos/nixpkgs/nixpkgs-unstable"; pkgs = import nixpkgs { config.allowUnfree = true; }; in (pkgs.python3.withPackages (ps: [ ps.jax ps.jaxlibWithCuda ]))' --command python -c "import jax; print(jax.devices())"
yields [cuda(id=0), cuda(id=1)].
But when numpyro is overridden to use jaxlibWithCuda, for some reason the propagated jaxlib is still the CPU version:
nix shell --refresh --impure --expr 'let nixpkgs = builtins.getFlake "github:nixos/nixpkgs/nixpkgs-unstable"; pkgs = import nixpkgs { config.allowUnfree = true; }; in (pkgs.python3.withPackages (ps: [ ((ps.numpyro.overridePythonAttrs (_: { doCheck = false; })).override { jaxlib = ps.jaxlibWithCuda; }) ]))' --command python -c "import jax; print(jax.devices())"
yields [CpuDevice(id=0)]. (Furthermore, if we try to add jaxlibWithCuda to the withPackages call, we get a collision error, so clearly something is propagating the CPU jaxlib 🤔)
Has anyone seen this, or have a better way to use the GPU-enabled jaxlib as a dependency?
Wait why is numpyro propagating jaxlib, I thought we had a convention not to propagate jaxlib | 08:10:45 |
