Deep Learning Training - Multiprocessing

hpourmodheji: <p>Hi, Let us assume I have a simple feedforward neural network, sketched below. I want to train this model using PyTorch. How can I train it in a way that “L1a” and “L1b” get trained (forward and backward) simultaneously on two different GPUs. Any comment is very much appreciated.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/6/4/64bf01ed200e3d1495f5068866489357e359a3a6.png" data-download-href="https://discuss.pytorch.org/uploads/default/64bf01ed200e3d1495f5068866489357e359a3a6" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/6/4/64bf01ed200e3d1495f5068866489357e359a3a6_2_690x478.png" alt="image" data-base62-sha1="eneYGqo4oqEkuytj6n8WC15E74y" width="690" height="478" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/6/4/64bf01ed200e3d1495f5068866489357e359a3a6_2_690x478.png, https://discuss.pytorch.org/uploads/default/optimized/3X/6/4/64bf01ed200e3d1495f5068866489357e359a3a6_2_1035x717.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/6/4/64bf01ed200e3d1495f5068866489357e359a3a6.png 2x" data-dominant-color="F9F9F9"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1271×881 24.4 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

hpourmodheji: <p>To be more specific, I have the following code. How can I execute the “branch 0” lines and “branch 1” lines parallel on two GPUs? I have tried torch.multiprocessing, but it gives me the following error “RuntimeError: Cowardly refusing to serialize non-leaf tensor which requires_grad, since autograd does not support crossing process boundaries.”.</p>
<pre><code>class Feedforward(torch.nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(Feedforward, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.output_size = output_size
        self.fc1 = torch.nn.Linear(self.input_size, self.hidden_size).to("cuda:0")
        self.relu1 = torch.nn.ReLU()
        self.fc2a = torch.nn.Linear(self.hidden_size, self.output_size).to("cuda:0")
        self.relu2a = torch.nn.ReLU()
        self.fc2b = torch.nn.Linear(self.hidden_size, self.output_size).to("cuda:1")
        self.relu2b = torch.nn.ReLU()
        self.fc3 = torch.nn.Linear(self.output_size, 1).to("cuda:0")
        self.sigmoid = torch.nn.Sigmoid()

    def forward(self, x):
        hidden = self.fc1(x)
        relu = self.relu1(hidden)

        hidden_a = self.fc2a(relu) # branch 0
        relu_a = self.relu2a(hidden_a) # branch 0

        hidden_b = self.fc2b(relu.to("cuda:1")) # branch 1
        relu_b = self.relu2a(hidden_b) # branch 1

        output = self.fc3(relu_a + relu_b.to("cuda:0"))
        output = self.sigmoid(output)

        return output</code></pre>

pritamdamania87: <p>Is there a specific reason you want to use multiprocessing here? If you are not going to use multiple machines, you can just run this script in a single process and I believe it should achieve what you are looking for.</p>

hpourmodheji: <p><a class="mention" href="/u/pritamdamania87">@pritamdamania87</a>, Thank you so much for your comment. I have a single machine with two GPUs. Since L1a (branch 0) and L1b (branch 1) are independent of each other, I believe I can do the forward pass on both branches simultaneously on two GPUs. In this way, I may decrease training time. Can I do this using PyTorch? Any comments are very much appreciated.</p>

pritamdamania87: <aside class="quote no-group" data-username="hpourmodheji" data-post="4" data-topic="145331">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/h/ce73a5/48.png" class="avatar"> hpourmodheji:</div>
<blockquote>
<p>Can I do this using PyTorch? Any comments are very much appreciated.</p>
</blockquote>
</aside>
<p>Yes you can do this with PyTorch. Note that GPU operations are asynchronous in PyTorch and the code you shared above should already work and use both GPUs simultaneously.</p>
<p>When you run this code:</p>
<pre><code class="lang-auto">        hidden_a = self.fc2a(relu) # branch 0
        relu_a = self.relu2a(hidden_a) # branch 0
</code></pre>
<p>The work is just enqueued asynchronously to GPU 0 and PyTorch continues execution on the CPU side and will run this code without waiting for branch 0 to complete on GPU 0:</p>
<pre><code class="lang-auto">        hidden_b = self.fc2b(relu.to("cuda:1")) # branch 1
        relu_b = self.relu2a(hidden_b) # branch 1
</code></pre>
<p>Now, finally when this code is run:</p>
<pre><code class="lang-auto">        output = self.fc3(relu_a + relu_b.to("cuda:0"))
</code></pre>
<p>It automatically would trigger a device synchronization (due to .to(“cuda:0”)) which would in term wait for execution on GPU 0 and GPU 1 to finish before combining their output.</p>
<p>One small detail is that if <code>relu</code> is already on GPU 0, you probably might need to move <code>relu.to("cuda:1")</code> to a line before branch 0 and save it in a temporary variable. That would ensure when branch 1 executes, there is no dependency on branch 0.</p>

hpourmodheji: <p><a class="mention" href="/u/pritamdamania87">@pritamdamania87</a>, I ran the code and checked GPUs computations with PyTorch Profiler. As graphs below show, GPU1 computation starts right after GPU0 finishes its own computations. Am I do something wrong here?</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/7/d/7d77d3391488fc8e9baf2cd205d145f01e12c020.png" data-download-href="https://discuss.pytorch.org/uploads/default/7d77d3391488fc8e9baf2cd205d145f01e12c020" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/7/d/7d77d3391488fc8e9baf2cd205d145f01e12c020_2_690x373.png" alt="image" data-base62-sha1="hTWphStELKvmXlIzUykTi5aoe6Q" width="690" height="373" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/7/d/7d77d3391488fc8e9baf2cd205d145f01e12c020_2_690x373.png, https://discuss.pytorch.org/uploads/default/optimized/3X/7/d/7d77d3391488fc8e9baf2cd205d145f01e12c020_2_1035x559.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/7/d/7d77d3391488fc8e9baf2cd205d145f01e12c020.png 2x" data-dominant-color="FAFAFB"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1227×664 26.4 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/e/f/ef1eb0c79f16d67d9c52aa5cf83c4ee66cbc8b74.png" data-download-href="https://discuss.pytorch.org/uploads/default/ef1eb0c79f16d67d9c52aa5cf83c4ee66cbc8b74" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/e/f/ef1eb0c79f16d67d9c52aa5cf83c4ee66cbc8b74_2_689x268.png" alt="image" data-base62-sha1="y7lRUQTQmO2KWw0uLoVnjRXXnM0" width="689" height="268" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/e/f/ef1eb0c79f16d67d9c52aa5cf83c4ee66cbc8b74_2_689x268.png, https://discuss.pytorch.org/uploads/default/optimized/3X/e/f/ef1eb0c79f16d67d9c52aa5cf83c4ee66cbc8b74_2_1033x402.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/e/f/ef1eb0c79f16d67d9c52aa5cf83c4ee66cbc8b74.png 2x" data-dominant-color="FBFBFB"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1346×524 16 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Here is my full code:</p>
<pre><code class="lang-auto">class Feedforward(torch.nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(Feedforward, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.output_size = output_size
        self.fc1 = torch.nn.Linear(self.input_size, self.hidden_size).to("cuda:0")
        self.relu1 = torch.nn.ReLU()
        self.fc2a = torch.nn.Linear(self.hidden_size, self.output_size).to("cuda:0")
        self.relu2a = torch.nn.ReLU()
        self.fc2b = torch.nn.Linear(self.hidden_size, self.output_size).to("cuda:1")
        self.relu2b = torch.nn.ReLU()
        self.fc3 = torch.nn.Linear(self.output_size, 1).to("cuda:0")
        self.sigmoid = torch.nn.Sigmoid()

    def forward_branch0(self, relu):
        for i in range(10):
            hidden_a = self.fc2a(relu)
            relu = self.relu2a(hidden_a)
        return relu

    def forward_branch1(self, relu):
        for i in range(10):
            hidden_b = self.fc2b(relu)
            relu = self.relu2b(hidden_b)
        return relu

    def forward(self, x):
        hidden = self.fc1(x)
        relu = self.relu1(hidden)

        relu_b = relu.to("cuda:1")

        relu_a = self.forward_branch0(relu)
        relu_b = self.forward_branch1(relu_b)

        output = self.fc3(relu_a + relu_b.to("cuda:0"))
        output = self.sigmoid(output)

        return output


# CREATE RANDOM DATA POINTS
def blob_label(y, label, loc):  # assign labels
    target = np.copy(y)
    for l in loc:
        target[y == l] = label
    return target


with torch.profiler.profile(
        schedule=torch.profiler.schedule(
            wait=2,
            warmup=2,
            active=6,
            repeat=1),
        activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
        profile_memory=True,
        record_shapes=True,
        on_trace_ready=torch.profiler.tensorboard_trace_handler(dir_name='path-to-results'),
        with_stack=True
) as profiler:
    x_train, y_train = make_blobs(n_samples=1000, n_features=1000, cluster_std=1.5, shuffle=True)
    x_train = torch.cuda.FloatTensor(x_train)
    y_train = torch.FloatTensor(blob_label(y_train, 0, [0]))
    y_train = torch.cuda.FloatTensor(blob_label(y_train, 1, [1, 2, 3]))

    model = Feedforward(1000, 1024, 1024)

    criterion = torch.nn.BCELoss()

    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

    model.train()
    epoch = 20

    for epoch in range(epoch):
        optimizer.zero_grad()

        # Forward pass
        y_pred = model(x_train)

        # Compute Loss
        loss = criterion(y_pred.squeeze(), y_train)

        print('Epoch {}: train loss: {}'.format(epoch, loss.item()))

        # Backward pass
        loss.backward()  # computes the gradients

        optimizer.step()  # updates the weights accordingly

        profiler.step()
</code></pre>

pritamdamania87: <p>I used a much simpler version of your script:</p>
<pre><code class="lang-auto">import torch

class Feedforward(torch.nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(Feedforward, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.output_size = output_size
        self.fc1 = torch.nn.Linear(self.input_size, self.hidden_size).to("cuda:0")
        self.relu1 = torch.nn.ReLU()
        self.fc2a = torch.nn.Linear(self.hidden_size, self.output_size).to("cuda:0")
        self.relu2a = torch.nn.ReLU()
        self.fc2b = torch.nn.Linear(self.hidden_size, self.output_size).to("cuda:1")
        self.relu2b = torch.nn.ReLU()
        self.fc3 = torch.nn.Linear(self.output_size, 1).to("cuda:0")
        self.sigmoid = torch.nn.Sigmoid()

    def forward_branch0(self, relu):
        for i in range(10):
            hidden_a = self.fc2a(relu)
            relu = self.relu2a(hidden_a)
        return relu

    def forward_branch1(self, relu):
        for i in range(10):
            hidden_b = self.fc2b(relu)
            relu = self.relu2b(hidden_b)
        return relu

    def forward(self, x):
        hidden = self.fc1(x)
        relu = self.relu1(hidden)

        relu_b = relu.to("cuda:1")

        relu_a = self.forward_branch0(relu)
        relu_b = self.forward_branch1(relu_b)

        output = self.fc3(relu_a + relu_b.to("cuda:0"))
        output = self.sigmoid(output)

        return output


with torch.profiler.profile(
        schedule=torch.profiler.schedule(
            wait=2,
            warmup=2,
            active=6,
            repeat=1),
        activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA],
        profile_memory=True,
        record_shapes=True,
        with_stack=True
) as profiler:
    model = Feedforward(1000, 1024, 1024)
    inp = torch.rand(1000, 1000, device='cuda:0')
    for i in range(100):
        print(f'Starting iter: {i}')
        model(inp)
        profiler.step()

    profiler.export_chrome_trace('/tmp/chrome_trace')
</code></pre>
<p>When I look at the chrome trace it looks like both GPUs are being used concurrently:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/a/ba9934e00dc0d3fe2f772ec337a8d01ee518315e.png" data-download-href="https://discuss.pytorch.org/uploads/default/ba9934e00dc0d3fe2f772ec337a8d01ee518315e" title="Screen Shot 2022-03-04 at 3.10.54 PM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/b/a/ba9934e00dc0d3fe2f772ec337a8d01ee518315e_2_690x219.png" alt="Screen Shot 2022-03-04 at 3.10.54 PM" data-base62-sha1="qCJ3hLHQt8IUcKLhXmwIX38tHA2" width="690" height="219" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/b/a/ba9934e00dc0d3fe2f772ec337a8d01ee518315e_2_690x219.png, https://discuss.pytorch.org/uploads/default/optimized/3X/b/a/ba9934e00dc0d3fe2f772ec337a8d01ee518315e_2_1035x328.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/b/a/ba9934e00dc0d3fe2f772ec337a8d01ee518315e_2_1380x438.png 2x" data-dominant-color="D1DCD8"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screen Shot 2022-03-04 at 3.10.54 PM</span><span class="informations">1920×611 100 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>I’d say you probably want to look at GPU utilization/GPU trace to see if both GPUs are being used concurrently and not something like memory view.</p>

hpourmodheji: <p><a class="mention" href="/u/pritamdamania87">@pritamdamania87</a>, thanks so much for your valuable response. I was able to regenerate the results.</p>