Repository: pyro-ppl/numpyro

Issue:
Multiple chains on Multiple GPU
Hi

I tried to run 4 chains with 4 GPU but I have received errors (it runs with 1 chain)

thanks

f


Compiling.. :   0%|                                                                                                                | 0/6000 [00:00<?, ?it/s2021-05-28 11:03:27.284436: E external/org_tensorflow/tensorflow/compiler/xla/status_macros.cc:56] Internal: RET_CHECK failure (external/org_tensorflow/tensorflow/compiler/xla/service/gpu/outfeed_thunk.cc:80) ShapeUtil::Equal(source_slices_[index].shape, output_shape) Mismatch between outfeed output buffer shape u32[2]{0} and outfeed source buffer shape s32[]                                                                                   | 0/6000 [00:00<?, ?it/s]
*** Begin stack trace ***

        clone
*** End stack trace ***

2021-05-28 11:03:27.284560: E external/org_tensorflow/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.cc:1881] Execution of replica 0 failed: Internal: RET_CHECK failure (external/org_tensorflow/tensorflow/compiler/xla/service/gpu/outfeed_thunk.cc:80) ShapeUtil::Equal(source_slices_[index].shape, output_shape) Mismatch between outfeed output buffer shape u32[2]{0} and outfeed source buffer shape s32[]
2021-05-28 11:03:27.375222: E external/org_tensorflow/tensorflow/compiler/xla/status_macros.cc:56] Internal: RET_CHECK failure (external/org_tensorflow/tensorflow/compiler/xla/service/gpu/outfeed_thunk.cc:80) ShapeUtil::Equal(source_slices_[index].shape, output_shape) Mismatch between outfeed output buffer shape u32[2]{0} and outfeed source buffer shape s32[]
*** Begin stack trace ***


        clone
*** End stack trace ***

2021-05-28 11:03:27.375336: E external/org_tensorflow/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.cc:1881] Execution of replica 1 failed: Internal: RET_CHECK failure (external/org_tensorflow/tensorflow/compiler/xla/service/gpu/outfeed_thunk.cc:80) ShapeUtil::Equal(source_slices_[index].shape, output_shape) Mismatch between outfeed output buffer shape u32[2]{0} and outfeed source buffer shape s32[]
2021-05-28 11:03:27.468534: E external/org_tensorflow/tensorflow/compiler/xla/status_macros.cc:56] Internal: RET_CHECK failure (external/org_tensorflow/tensorflow/compiler/xla/service/gpu/outfeed_thunk.cc:80) ShapeUtil::Equal(source_slices_[index].shape, output_shape) Mismatch between outfeed output buffer shape s32[] and outfeed source buffer shape u32[2]{0}
*** Begin stack trace ***



        clone
*** End stack trace ***

2021-05-28 11:03:27.468676: E external/org_tensorflow/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.cc:1881] Execution of replica 2 failed: Internal: RET_CHECK failure (external/org_tensorflow/tensorflow/compiler/xla/service/gpu/outfeed_thunk.cc:80) ShapeUtil::Equal(source_slices_[index].shape, output_shape) Mismatch between outfeed output buffer shape s32[] and outfeed source buffer shape u32[2]{0}
2021-05-28 11:03:27.559775: E external/org_tensorflow/tensorflow/compiler/xla/status_macros.cc:56] Internal: RET_CHECK failure (external/org_tensorflow/tensorflow/compiler/xla/service/gpu/outfeed_thunk.cc:80) ShapeUtil::Equal(source_slices_[index].shape, output_shape) Mismatch between outfeed output buffer shape s32[] and outfeed source buffer shape u32[2]{0}
*** Begin stack trace ***

        clone
*** End stack trace ***

2021-05-28 11:03:27.559911: E external/org_tensorflow/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.cc:1881] Execution of replica 3 failed: Internal: RET_CHECK failure (external/org_tensorflow/tensorflow/compiler/xla/service/gpu/outfeed_thunk.cc:80) ShapeUtil::Equal(source_slices_[index].shape, output_shape) Mismatch between outfeed output buffer shape s32[] and outfeed source buffer shape u32[2]{0}
Traceback (most recent call last):
  File "prova.py", line 121, in <module>
    m_vpi.run(random.PRNGKey(11), **dat)
  File "/usr/local/lib/python3.8/dist-packages/numpyro/infer/mcmc.py", line 504, in run
    states, last_state = pmap(partial_map_fn)(map_args)
