33

据我了解，内置的 PyTorch 操作都通过隐式矢量化自动处理批处理，允许跨多个 GPU 进行并行处理。

但是，当根据文档在 CUDA 中编写自定义操作时，给出的 LLTM 示例执行批量不变的操作，例如逐元素计算 Sigmoid 函数的梯度。

但是，我有一个不是批处理元素不变且不可矢量化的用例。在单个 GPU 上运行，我目前（效率低下）循环遍历批处理中的每个元素，为每个元素执行内核启动，如下所示（在浏览器中编写，只是为了演示）：

std::vector<at::Tensor> op_cuda_forward(at::Tensor input, 
                                        at::Tensor elementSpecificParam) {

    auto output = at::zeros(torch::CUDA(/* TYPE */), {/* DIMENSIONS */});

    const size_t blockDim = //
    const size_t gridDim = //
    const size_t = numBatches = //

    for (size_t i = 0; i < numBatches; i++) {
        op_cuda_forward_kernel<T><<<gridDim, blockDim>>>(input[i],
                                                         elementSpecificParam[i], 
                                                         output[i]);
    }

    return {output};
}

但是，我希望通过批处理元素将此操作拆分到多个 GPU 上。

output在多 GPU 场景中，张量的分配如何工作？

当然，可以在启动适当的内核之前在每个 GPU 上创建中间张量，但是，将输入数据复制到每个 GPU 并再次返回的开销将是有问题的。

有没有一种更简单的方法来启动内核而无需先探测环境以获取 GPU 信息（# GPU 等）？

最终目标是让 CUDA 操作与torch.nn.DataParallel一起使用。

4

0 回答 0