我可以将物理 GPU 拆分为多个逻辑/虚拟 GPU 并将它们传递给 dask_cuda.LocalCUDACluster 吗？

如何解决我可以将物理 GPU 拆分为多个逻辑/虚拟 GPU 并将它们传递给 dask_cuda.LocalCUDACluster 吗？

我的工作流极大地受益于 GPU 加速，但每个任务的内存要求相对较低（2-4 GB）。我正在使用 dask.dataframe、dask.distributed.Client 和 dask_cuda.LocalCUDACluster 的组合。该过程将从更多的 CUDA 工作器中受益，因此我想将物理 GPU（Nvidia RTX A600、V100、A100）拆分为多个虚拟/逻辑 GPU，以增加我的 dask_cuda LocalCUDACluster 中的工作器数量。我最初的想法是尝试将在 TensorFlow 中创建的 logical_gpus 传递到 LocalCUDACluster，但我似乎无法将它们传递到集群中。

我在 docker 环境中工作，我想将这些拆分保留在 python 中。理想情况下，此工作流程可以从本地工作站扩展到多节点 MPI 作业，但我不确定这是否可行，我愿意接受任何建议。

添加一个类似的例子。

from dask.distributed import Client
from dask_cuda import LocalCUDACluster
from dask_cuda.initialize import initialize
import pandas as pd
import dask.dataframe as dd
import time

# fake function
def my_gpu_sim(x):
   """ 
   GPU simulation which is independent of any others (calls a c++ program in real-world,which saves a 
   file.)
   """
   ...
   return None

# fake data creation 
dic = {'random':['apple' for i in range(40)],'main':[i for i in range(40)]}
df = pd.DataFrame.from_dict(dic)
ddf = dd.from_pandas(df,npartitions=4) 

 
# Configurations
protocol = "ucx"
enable_tcp_over_ucx = True
enable_nvlink = True
enable_infiniband = False
initialize(
    create_cuda_context=True,enable_tcp_over_ucx=enable_tcp_over_ucx,enable_infiniband=enable_infiniband,enable_nvlink=enable_nvlink,)
cluster = LocalCUDACluster(local_directory="/tmp/USERNAME",protocol=protocol,rmm_pool_size="35GB"
                    )
client = Client(cluster)

# Simulation
ddf.map_partitions(lambda df: df.apply(lambda x: my_gpu_sim(x.main),axis=1)).compute(scheduler=client)