FSDP Unit Comprise the head and tail modules into a single flat parameter

111357: <p>As I understand, the FSDP splits the model into several FSDP units. However, there appears to be some inefficiency in the parameters grouping strategy. It may comprise the parameters of both head and tail modules into a single flat parameter. To illustrate this, let’s consider the following tiny demo.</p>
<pre><code class="lang-python">
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout(0.25)
        self.dropout2 = nn.Dropout(0.5)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):

        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        output = F.log_softmax(x, dim=1)
        return output


  my_auto_wrap_policy = functools.partial(
      size_based_auto_wrap_policy, min_num_params=2000
  )
  model = Net().to(rank)
  model = FSDP(model, auto_wrap_policy=my_auto_wrap_policy)

# FullyShardedDataParallel(
#   (_fsdp_wrapped_module): Net(
#     (conv1): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1))
#     (conv2): FullyShardedDataParallel(
#       (_fsdp_wrapped_module): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1))
#     )
#     (dropout1): Dropout(p=0.25, inplace=False)
#     (dropout2): Dropout(p=0.5, inplace=False)
#     (fc1): FullyShardedDataParallel(
#       (_fsdp_wrapped_module): Linear(in_features=9216, out_features=128, bias=True)
#     )
#     (fc2): Linear(in_features=128, out_features=10, bias=True)
#   )
# )
</code></pre>
<p>The first FSDP unit combines the <code>conv1</code> and <code>fc2</code> layers in this demonstration. When we perform the forward pass for <code>conv1</code>, we only need the parameters of <code>conv1</code>, but the current setup fetches <code>fc2</code> parameters, even though they are not required.</p>
<p>Are there any configs to eliminate this unnecessary parameter fetch and improve efficiency?"</p>

H-Huang: <p>cc <a class="mention" href="/u/agu">@agu</a> for fsdp autowrap</p>

agu: <p>Thanks for asking! The FSDP API does have limitation around how it can be applied to group parameters together for communication, namely that it must be specified on a single <code>nn.Module</code>.</p>
<p>Auto wrapping (i.e. <code>auto_wrap_policy</code>) is a syntactic sugar to help users apply FSDP. The size-based policy can often lead to inefficient constructions like you have pointed out, so we do not recommend it.</p>
<p>In cases like this, if you want “optimal” grouping, then you may need to rewrite your <code>nn.Module</code> to construct dummy parent modules for the modules you want to group together and apply FSDP to those dummy parents.</p>

111357: <p>Hi <a class="mention" href="/u/agu">@agu</a>, thanks for your reply. As per your response, auto-wrapping is not recommended. Can you suggest an alternative?</p>
<p>Additionally, for the usage of <code>dummy module</code> and <code>dummy pattern</code>. Does it mean adding some dummy modules to adhere to size-based group rules? It would be helpful if you could provide some demonstrations or unit tests to clarify the concept.</p>
<p>Thanks a lot.</p>