Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Error when donating buffer with under-specified sharding #25346

Open
jaro-sevcik opened this issue Dec 9, 2024 · 0 comments
Open

Error when donating buffer with under-specified sharding #25346

jaro-sevcik opened this issue Dec 9, 2024 · 0 comments
Assignees
Labels
bug Something isn't working

Comments

@jaro-sevcik
Copy link
Contributor

Description

When donating a buffer and the donation target has under-specified sharding, one can get an input-output aliasing mismatch error if XLA infers output-sharding for the donation target that is different from the donation source.

Here is an example of the error:

jaxlib.xla_extension.XlaRuntimeError: INTERNAL: Expected aliased input 0 at index {} and output at index {0} to have the same size. Input sub-shape is s32[16]{0} with size 64, output sub-shape is s32[8]{0} with size 32

Either the message should be better (recommend providing concrete sharding for everything), or JAX should ensure that the mismatch never happens (by passing the same fixed sharding for aliased input-output pairs to XLA, or perhaps XLA's sharding inference should respect aliasing in SPMD partitioning).

Repro:

import jax
from jax.sharding import Mesh
from jax.sharding import PartitionSpec as P

mesh = Mesh(jax.devices()[:2], ('x', ))
s = jax.sharding.NamedSharding(mesh, P('x'))

def f(x, y):
  x = x + 1
  x = jax.lax.with_sharding_constraint(x, s)
  return x, y

f = jax.jit(f, donate_argnums=(0,), in_shardings=(None, None), out_shardings=(None, s))

lowered = f.lower(jax.numpy.arange(16), jax.numpy.arange(16))
print(lowered.as_text())
print(lowered.compile().as_text())

System info (python version, jaxlib version, accelerator, etc.)

>>> import jax; jax.print_environment_info()
jax:    0.4.37.dev20241206+861115ad4
jaxlib: 0.4.37.dev20241209
numpy:  1.26.4
python: 3.12.3 (main, Nov  6 2024, 18:32:19) [GCC 13.2.0]
device info: NVIDIA A100X-2, 2 local devices"
process_count: 1
platform: uname_result(system='Linux', node='...', release='5.4.0-177-generic', version='#197-Ubuntu SMP Thu Mar 28 22:45:47 UTC 2024', machine='x86_64')


$ nvidia-smi
Mon Dec  9 07:39:31 2024
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 550.54.15              Driver Version: 550.54.15      CUDA Version: 12.6     |
|-----------------------------------------+------------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
|                                         |                        |               MIG M. |
|=========================================+========================+======================|
|   0  NVIDIA A100X                   On  |   00000000:06:00.0 Off |                    0 |
| N/A   38C    P0             70W /  300W |     428MiB /  81920MiB |      2%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   1  NVIDIA A100X                   On  |   00000000:46:00.0 Off |                    0 |
| N/A   38C    P0             70W /  300W |     424MiB /  81920MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+

+-----------------------------------------------------------------------------------------+
| Processes:                                                                              |
|  GPU   GI   CI        PID   Type   Process name                              GPU Memory |
|        ID   ID                                                               Usage      |
|=========================================================================================|
|    0   N/A  N/A   2549129      C   python                                        420MiB |
|    1   N/A  N/A   2549129      C   python                                        416MiB |
+-----------------------------------------------------------------------------------------+
@jaro-sevcik jaro-sevcik added the bug Something isn't working label Dec 9, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
bug Something isn't working
Projects
None yet
Development

No branches or pull requests

2 participants