Using `TensorParallel` with a mix of supported/non-supported layer ypes

username: Jacob_Platin
cooked: <p>Hi!</p>
<p>I am attempting to follow the <a href="https://pytorch.org/tutorials/intermediate/TP_tutorial.html" rel="noopener nofollow ugc">TensorParallel</a> tutortial, but my model, beyond containing the typical <code>nn.Linear</code> and <code>nn.Embedding</code> layers also contains <code>nn.Conv2d</code> layers.  In reading through the <code>TensorParallel</code> documentation, I didn’t see any mention of support for that layer type, so I am curious: will TensorParallel work if I don’t wrap all modules in my model via <code>parallelize_module</code>?</p>
<p>For example, my model is something like:</p>
<pre><code class="lang-auto">class Net(nn.Module):
    def __init__(self):
      super(Net, self).__init__()

      # First 2D convolutional layer, taking in 1 input channel (image),
      # outputting 32 convolutional features, with a square kernel size of 3
      self.conv1 = nn.Conv2d(1, 32, 3, 1)
      # Second 2D convolutional layer, taking in the 32 input layers,
      # outputting 64 convolutional features, with a square kernel size of 3
      self.conv2 = nn.Conv2d(32, 64, 3, 1)

      # Designed to ensure that adjacent pixels are either all 0s or all active
      # with an input probability
      self.dropout1 = nn.Dropout2d(0.25)
      self.dropout2 = nn.Dropout2d(0.5)

      # First fully connected layer
      self.fc1 = nn.Linear(9216, 128)
      # Second fully connected layer that outputs our 10 labels
      self.fc2 = nn.Linear(128, 10)

  # x represents our data
    def forward(self, x):
      # Pass data through conv1
      x = self.conv1(x)
      # Use the rectified-linear activation function over x
      x = F.relu(x)

      x = self.conv2(x)
      x = F.relu(x)

      # Run max pooling over x
      x = F.max_pool2d(x, 2)
      # Pass data through dropout1
      x = self.dropout1(x)
      # Flatten x with start_dim=1
      x = torch.flatten(x, 1)
      # Pass data through ``fc1``
      x = self.fc1(x)
      x = F.relu(x)
      x = self.dropout2(x)
      x = self.fc2(x)

      # Apply softmax to x
      output = F.log_softmax(x, dim=1)
      return output

</code></pre>
<p>Is it possible to only call <code>parallelize_module</code> on <code>fc1/fc2</code>?</p>
<p>Thanks!</p>

username: tianyu
cooked: <p>Thanks for the question!</p>
<p>You should be able to do that.<br>
Internally for parallelized module, the parameters would be represented by DTensor, instead of <code>torch.Tensor</code>. For each <code>ParallelStyle</code> (e.g. <code>ColwiseParallel</code>), there is the arg <code>use_local_output</code>. If it’s set to <code>True</code> (the default for most <code>ParallelStyle</code>), then the output from the parallelized module would be <code>torch.Tensor</code> and can work seamlessly with any non-parallelized module.</p>
<p>On a side note, for intermediate results to work with non-parallelized module, you need to make sure the sharding of activations out of a parallelized module is acceptable for the following non-parallelized module.</p>

username: Jacob_Platin
cooked: <p><a class="mention" href="/u/tianyu">@tianyu</a>  Thanks for the help! Any idea if that would work with FSDP and DDP as well? I tried using DDP, as follows, but ran into the following error:</p>
<p>DDP Code (using the model defined above):</p>
<pre><code class="lang-auto"># create a device mesh based on the given world_size.
_world_size = 2  # TODO

device_mesh = init_device_mesh(device_type="cuda", mesh_shape=(_world_size,))
_rank = device_mesh.get_rank()


print(f"Starting PyTorch TP example on rank {_rank}.")
assert _world_size % 2 == 0, f"TP examples require even number of GPUs, but got {_world_size} gpus"

