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 获取更多信息。
Tensor 通过指定 gpu=True 来指定在 GPU 上创建。创建 Tensor 和 随机抽样 中的方法都支持这个参数。
gpu=True
>>> 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 上才会有真正的内存分配和计算。
.execute()
对于一个主存上创建的 tensor,调用 .to_gpu() 来指示把数据移至 GPU。
.to_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() 来指示把数据移到主存。
.to_cpu()
>>> c = b.to_cpu() # b is allocated on GPU, move back to main memory >>> c.sum().execute() # execution will happen on CPU
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
>>> df2 = df.to_cpu() # df is allocated on GPU, move back to main memory >>> df2.sum().execute() # execution will happen on CPU
本地基于多线程的调度器 能在单卡上正常运行。以上的例子都可以在单卡上工作。
对于 Mars tensor,可以利用单机多卡。
>>> import mars.tensor as mt >>> t = mt.random.rand(10000, 10000, gpu=True) >>> t.sum().execute()
如上代码会利用所有可见的显卡来计算。
如果要限制使用的显卡,可以设置环境变量 CUDA_VISIBLE_DEVICES。
CUDA_VISIBLE_DEVICES
CUDA_VISIBLE_DEVICES=0,3,5 ipython
这会让 ipython 只使用 0、3、5 显卡。因此在 ipython 中跑的 Mars tensor 任务只会使用这些显卡。
对于 Mars DataFrame,单机基于线程的调度器不能利用多卡来计算,在这种情况下,请使用分布式调度器。
Mars scheduler 和 web 的启动命令完全相同,参考 在集群中部署。
Mars worker 只能绑定到一张显卡上,因此如果你想使用多张显卡,启动和显卡个数一样多的 worker。
绑定到某个显卡的 worker 可以按如下命令启动:
mars-worker -a <worker_ip> -p <worker_port> -s <scheduler_ip>:<scheduler_port> --cuda-device 0
这个 worker 会绑定到显卡 0。
参考 启动 Worker 的其他命令 获取更多信息。
一旦一个集群创建好了,可以运行如下代码。
>>> import mars.tensor as mt >>> from mars.session import new_session >>> new_session('http://<web_ip>:<web_port>').as_default() >>> t = mt.random.rand(20, 20, gpu=True) >>> t.sum().execute() # run on workers which are bind to GPU