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