Parallelizing Two Concurrent Blocks By Combining Different Parallelization Strategies

Qikai_Lu: <p>Suppose that I have the following DDP model:</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
from torch.nn.parallel import DistributedDataParallel as DDP

class Foo(nn.Module):  # this is a demo, with dims only provided for context
    def __init__(self):
        super().__init__()  # Initialize the parent class
        self.block_a = nn.Linear(256, 32)
        self.block_b = nn.Sequential(
            nn.Linear(256, 128),
            nn.Linear(128, 64),
            nn.Linear(64, 32)
        )

    def forward(self, x):
        x_a = self.block_a(x)
        x_b = self.block_b(x)
        return x_a + x_b

# Assuming rank is defined and initialized elsewhere in your distributed setup
rank = 0  # Replace with the actual rank if needed
foo = Foo()
ddp_foo = DDP(foo, device_ids=[rank])
</code></pre>
<p>I have the following questions:</p>
<ol>
<li>Is there a way, so that <code>block_a</code> and <code>block_b</code> are concurrently executed in each DDP process? Thus, for each rank, the model computation is also parallelize along <code>x_a</code> and <code>x_b</code>.</li>
<li>Since <code>block_b</code> is large than <code>block_a</code>, I am considering running <code>block_b</code> (which has three layers) with FSDP while keeping <code>block_a</code> DDP. Is this doable, and how would I approach this?</li>
<li>This is not directly related to the model shown, but I am wondering if it is possible to mix data parallelism and tensor parallelism together? If so, are there any intuitive examples?</li>
</ol>
<p>Thanks.</p>

tianyu: <p>Using data parallel alone (FSDP or DDP), your model would be replicated on each device, and would always execute the whole computation. It is the data that is sharded across devices and being computed simultaneously.</p>
<p>In other words, do you really need to process the same data but with parallelized computations? Can you shard your data so that DDP alone would solve your problem?</p>

Qikai_Lu: <p>I suppose sharding my data does not really solve the issue. On further examination, I think this is more of a problem of tensor parallelism on block_a and block_b concurrently.</p>
<p>Let me re-frame the problem. Suppose that I have 8 GPUs, with 1 process launched per GPU. I want:</p>
<ul>
<li>process rank 0-1 to run tensor parallelism on block_a,</li>
<li>process rank 2-7 to run tensor parallelism on block_b,</li>
<li>block_a and block_b should be concurrently computed maximize overlap time.<br>
How would this be done?</li>
</ul>

tianyu: <p>So what you’d like to achieve is a special form of Model Parallel, not Data Parallel. AFAIK pytorch doesn’t have high level APIs for such use cases.</p>
<p>I believe you can write customized (autograd) functions to explicitly insert correct communication collectives to achieve it. E.g.</p>
<pre><code class="lang-auto">if current_rank == 0 or current_rank == 1:
  x_a = self.block_a(x)
  x_b = 0
else:
  x_a = 0
  x_b = self.block_b(x)

# run all-reduce on x_a and x_b

return x_a + x_b
</code></pre>
<p>Besides, using PyTorch DTensor would help you avoid writing backward functions on your own.</p>