Pipeline Parallelism performance with distributed-rpc on Jetson Nano devices

minh_nn: <p>Hi, i’m a student and currently working on pipelining models on edge devices.</p>
<p>My current setup is: 1 Master linux machine, 2 Jetson Nano as workers.<br>
I used the example on:<br>
<a href="https://pytorch.org/tutorials/intermediate/dist_pipeline_parallel_tutorial.html" class="onebox" target="_blank" rel="noopener nofollow ugc">https://pytorch.org/tutorials/intermediate/dist_pipeline_parallel_tutorial.html</a></p>
<p>But the difference is I trained the DenseNet(small) model with CIFAR10 dataset.<br>
I split the model as follows:</p>
<pre><code class="lang-auto">class Shard1(nn.Module):
    def __init__(self, *args, **kwargs):
        super(Shard1, self).__init__()

        block = Bottleneck
        nblocks = [6,12,24,16]
        growth_rate = 12
        reduction = 0.5
        num_classes = 10

        self.growth_rate = growth_rate
        self._lock = threading.Lock()
        self.device = "cuda:0" if torch.cuda.is_available() else "cpu"
        print(f"Using device: {torch.cuda.get_device_name(self.device)}")

        num_planes = 2*growth_rate

        self.conv1 = nn.Conv2d(3, num_planes, kernel_size=3, padding=1, bias=False).to(self.device)

        self.dense1 = self._make_dense_layers(block, num_planes, nblocks[0]).to(self.device)
        num_planes += nblocks[0]*growth_rate
        out_planes = int(math.floor(num_planes*reduction))
        self.trans1 = Transition(num_planes, out_planes).to(self.device)
        num_planes = out_planes

        self.dense2 = self._make_dense_layers(block, num_planes, nblocks[1]).to(self.device)
        num_planes += nblocks[1]*growth_rate
        out_planes = int(math.floor(num_planes*reduction))
        self.trans2 = Transition(num_planes, out_planes).to(self.device)
        
    def _make_dense_layers(self, block, in_planes, nblock):
        layers = []
        for i in range(nblock):
            layers.append(block(in_planes, self.growth_rate))
            in_planes += self.growth_rate
        return nn.Sequential(*layers)

    def forward(self, x_rref):
        x = x_rref.to_here().to(self.device)
        with self._lock:
            out = self.conv1(x)
            out = self.trans1(self.dense1(out))
            out = self.trans2(self.dense2(out))
        return out.cpu()

    def parameter_rrefs(self):
        r"""
        Create one RRef for each parameter in the given local module, and return a
        list of RRefs.
        """
        return [RRef(p) for p in self.parameters()]

class Shard2(nn.Module):
    def __init__(self, *args, **kwargs):
        super(Shard2, self).__init__()

        block = Bottleneck
        nblocks = [6,12,24,16]
        growth_rate = 12
        reduction = 0.5
        num_classes = 10

        self.growth_rate = growth_rate
        self._lock = threading.Lock()
        self.device = "cuda:0" if torch.cuda.is_available() else "cpu"
        print(f"Using device: {torch.cuda.get_device_name(self.device)}")

        num_planes = 8*growth_rate
    
        self.dense3 = self._make_dense_layers(block, num_planes, nblocks[2]).to(self.device)
        num_planes += nblocks[2]*growth_rate
        out_planes = int(math.floor(num_planes*reduction))
        self.trans3 = Transition(num_planes, out_planes).to(self.device)
        num_planes = out_planes

        self.dense4 = self._make_dense_layers(block, num_planes, nblocks[3]).to(self.device)
        num_planes += nblocks[3]*growth_rate

        self.bn = nn.BatchNorm2d(num_planes).to(self.device)
        self.linear = nn.Linear(num_planes, num_classes).to(self.device)

    def _make_dense_layers(self, block, in_planes, nblock):
        layers = []
        for i in range(nblock):
            layers.append(block(in_planes, self.growth_rate))
            in_planes += self.growth_rate
        return nn.Sequential(*layers)

    def forward(self, x_rref):
        x = x_rref.to_here().to(self.device)
        with self._lock:
            out = self.trans3(self.dense3(x))
            out = self.dense4(out)
            out = F.avg_pool2d(F.relu(self.bn(out)), 4)
            out = out.view(out.size(0), -1)
            out = self.linear(out)
        return out.cpu()

    def parameter_rrefs(self):
        r"""
        Create one RRef for each parameter in the given local module, and return a
        list of RRefs.
        """
        return [RRef(p) for p in self.parameters()]
</code></pre>
<p>I also trying to investigate the impact of split size (for parallel execution) on the performance.</p>
<p>Here is the plot of average GPU usage (recorded with <code>tegrastats</code>) on 2 Jetson devices:<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/2/4/243945eca0ddce16768cab08da8b3a01e74a68c3.png" alt="avg_gpu_all_split_size_cifar" data-base62-sha1="5arTywbNlgbZFL94ZixNVWjV4bx" width="640" height="480"></p>
<p>In <a href="https://pytorch.org/docs/stable/pipeline.html" class="inline-onebox" rel="noopener nofollow ugc">Pipeline Parallelism — PyTorch 2.1 documentation</a>, ideally, the smaller the split size, the higher the number of microbatches, thus should increase the GPU utilization. However, the figure above tell the opposite. Even for the 2nd Jetson Nano, the result is unexpected. So i suspect that there are definitely some kind of a bottleneck. One possible cause is the network bandwidth and another one is the CPU when we keep changing tensor’s device.</p>
<p>How can I address these problems?<br>
Thank you very much.</p>

minh_nn: <p>And here is the training time regarding to the split size:<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/0/7/0746d464cba35017c2067d2ce12d12e2e0c3e641.png" alt="total_training_time_cifar" data-base62-sha1="12n5PTH8ibQzRXK97bDGFe3c3Rf" width="640" height="480"></p>

Nishant_Madan: <p>Hi, I am a student and I am doing the same task for my master thesis and I am currently struggling to connect the edge devices with my Pytorch RPC module. Since, I have never done that before. I need a help. Do you have any idea to how to do that? It will be a great help for me. I have exactly the same setup.</p>
<p>1 Master linux machine, 2 Jetson Nano as workers.</p>