Repository: pytorch
Issue: [libTorch] Model initialization on multi-device is slow. It seems to run sequentially in multi-thread
Labels: module: cpp,triaged

> Originally posted at https://discuss.pytorch.org/t/x/215093 

I am using libTorch for inference on multiple GPU devices. I use one-thread-per-device to initialize and then to run inference. Inference (i.e. `forward()` ) works fast as expected, however the initialization step seems to run sequentially. Once the initialization is complete, the rest of the code runs concurrently as expected. This is problematic for bigger models, where each thread takes several minutes. How to initialize models on multiple devices using libtorch?

Here is a minimal, reproducible example:

```cpp
#include <torch/torch.h>
#include <spdlog/spdlog.h>


using namespace torch;
namespace nn = torch::nn;
const torch::Device DEVICE = torch::Device(torch::cuda::is_available() ? torch::kCUDA : torch::kCPU); 

// a dummy model for demonstration
struct NetImpl : nn::Module {
    nn::Sequential layers;

    NetImpl(std::vector<int64_t> sizes, torch::Device device = DEVICE)
   : layers{ register_module("layers", torch::nn::Sequential()) } 
    {
        for (size_t i = 0; i < sizes.size() - 1; i++) {
            layers->push_back(nn::Linear(sizes[i], sizes[i + 1]));
            layers->push_back(nn::Functional(torch::relu));
        }
        this->to(device);
    }

    auto forward(Tensor x) -> Tensor {
        x = layers->forward(x);
        return x;
    }
};
TORCH_MODULE(Net);

struct Timer {

    std::string name;
    std::chrono::time_point<std::chrono::high_resolution_clock> start;

    Timer(std::string name="")
    : name {name}, start {std::chrono::high_resolution_clock::now()}
    {
        spdlog::info("Timer {} started", name);
    }

    double elapsed() {
        auto now = std::chrono::high_resolution_clock::now();
        return std::chrono::duration_cast<std::chrono::seconds>(now - start).count();
    }

    ~Timer() {
        spdlog::info("Timer {} ended: {:.3f}s", name, elapsed());
    }
};


int main() {
    spdlog::info("torch version {}", TORCH_VERSION);
    // deep network; FFN with a lot of layers to make it deep
    std::vector<int64_t> dims = { 
        1024, 4096, 8192, 16384, 8192, 4096, 1024, 512, 256, 512,
        1024, 4096, 8192, 16384, 8192, 4096, 1024, 512, 256, 512,
        1024, 4096, 8192, 16384, 8192, 4096, 1024, 512, 256, 512,
        1024, 4096, 8192, 16384, 8192, 4096, 1024, 512, 256, 512,
        1024, 4096, 8192, 16384, 8192, 4096, 1024, 512, 256, 512,
        };

    if (!torch::cuda::is_available()) {
        throw std::runtime_error("CUDA is not available");
    }
    std::vector<torch::Device> devices;
    for (auto i = 0; i < torch::cuda::device_count(); i++) {
        devices.push_back(torch::Device(torch::kCUDA, i));
    }
    { // scope for timer 
        int n_threads = devices.size();
        Timer timer(fmt::format("[{}-threaded initializer]", n_threads));
        std::vector<std::jthread> threads;
        for (int i = 0; i < n_threads; i++) {
            auto t = std::jthread([i, &dims, &devices] {
                auto device = devices[i];
                Timer timer(fmt::format("{}", device.str()));
                auto model = Net(dims, device);
            });
            threads.push_back(std::move(t));
        }
    }
    return 0;
}
```
With a single GPU, i.e. `CUDA_VISIBLE_DEVICES=0`

```
[250108 04:12:39|t1753841][info] Timer [1-threaded initializer] started
[250108 04:12:39|t1753854][info] Timer cuda:0 started
[250108 04:12:53|t1753854][info] Timer cuda:0 ended: 14.000s
[250108 04:12:53|t1753841][info] Timer [1-threaded initializer] ended: 14.000s
```
Now, with `CUDA_VISIBLE_DEVICES=0,1,` the time is almost doubled
```
[250108 04:13:02|t1754149][info] Timer [2-threaded initializer] started
[250108 04:13:02|t1754163][info] Timer cuda:0 started
[250108 04:13:02|t1754164][info] Timer cuda:1 started
[250108 04:13:26|t1754164][info] Timer cuda:1 ended: 24.000s
[250108 04:13:27|t1754163][info] Timer cuda:0 ended: 24.000s
[250108 04:13:27|t1754149][info] Timer [2-threaded initializer] ended: 24.000s
```
And with `CUDA_VISIBLE_DEVICES=0,1,2,3`, the pattern continues:
```
[250108 04:14:04|t1754791][info] Timer [4-threaded initializer] started
[250108 04:14:04|t1754795][info] Timer cuda:0 started
[250108 04:14:04|t1754796][info] Timer cuda:1 started
[250108 04:14:04|t1754797][info] Timer cuda:2 started
[250108 04:14:04|t1754798][info] Timer cuda:3 started
[250108 04:14:52|t1754796][info] Timer cuda:1 ended: 47.000s
[250108 04:14:52|t1754795][info] Timer cuda:0 ended: 48.000s
[250108 04:14:58|t1754797][info] Timer cuda:2 ended: 54.000s
[250108 04:14:58|t1754798][info] Timer cuda:3 ended: 54.000s
[250108 04:14:58|t1754791][info] Timer [4-threaded initializer] ended: 54.000s
```
Finally, with all 8 devices:
```
[250108 04:15:50|t1755936][info] Timer [8-threaded initializer] started
[250108 04:15:50|t1755959][info] Timer cuda:0 started
[250108 04:15:50|t1755960][info] Timer cuda:1 started
[250108 04:15:50|t1755961][info] Timer cuda:2 started
[250108 04:15:50|t1755962][info] Timer cuda:3 started
[250108 04:15:50|t1755963][info] Timer cuda:4 started
[250108 04:15:50|t1755964][info] Timer cuda:5 started
[250108 04:15:50|t1755965][info] Timer cuda:6 started
[250108 04:15:50|t1755966][info] Timer cuda:7 started
[250108 04:17:23|t1755960][info] Timer cuda:1 ended: 92.000s
[250108 04:17:23|t1755965][info] Timer cuda:6 ended: 93.000s
[250108 04:17:24|t1755964][info] Timer cuda:5 ended: 93.000s
[250108 04:17:24|t1755959][info] Timer cuda:0 ended: 94.000s
[250108 04:17:24|t1755963][info] Timer cuda:4 ended: 94.000s
[250108 04:17:25|t1755966][info] Timer cuda:7 ended: 94.000s
[250108 04:17:25|t1755961][info] Timer cuda:2 ended: 95.000s
[250108 04:17:28|t1755962][info] Timer cuda:3 ended: 97.000s
[250108 04:17:28|t1755936][info] Timer [8-threaded initializer] ended: 97.000s
```
I can’t see where in `NetImpl` or `nn::LinearImpl` the locking is enforcing sequential execution.
It looks like some internal API (ATen/C10) is at play and I am clueless how to resolve it. How to improve the parallelization in this case?

cc @jbschlosser