UNIMPLEMENTED: Requested AllToAll not implemented on GPU;

[Cascol-Chen]

Here is the error log

Traceback (most recent call last):
  File "vmoe/main.py", line 71, in <module>
    app.run(main)
  File "/mnt/cephfs/home/cascol/anaconda3/envs/vmoe/lib/python3.7/site-packages/absl/app.py", line 303, in run
    _run_main(main, args)
  File "/mnt/cephfs/home/cascol/anaconda3/envs/vmoe/lib/python3.7/site-packages/absl/app.py", line 251, in _run_main
    sys.exit(main(argv))
  File "vmoe/main.py", line 64, in main
    trainer.train_and_evaluate(config=FLAGS.config, workdir=FLAGS.workdir)
  File "/mnt/cephfs/home/cascol/working-directory/MoE-in-MoE/vmoe/vmoe/train/trainer.py", line 492, in train_and_evaluate
    datasets['train'].element_spec['labels']))).compile()
  File "/mnt/cephfs/home/cascol/anaconda3/envs/vmoe/lib/python3.7/site-packages/jax/_src/api.py", line 521, in compile
    self._lowering.compile(), self.in_tree, self.out_tree,
  File "/mnt/cephfs/home/cascol/anaconda3/envs/vmoe/lib/python3.7/site-packages/jax/interpreters/pxla.py", line 2056, in compile
    _allow_compile_replicated=_allow_compile_replicated)  # type: ignore
  File "/mnt/cephfs/home/cascol/anaconda3/envs/vmoe/lib/python3.7/site-packages/jax/interpreters/pxla.py", line 2142, in from_hlo
    compiled = dispatch.compile_or_get_cached(backend, computation, compile_options)
  File "/mnt/cephfs/home/cascol/anaconda3/envs/vmoe/lib/python3.7/site-packages/jax/_src/dispatch.py", line 582, in compile_or_get_cached
    return backend_compile(backend, computation, compile_options)
  File "/mnt/cephfs/home/cascol/anaconda3/envs/vmoe/lib/python3.7/site-packages/jax/_src/profiler.py", line 206, in wrapper
    return func(*args, **kwargs)
  File "/mnt/cephfs/home/cascol/anaconda3/envs/vmoe/lib/python3.7/site-packages/jax/_src/dispatch.py", line 537, in backend_compile
    return backend.compile(built_c, compile_options=options)
RuntimeError: UNIMPLEMENTED: Requested AllToAll not implemented on GPU; replica_count: 1; partition_count: 8, group_mode: kCrossPartition, operand_count: 8; NCCL support: 1; first operand array element-type: BF16

what may cause the problem, and how may I solve it?
Please let me know if anything is needed.

My Enviroment

