Tensor Parallel Single Layer

CJog: <p>I would like to train a model which has a large number of classes, making the linear layer too large to fit on a single gpu. Using <a href="https://pytorch.org/docs/stable/distributed.tensor.parallel.html" rel="noopener nofollow ugc">tensor parallel</a>, how can I parallelize just the linear layer while keeping the rest of the network on each gpu like in distributed data parallel?</p>
<p>The model structure as shown below gives an idea of what I want to achieve.</p>
<pre><code class="lang-auto">import torch.nn as nn
import torch
import torch.distributed as dist


class DummyModel(nn.Module):
    def __init__(self):
        super(DummyModel, self).__init__()

        # some conv2d layer and relu
        self.conv = nn.Conv2d()
        self.relu = nn.ReLU()

        # Large linear layer
        self.linear = nn.Linear(in_channels, 1000000)

    def forward(self, x):
        x = self.conv(x)
        x = self.relu(x)
        x = torch.flatten(x, 1)
        
        # Want to parallelize
        x = self.linear(x)
        return x

# Example train function
def train(model, device, rank, world_size):
    # Initialize distributed training
    dist.init_process_group(backend='nccl')

    # Set device for the model
    model.to(device)

    # Wrap the model with DistributedDataParallel
    model = nn.parallel.DistributedDataParallel(model, device_ids=[device])

    # Define loss function and optimizer
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=0.01)

    # Dummy input and target
    input_data = torch.randn(64, 3, 32, 32).to(device)
    target = torch.randint(0, 1000000, (64,)).to(device)

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

        # How to tensor parallelize linear layer but make sure output is gathered and synced on all gpus?
        output = model(input_data)

        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        print(f"Epoch {epoch+1}: Loss={loss.item()}")

    # Clean up distributed training
    dist.destroy_process_group()
</code></pre>

J_Johnson: <p>Welcome to the forums! Saw your question a few days ago but was too preoccupied to answer, but made a mental note to check back if anyone replied.</p>
<p>Let’s just start by saying there is no clean and simple one-line code solution to your issue. However, it is doable, with a bit of understanding of both Pytorch and how linear algebra operations work.</p>
<p>First off, the vanilla Linear layer is not divisible, as each tensor must be assigned to a device.</p>
<p>So we will need to use something a little more basic than the Linear layer and effectively reconstruct it, but split it as many ways as desired.</p>
<p>Please see the code below with comments, and a demonstration that the two mathematical operations are identical, if taking into account rounding errors.</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.nn.functional as F

# demonstrate an equivalent linear algebra operation with split weights
big_n = 1000000
# Create a control layer
big_layer = nn.Linear(big_n, 10)

input_data = torch.rand((12, big_n))

expected_output=big_layer(input_data)
print("Control layer output size", expected_output.size())

#test layers, assign weights to separate learnable layers, these can be assigned to different devices, can also assign a bias vector
#this would be in your __init__
layer_a_weight = nn.Parameter(big_layer.weight.data[:,:big_n//2])
layer_b_weight = nn.Parameter(big_layer.weight.data[:,big_n//2:])
layer_bias = nn.Parameter(big_layer.bias.data)
print("Split weight sizes", layer_a_weight.size(), layer_b_weight.size())

#split data
#the rest would be in your forward pass, you'd need to assign data to the respective devices you put your learnable parameters on
data_a = input_data[:,:big_n//2]
data_b = input_data[:,big_n//2:]
print("Split data sizes", data_a.size(), data_b.size())


#calculate separate matmuls
data_a = data_a@layer_a_weight.T
data_b = data_b@layer_b_weight.T
print("Intermediate split sizes", data_a.size(), data_b.size())

#add both elementwise, you can also add the bias at this point, if you are using it
data_test = data_a + data_b + layer_bias
print("Final output size", data_test.size())

#check if equivalent to expected_output
with torch.no_grad():
    print("Mean squared distance between both", F.mse_loss(data_test, expected_output))
</code></pre>
<blockquote>
<p>Control layer output size torch.Size([12, 10])<br>
Split weight sizes torch.Size([10, 500000]) torch.Size([10, 500000])<br>
Split data sizes torch.Size([12, 500000]) torch.Size([12, 500000])<br>
Intermediate split sizes torch.Size([12, 10]) torch.Size([12, 10])<br>
Final output size torch.Size([12, 10])<br>
Mean squared distance between both tensor(3.5189e-13)</p>
</blockquote>
<p>As can be seen, the mean squared distance between both is tiny and just due to rounding error.</p>
<p>Hope this helps you accomplish what you’re attempting.</p>

CJog: <p>Thank you! I thought the tensor parallel module will do exactly this under the hood. My understanding of that module seems incorrect. Can you explain what that module does exactly?</p>

J_Johnson: <p>You can read it here:</p>
<p><a href="https://pytorch.org/docs/stable/generated/torch.nn.DataParallel.html" class="onebox" target="_blank" rel="noopener nofollow ugc">https://pytorch.org/docs/stable/generated/torch.nn.DataParallel.html</a></p>
<p>Basically, DataParallel makes a copy of the entire model on each device. And then divides the batch size up to different devices, and combines the gradients of all copies before back propagation.</p>

CJog: <p>Yes, but I assume <a href="https://pytorch.org/docs/stable/distributed.tensor.parallel.html" rel="noopener nofollow ugc">TensorParallel</a> is different than DataParallel?</p>

J_Johnson: <p>I see. That must be new. I’m not clear how that works, yet. But will definitely check it out when I get a chance.</p>

CJog: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> can you shed some light on this or tag someone who can?</p>