Model parallelism in Multi-GPUs: forward/backward graph

autograd
4

Yes, you can do “model parallelism” for instance with 2 gpus like this:

class Network(nn.Module):
    def __init__(self, split_gpus):
        self.module1 = (some layers)
        self.module2 = (some layers)

        self.split_gpus = split_gpus
        if self.split_gpus:
            self.module1.device("cuda:0")
            self.module2.device("cuda:1")

    def forward(self, x):
        x = self.module1(x)
        if self.split_gpus:
            x = x.device("cuda:1") 
        return self.module2(x)
3 Likes