GPU 指南#

Mars 可以利用 NVIDIA 显卡执行，但对于 Mars 里的不同模块，需要一些必要的依赖。

安装#

Mars tensor 依赖于 CuPy，假设你的 CUDA 驱动是 10.1，用如下方式安装 cupy：

pip install cupy-cuda101

参考安装 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

参考安装 cuDF 获取更多信息。

CUDA 上运行 Mars tensor#

Tensor 通过指定 gpu=True 来指定在 GPU 上创建。创建 Tensor 和随机抽样中的方法都支持这个参数。

>>> 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

记住，创建 tensor 的时候，并没有实际的 GPU 内存分配。当 .execute() 触发的时候，GPU 上才会有真正的内存分配和计算。

对于一个主存上创建的 tensor，调用 .to_gpu() 来指示把数据移至 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() 来指示把数据移到主存。

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

CUDA 上运行 Mars DataFrame#

Mars 可以直接读取 CSV 文件至显存。

>>> 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() 指示把数据移至显存。

>>> 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() 来指示把数据移到主存。

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

多卡#

对于 Mars tensor 和 DataFrame，可以利用单机多卡。

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

如上代码会利用所有可见的显卡来计算。

如果要限制使用的显卡，可以设置环境变量 CUDA_VISIBLE_DEVICES。

CUDA_VISIBLE_DEVICES=0,3,5 ipython

这会让 ipython 只使用 0、3、5 显卡。因此在 ipython 中跑的 Mars tensor 任务只会使用这些显卡。

分布式#

Mars Supervisor 启动命令完全相同，参考在集群中部署。

Mars worker 可以绑定到一个或者多张显卡。

绑定到某个显卡的 Worker 可以按如下命令启动：

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

这个 Worker 会绑定到显卡 0、1 和 2。

参考启动 Worker 的其他命令获取更多信息。

一旦一个集群创建好了，可以运行如下代码。

>>> 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