多卡环境下的资源确定

情况是一个节点上有多个 GPU,每个 GPU 运行一个 actor(指定 num_gpus=1),每个 actor 会周期性地从 object store 中加载新的模型参数,所以需要用到 model.to(device) 这样的接口。请问每个 actor 调用 model.to(device) 时如何保证参数加载到自己所在的 GPU 显存中呢?

如果你使用的是基于pytorch distributed的多卡通信方案,在actor启动时通过设置环境变量RANK等指定其使用的GPU编号即可

没有用这个,仅针对 ray 的话,有什么办法获得 actor 被调度到的 GPU 编号吗?

拿这个环境变量 os.environ[“CUDA_VISIBLE_DEVICES”]
https://docs.ray.io/en/latest/ray-core/tasks/using-ray-with-gpus.html?highlight=gpu#using-gpus-in-tasks-and-actors

ray.get_gpu_ids();

1 Like