Насколько я понимаю, встроенные операции pytorch автоматически обрабатывают партии посредством неявной векторизации, позволяя параллелизм между несколькими GPU.
Однако при написании пользовательской операции в CUDA в соответствии с Документацией приведенный пример LLTM выполняет операции, которые являются пакетными, например, вычисляя градиент функции Sigmoid по-разному.
Тем не менее, у меня есть вариант использования, который не является периодическим и не является векторизованным. Находясь на одном графическом процессоре, я в настоящее время (неэффективно) перебираю каждый элемент в пакете, выполняя запуск ядра для каждого, например (написанный в браузере, чтобы продемонстрировать):
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};
}
Однако я хочу разделить эту операцию на несколько графических процессоров с помощью пакетного элемента.
Как распределение output
тензора будет работать в сценарии с несколькими GPU?
Конечно, можно создать промежуточные тензоры на каждом графическом процессоре, прежде чем запускать соответствующее ядро, однако накладные расходы на копирование входных данных на каждый графический процессор и обратно будут проблематичными.
Есть ли более простой способ запуска ядра без предварительного тестирования среды для информации о GPU (# GPU и т.д.)?
Конечной целью является работа CUDA, которая работает с torch.nn.DataParallel.