Mars在GPU上#

Mars可以在NVIDIA GPU上运行。然而,不同模块需要额外的要求。

安装#

对于Mars张量,需要CuPy。假设您的CUDA驱动程序是10.1,请通过以下方式安装cupy:

pip install cupy-cuda101

有关更多信息,请参阅 install cupy

对于Mars DataFrame,RAPIDS cuDF是必需的。通过conda安装cuDF:

conda install -c rapidsai -c nvidia -c conda-forge \
 -c defaults cudf=0.13 python=3.7 cudatoolkit=10.1

有关更多信息,请参考 install cuDF

Mars 张量在 CUDA#

可以通过指定 gpu=True 在GPU上创建张量。 包含的方法在张量创建随机数据中提到。

>>> import mars.tensor as mt
>>> a = mt.random.rand(10, 10, gpu=True)  # indicate to create tensor on CUDA
>>> a.sum().execute()                     # execution will happen on CUDA

请记住,当创建张量时,尚未发生GPU内存分配。当 .execute() 被触发时,真正的内存分配和GPU上的计算将会发生。

对于主机内存上的张量,调用 .to_gpu() 来告诉 Mars 将数据移动到 GPU。

>>> b = mt.random.rand(10, 10)  # indicate to create on main memory
>>> b = b.to_gpu()              # indicate to move data to GPU memory
>>> b.sum().execute()

调用 .to_cpu() 告诉 Mars 将数据移动到主机内存。

>>> c = b.to_cpu()     # b is allocated on GPU, move back to main memory
>>> c.sum().execute()  # execution will happen on CPU

Mars 数据框在 CUDA#

Mars 可以直接将 CSV 文件读取到 GPU 中。

>>> import mars.dataframe as md
>>> df = md.read_csv('data.csv', gpu=True)  # indicates to read csv into GPU memory
>>> df.groupby('a').sum().execute()         # execution will happen on GPU

对于在主机内存上的DataFrame,调用 .to_gpu() 以告诉Mars将数据移动到GPU。

>>> import mars.tensor as mt
>>> import mars.dataframe as md
>>> df = md.DataFrame(mt.random.rand(10, 10))  # indicate to create on main memory
>>> df = df.to_gpu()                            # indicate to move data to GPU memory

调用 .to_cpu() 来告诉Mars将数据移动到主机内存。

>>> df2 = df.to_cpu()     # df is allocated on GPU, move back to main memory
>>> df2.sum().execute()     # execution will happen on CPU

多GPU#

对于Mars张量和DataFrame,可以在单台机器上使用多个GPU。

>>> import mars.tensor as mt
>>> t = mt.random.rand(10000, 10000, gpu=True)
>>> t.sum().execute()

上面的代码将尝试利用所有可见的GPU卡来执行计算。

如果您想将计算限制在某些GPU卡上,您可以设置环境变量 CUDA_VISIBLE_DEVICES

CUDA_VISIBLE_DEVICES=0,3,5 ipython

这将只限制ipython到GPU 0、3和5。因此在ipython中执行的所有Mars张量将仅在可见的GPU上运行。

分布式#

对于Mars监控者,启动命令是相同的。请参考 在集群上运行

对于Mars工作者,一个工作者可以绑定一个或多个GPU。

启动一个绑定到某个GPU的工作者的基本命令是:

mars-worker -H <worker_ip> -p <worker_port> -s <supervisor_ip>:<supervisor_port> --cuda-devices 0,1,2

启动的工作者将绑定到 GPU 0、1 和 2。

有关启动工作者的额外参数,请参阅 extra arguments for starting worker 获取更多信息。

一旦启动了Mars集群,您可以运行下面的代码。

>>> import mars
>>> import mars.tensor as mt
>>> new_session('http://<web_ip>:<web_port>')
>>> t = mt.random.rand(20, 20, gpu=True)
>>> t.sum().execute()  # run on workers which are bind to GPU