Comments (4)
MasterSkepticista
commented on August 15, 2024
1
@vfdev-5 Your suggestion worked.
@sergachev I observed that JAX was built against openxla/xla@95e3eea, which was before the revert
from jax.
vfdev-5
commented on August 15, 2024
@MasterSkepticista the error is related with getting cuda:gemm_fusion_autotuning_results on shards and maybe related to openxla/xla#13108 (cc @sergachev). To disable the autotuning and to make your MWE work, you could try to run it with:
XLA_FLAGS=--xla_gpu_shard_autotuning=false mpirun -n 8 python error.pyLet me know if this workaround helps
from jax.
sergachev
commented on August 15, 2024
openxla/xla#13108 was reverted.
--xla_gpu_shard_autotuning=false disables sharding of autotuning, not the autotuning itself.
from jax.
sergachev
commented on August 15, 2024
I can reproduce with jax==0.4.31 and --xla_gpu_shard_autotuning=false helps - looks like openxla/xla#13108 got into this JAX release before it got reverted. Thank you for cc'ing me, I'll investigate why does it fail.
from jax.
Related Issues (20)
- Numerical precision differences in jitted vs. non-jitted code! HOT 3
- JAX indexing should support zero-dimensional boolean masks
- Memory profiler missing in trace HOT 6
- Getting some problems in the usage of jnp.linalg.inv HOT 3
- Installing jaxlib with CUDA Version: 11.7 HOT 5
- AssertionError raised due to failing assertion in partial_eval.py HOT 4
- Extending JAX with CUDA: cstdint: error: the global scope has no "..." HOT 3
- Complex gradient HOT 3
- Reconstructing a custom registered dataclass seems to give a segmentation fault
- AttributeError: module 'jax.random' has no attribute 'KeyArray' HOT 5
- Accessing incoming gradient and activations in custom_vjp HOT 4
- Finite precision error calculations always 0 under JIT with bfloat16 HOT 3
- Errors from interaction between `shard_map` with `auto` and `lax.map`
- Document Named Collectives HOT 1
- local_device_count fails on wsl with 'cpu' HOT 4
- Scatter on Sharded Matrices has bugs HOT 1
- Port all custom calls to FFI
- custom_jvp leaks tracers if they're marked as nondiff_argnum HOT 2
- Add support nogil mode in JAX for Python 3.13
Recommend Projects
-
React
A declarative, efficient, and flexible JavaScript library for building user interfaces.
-
Vue.js
🖖 Vue.js is a progressive, incrementally-adoptable JavaScript framework for building UI on the web.
-
Typescript
TypeScript is a superset of JavaScript that compiles to clean JavaScript output.
-
TensorFlow
An Open Source Machine Learning Framework for Everyone
-
Django
The Web framework for perfectionists with deadlines.
-
Laravel
A PHP framework for web artisans
-
D3
Bring data to life with SVG, Canvas and HTML. 📊📈🎉
-
Recommend Topics
-
javascript
JavaScript (JS) is a lightweight interpreted programming language with first-class functions.
-
web
Some thing interesting about web. New door for the world.
-
server
A server is a program made to process requests and deliver data to clients.
-
Machine learning
Machine learning is a way of modeling and interpreting data that allows a piece of software to respond intelligently.
-
Visualization
Some thing interesting about visualization, use data art
-
Game
Some thing interesting about game, make everyone happy.
Recommend Org
-
Facebook
We are working to build community through open source technology. NB: members must have two-factor auth.
-
Microsoft
Open source projects and samples from Microsoft.
-
Google
Google ❤️ Open Source for everyone.
-
Alibaba
Alibaba Open Source for everyone
-
D3
Data-Driven Documents codes.
-
Tencent
China tencent open source team.
from jax.