rapidsai / dask-cuda

Utilities for Dask and CUDA interactions
https://docs.rapids.ai/api/dask-cuda/stable/
Apache License 2.0
292 stars 93 forks source link

Setting `CUPY_ACCELERATORS` #972

Open jakirkham opened 2 years ago

jakirkham commented 2 years ago

Something we might want to explore is whether CUPY_ACCELERATORS should be set as part of cluster startup (either here or in Distributed). In particular this might be useful to pick up tensor optimizations for things like einsum and matmul.

cc @leofang (who may have thoughts on this) cc @jrbourbeau (in case you have any thoughts on whether this could live in Dask somewhere)

leofang commented 2 years ago

I would definitely suggest setting CUPY_ACCELERATORS=cub,cutensor and pinging at cuTENSOR 1.5.0. This should help a wide verity of elementwise and reduction operations.

jakirkham commented 2 years ago

What happens if cuTENSOR is not installed? Is there some kind of fallback? Asking as some users install dask-cuda from PyPI, so may not get cuTENSOR (unless they've installed it some other way)

leofang commented 2 years ago

CuPy provides a mechanism to check optional modules' availability and provide guidance to users:

>>> import cupy.cuda.cutensor
/home/leof/miniforge3/envs/opt_einsum_dev/lib/python3.8/site-packages/cupy/_environment.py:399: UserWarning: 
cutensor library could not be loaded.

Reason: ImportError (libcutensor.so.1: cannot open shared object file: No such file or directory)

You can install the library by:

  $ python -m cupyx.tools.install_library --library cutensor --cuda 11.4

  warnings.warn(msg)
>>> cupy.cuda.cutensor.available
False

So what happens is:

  1. if users explicitly use the cuTENSOR functionalities through CuPy, they'll be guided to install cuTENSOR using CuPy's automation script,
  2. if users request cutensor via CUPY_ACCELERATORS, this backend would just be skipped if available is False, and the fallback is based on the order of cutensor appearing in the string.

What CuPy has not done is to update the automation script to use the cuTENSOR wheel, though currently I don't encourage downstream other than cuQuantum to rely on it, as it's fat and not straightforward to adopt (ref). But it's there.

jakirkham commented 2 years ago

Thanks Leo! This is very helpful information 🙏

github-actions[bot] commented 2 years ago

This issue has been labeled inactive-30d due to no recent activity in the past 30 days. Please close this issue if no further response or action is needed. Otherwise, please respond with a comment indicating any updates or changes to the original issue and/or confirm this issue still needs to be addressed. This issue will be labeled inactive-90d if there is no activity in the next 60 days.