# create model and move it to GPU - init"cuda"_mesh has already mapped GPU ids.
tp_model = Net().to("cuda")


# Create an optimizer for the parallelized module.
lr = 0.25
optimizer = torch.optim.AdamW(tp_model.parameters(), lr=lr, foreach=True)

# Custom parallelization plan for the model
tp_model = parallelize_module(
    module=tp_model,
    device_mesh=device_mesh,
    parallelize_plan={
        "fc1": ColwiseParallel(),
        "fc2": RowwiseParallel(),
    },
)
tp_model = DDP(tp_model, device_mesh=device_mesh)
</code></pre>
<p>Which throws error:</p>
<pre><code class="lang-auto">[rank0]: Traceback (most recent call last):
[rank0]:   File "/scratch/real-csd/test_tp.py", line 94, in &lt;module&gt;
[rank0]:     tp_model = DDP(tp_model, device_mesh=device_mesh)
[rank0]:   File "/opt/conda/envs/ptca/lib/python3.10/site-packages/torch/nn/parallel/distributed.py", line 784, in __init__
[rank0]:     self._log_and_throw(
[rank0]:   File "/opt/conda/envs/ptca/lib/python3.10/site-packages/torch/nn/parallel/distributed.py", line 1127, in _log_and_throw
[rank0]:     raise err_type(err_msg)
[rank0]: RuntimeError: Modules with uninitialized parameters can't be used with `DistributedDataParallel`. Run a dummy forward pass to correctly initialize the modules
[rank1]: Traceback (most recent call last):
[rank1]:   File "/scratch/real-csd/test_tp.py", line 94, in &lt;module&gt;
[rank1]:     tp_model = DDP(tp_model, device_mesh=device_mesh)
[rank1]:   File "/opt/conda/envs/ptca/lib/python3.10/site-packages/torch/nn/parallel/distributed.py", line 784, in __init__
[rank1]:     self._log_and_throw(
[rank1]:   File "/opt/conda/envs/ptca/lib/python3.10/site-packages/torch/nn/parallel/distributed.py", line 1127, in _log_and_throw
[rank1]:     raise err_type(err_msg)
[rank1]: RuntimeError: Modules with uninitialized parameters can't be used with `DistributedDataParallel`. Run a dummy forward pass to correctly initialize the modules

</code></pre>
<pre><code class="lang-auto"></code></pre>

username: tianyu
cooked: <p>As far as I know, DDP doesn’t work with other parallelisms yet.</p>
<p>FSDP can be composed with TP. You can do TP first and then do FSDP wrapping.</p>
<p>For a concrete example, maybe the pytorch/torchtitan repo would help. In particular, here’s how FSDP + TP works for the llama model.</p><aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/torchtitan/blob/main/torchtitan/parallelisms/parallelize_llama.py">
  <header class="source">

      <a href="https://github.com/pytorch/torchtitan/blob/main/torchtitan/parallelisms/parallelize_llama.py" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/torchtitan/blob/main/torchtitan/parallelisms/parallelize_llama.py" target="_blank" rel="noopener nofollow ugc">pytorch/torchtitan/blob/main/torchtitan/parallelisms/parallelize_llama.py</a></h4>


      <pre><code class="lang-py"># Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

# This file applies the PT-D parallelisms (except pipeline parallelism) and various
# training techniques (e.g. activation checkpointing and compile) to the Llama model.

from collections import defaultdict

import torch
import torch.nn as nn
from torch.distributed import DeviceMesh
from torch.distributed._composable.fsdp import fully_shard, MixedPrecisionPolicy
from torch.distributed._composable.replicate import replicate
from torch.distributed._tensor import Replicate, Shard
from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
    checkpoint_wrapper as ptd_checkpoint_wrapper,
)
</code></pre>



  This file has been truncated. <a href="https://github.com/pytorch/torchtitan/blob/main/torchtitan/parallelisms/parallelize_llama.py" target="_blank" rel="noopener nofollow ugc">show original</a>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>