_libgcc_mutex             0.1                        main  
_openmp_mutex             4.5                       1_gnu  
absl-py                   0.12.0                   pypi_0    pypi
astroid                   2.9.0            py37h06a4308_0  
astunparse                1.6.3                    pypi_0    pypi
attrs                     21.4.0                   pypi_0    pypi
blas                      1.0                         mkl  
bzip2                     1.0.8                h7f98852_4    conda-forge
ca-certificates           2021.10.26           h06a4308_2  
cached-property           1.5.2                    pypi_0    pypi
cachetools                5.0.0                    pypi_0    pypi
certifi                   2021.10.8        py37h06a4308_2  
charset-normalizer        2.0.10                   pypi_0    pypi
chex                      0.1.0                    pypi_0    pypi
cloudpickle               2.0.0                    pypi_0    pypi
clu                       0.0.6                    pypi_0    pypi
contextlib2               21.6.0                   pypi_0    pypi
cudatoolkit               11.1.1               h6406543_8    conda-forge
cycler                    0.11.0                   pypi_0    pypi
decorator                 5.1.1                    pypi_0    pypi
dill                      0.3.4                    pypi_0    pypi
dm-tree                   0.1.6                    pypi_0    pypi
einops                    0.3.2                    pypi_0    pypi
ffmpeg                    4.3                  hf484d3e_0    pytorch
flatbuffers               2.0                      pypi_0    pypi
flax                      0.3.6                    pypi_0    pypi
fonttools                 4.28.5                   pypi_0    pypi
freetype                  2.10.4               h0708190_1    conda-forge
future                    0.18.2                   pypi_0    pypi
gast                      0.4.0                    pypi_0    pypi
gmp                       6.2.1                h58526e2_0    conda-forge
gnutls                    3.6.13               h85f3911_1    conda-forge
google-auth               1.6.3                    pypi_0    pypi
google-auth-oauthlib      0.4.6                    pypi_0    pypi
google-pasta              0.2.0                    pypi_0    pypi
googleapis-common-protos  1.54.0                   pypi_0    pypi
grpcio                    1.43.0                   pypi_0    pypi
h5py                      3.6.0                    pypi_0    pypi
idna                      3.3                      pypi_0    pypi
importlib-metadata        4.10.1                   pypi_0    pypi
importlib-resources       5.4.0                    pypi_0    pypi
intel-openmp              2021.4.0          h06a4308_3561  
isort                     5.9.3              pyhd3eb1b0_0  
jax                       0.2.28                   pypi_0    pypi
jaxlib                    0.1.75+cuda11.cudnn805          pypi_0    pypi
jpeg                      9b                   h024ee3a_2  
keras                     2.7.0                    pypi_0    pypi
keras-preprocessing       1.1.2                    pypi_0    pypi
kiwisolver                1.3.2                    pypi_0    pypi
lame                      3.100             h7f98852_1001    conda-forge
lazy-object-proxy         1.6.0            py37h27cfd23_0  
ld_impl_linux-64          2.35.1               h7274673_9  
libclang                  12.0.0                   pypi_0    pypi
libffi                    3.3                  he6710b0_2  
libgcc-ng                 9.3.0               h5101ec6_17  
libgomp                   9.3.0               h5101ec6_17  
libiconv                  1.16                 h516909a_0    conda-forge
libpng                    1.6.37               h21135ba_2    conda-forge
libstdcxx-ng              9.3.0               hd4cf53a_17  
libtiff                   4.0.9                he6b73bb_1    conda-forge
libuv                     1.42.0               h7f98852_0    conda-forge
markdown                  3.3.6                    pypi_0    pypi
matplotlib                3.5.1                    pypi_0    pypi
mccabe                    0.6.1                    py37_1  
mkl                       2020.2                      256  
mkl-service               2.3.0            py37h8f50634_2    conda-forge
mkl_fft                   1.2.0            py37h161383b_1    conda-forge
mkl_random                1.2.0            py37h9fdb41a_1    conda-forge
ml-collections            0.1.0                    pypi_0    pypi
msgpack                   1.0.3                    pypi_0    pypi
ncurses                   6.3                  h7f8727e_2  
nettle                    3.6                  he412f7d_0    conda-forge
ninja                     1.10.2               h4bd325d_0    conda-forge
numpy                     1.21.5                   pypi_0    pypi
numpy-base                1.19.2           py37hfa32c7d_0  
oauthlib                  3.1.1                    pypi_0    pypi
olefile                   0.46               pyh9f0ad1d_1    conda-forge
openh264                  2.1.1                h780b84a_0    conda-forge
openssl                   1.1.1m               h7f8727e_0  
opt-einsum                3.3.0                    pypi_0    pypi
optax                     0.1.0                    pypi_0    pypi
packaging                 21.3                     pypi_0    pypi
pandas                    1.3.5                    pypi_0    pypi
pillow                    9.0.0                    pypi_0    pypi
pip                       21.2.2           py37h06a4308_0  
platformdirs              2.4.0              pyhd3eb1b0_0  
promise                   2.3                      pypi_0    pypi
protobuf                  3.19.3                   pypi_0    pypi
pyasn1                    0.4.8                    pypi_0    pypi
pyasn1-modules            0.2.8                    pypi_0    pypi
pylint                    2.12.2           py37h06a4308_0  
pyparsing                 3.0.6                    pypi_0    pypi
pysocks                   1.7.1                    pypi_0    pypi
python                    3.7.11               h12debd9_0  
python-dateutil           2.8.2                    pypi_0    pypi
python_abi                3.7                     2_cp37m    conda-forge
pytorch                   1.8.0           py3.7_cuda11.1_cudnn8.0.5_0    pytorch
pytz                      2021.3                   pypi_0    pypi
pyyaml                    6.0                      pypi_0    pypi
readline                  8.1.2                h7f8727e_1  
requests                  2.27.1                   pypi_0    pypi
requests-oauthlib         1.3.0                    pypi_0    pypi
rsa                       4.8                      pypi_0    pypi
scipy                     1.7.3                    pypi_0    pypi
setuptools                58.0.4           py37h06a4308_0  
six                       1.16.0             pyh6c4a22f_0    conda-forge
socks                     0                        pypi_0    pypi
sqlite                    3.37.0               hc218d9a_0  
tensorboard               2.7.0                    pypi_0    pypi
tensorboard-data-server   0.6.1                    pypi_0    pypi
tensorboard-plugin-wit    1.8.1                    pypi_0    pypi
tensorflow                2.7.0                    pypi_0    pypi
tensorflow-cpu            2.7.0                    pypi_0    pypi
tensorflow-datasets       4.4.0                    pypi_0    pypi
tensorflow-estimator      2.7.0                    pypi_0    pypi
tensorflow-io-gcs-filesystem 0.23.1                   pypi_0    pypi
tensorflow-metadata       1.5.0                    pypi_0    pypi
tensorflow-probability    0.15.0                   pypi_0    pypi
termcolor                 1.1.0                    pypi_0    pypi
tfds-nightly              4.4.0.dev202201180108          pypi_0    pypi
tk                        8.6.11               h1ccaba5_0  
toml                      0.10.2             pyhd3eb1b0_0  
toolz                     0.11.2                   pypi_0    pypi
torchaudio                0.8.0                      py37    pytorch
torchvision               0.9.0                py37_cu111    pytorch
tqdm                      4.62.3                   pypi_0    pypi
typed-ast                 1.4.3            py37h7f8727e_1  
typing-extensions         3.10.0.2             hd3eb1b0_0  
typing_extensions         3.10.0.2           pyh06a4308_0  
urllib3                   1.26.8                   pypi_0    pypi
werkzeug                  2.0.2                    pypi_0    pypi
wheel                     0.37.1             pyhd3eb1b0_0  
wrapt                     1.13.3           py37h7f8727e_2  
xz                        5.2.5                h7b6447c_0  
zipp                      3.7.0                    pypi_0    pypi
zlib                      1.2.11               h7f8727e_4

[mattiasmar]

I encountered this with another workload and with a slightly different error message... Nevertheless, it seemed like the failing factor then was that my particular GPU card didn't support BF16. Check if your card does. "BF16" is mentioned in the very end of your error message.

[Cascol-Chen]

How shall I check if it support..? I have no idea on it.

[mattiasmar]

As for now I think only the Ampere architecture supports it.

[mattiasmar]

https://machine-learning-note.readthedocs.io/en/latest/blog/cheatsheets/Nvidia_GPU.html