jax._src.traceback_util.FilteredStackTrace: RuntimeError: Internal: RET_CHECK failure (external/org_tensorflow/tensorflow/compiler/xla/service/gpu/outfeed_thunk.cc:80) ShapeUtil::Equal(source_slices_[index].shape, output_shape) Mismatch between outfeed output buffer shape u32[2]{0} and outfeed source buffer shape s32[]: while running replica 0 and partition 0 of a replicated computation (other replicas may have failed as well).

The stack trace above excludes JAX-internal frames.
The following is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "prova.py", line 121, in <module>
    m_vpi.run(random.PRNGKey(11), **dat)
  File "/usr/local/lib/python3.8/dist-packages/numpyro/infer/mcmc.py", line 504, in run
    states, last_state = pmap(partial_map_fn)(map_args)
  File "/usr/local/lib/python3.8/dist-packages/jax/_src/traceback_util.py", line 139, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/jax/api.py", line 1582, in f_pmapped
    out = pxla.xla_pmap(
  File "/usr/local/lib/python3.8/dist-packages/jax/core.py", line 1453, in bind
    return call_bind(self, fun, *args, **params)
  File "/usr/local/lib/python3.8/dist-packages/jax/core.py", line 1385, in call_bind
    outs = primitive.process(top_trace, fun, tracers, params)
  File "/usr/local/lib/python3.8/dist-packages/jax/core.py", line 1456, in process
    return trace.process_map(self, fun, tracers, params)
  File "/usr/local/lib/python3.8/dist-packages/jax/core.py", line 625, in process_call
    return primitive.impl(f, *tracers, **params)
  File "/usr/local/lib/python3.8/dist-packages/jax/interpreters/pxla.py", line 621, in xla_pmap_impl
    return compiled_fun(*args)
  File "/usr/local/lib/python3.8/dist-packages/jax/interpreters/pxla.py", line 1168, in execute_replicated
    out_bufs = compiled.execute_sharded_on_local_devices(input_bufs)
RuntimeError: Internal: RET_CHECK failure (external/org_tensorflow/tensorflow/compiler/xla/service/gpu/outfeed_thunk.cc:80) ShapeUtil::Equal(source_slices_[index].shape, output_shape) Mismatch between outfeed output buffer shape u32[2]{0} and outfeed source buffer shape s32[]: while running replica 0 and partition 0 of a replicated computation (other replicas may have failed as well).



Discussion:
Hi @hyperfra, could you try to run `pmap` with some simple models and paste the code below? I'm not sure what leads to the issue but looking at the shape, it seems to relate to random PRNGKey.
Hi  fehiepsi

I experimented with pmap, the issue appears to be caused by the progress_bar (?) if I put progress_bar=False, the code runs on multiple GPU

f


Thanks, @hyperfra! It seems that host_callback for the progress bar is [not working yet](https://github.com/google/jax/issues/6447) on multi-GPU. I'll add a PR to disable the progress bar on GPU.

Resolving patch:
diff --git a/numpyro/util.py b/numpyro/util.py
--- a/numpyro/util.py
+++ b/numpyro/util.py
@@ -6,6 +6,7 @@
 import os
 import random
 import re
+import warnings
 
 import numpy as np
 import tqdm
@@ -300,6 +301,13 @@ def fori_collect(
     init_val_flat, unravel_fn = ravel_pytree(transform(init_val))
     start_idx = lower + (upper - lower) % thinning
     num_chains = progbar_opts.pop("num_chains", 1)
+    # host_callback does not work yet with multi-GPU platforms
+    # See: https://github.com/google/jax/issues/6447
+    if num_chains > 1 and jax.default_backend() == "gpu":
+        warnings.warn(
+            "We will disable progress bar because it does not work yet on multi-GPUs platforms."
+        )
+        progbar = False
 
     @cached_by(fori_collect, body_fun, transform)
     def _body_fn(i, vals):


Test patch: