Sub-modules in FSDP

Yu_I: <p>I am new to FSDP.</p>
<p>This is my understanding of the sub-module FSDP.<br>
FSDP relies on <code>torch.nn.Module.forward()</code> because of all-gather parameters. So, when I create sub modules with <code>torch.nn.Sequential()</code> in <code>torch.nn.Module</code> class, it does not work.<br>
Also, FSDP class says</p>
<blockquote>
<p>FSDP does not support running the forward pass of a submodule<br>
that is contained in an FSDP instance. This is because the<br>
submodule’s parameters will be sharded, but the submodule itself<br>
is not an FSDP instance, so its forward pass will not all-gather<br>
the full parameters appropriately.</p>
</blockquote>
<p>For example, the FSDP tutorial (Pytorch) demonstrates Deep-ViT, Cross-ViT, and Cait on FSDP.<br>
This tutorial wrapped them with <code>transformer_auto_wrapper_policy</code>, but it does not work because of sub-modules.<br>
(<a href="https://github.com/lessw2020/transformer_central/blob/main/transformer_wrapping_tutorial/transformer_wrapper_tutorial.ipynb" class="inline-onebox" rel="noopener nofollow ugc">transformer_central/transformer_wrapping_tutorial/transformer_wrapper_tutorial.ipynb at main · lessw2020/transformer_central · GitHub</a>)</p>
<pre><code class="lang-auto">class DeepViT(nn.Module):
  def __init__(self, *, image_size, ...):
          super().__init__()
          self.to_patch_embedding = nn.Sequential(
                      Rearrange('b c (h p1) (w p2) -&gt; b (h w) (p1 p2 c)', p1 = patch_size, p2 = patch_size),
                      nn.Linear(patch_dim, dim),
                  )  &lt;----- I think this does not work because of the sub-module.
</code></pre>
<p>In this case, we need to rewrite the sub-module (self.to_patch_embedding) with <code>torch.nn.module</code> like the following. (Assume add the following lines in Deep_ViT)</p>
<pre><code class="lang-auto">class PatchEmbedding(nn.Module):
    def __init__(self, in_channels: int = 3, patch_size: int = 16, emb_size: int = 768, img_size: int = 224):
        self.patch_size = patch_size
        super().__init__()
        self.projection = nn.Sequential(
            # using a conv layer instead of a linear one -&gt; performance gains
            nn.Conv2d(in_channels, emb_size, kernel_size=patch_size, stride=patch_size),
            Rearrange('b e (h) (w) -&gt; b (h w) e'),
        )
        self.cls_token = nn.Parameter(torch.randn(1, 1, emb_size))
        self.positions = nn.Parameter(torch.randn((img_size // patch_size) ** 2 + 1, emb_size))

    def forward(self, x: Tensor) -&gt; Tensor:
        b, _, _, _ = x.shape
        x = self.projection(x)
        cls_tokens = repeat(self.cls_token, '() n e -&gt; b n e', b=b)
        # prepend the cls token to the input
        x = torch.cat([cls_tokens, x], dim=1)
        # add position embedding
        x += self.positions
        return x

class DeepViT(nn.Module):
      def __init__(self, *, image_size, ...):
          super().__init__()
          self.to_patch_embedding = PatchEmbedding()
          ...

     def forward(self, img):
          x = self.to_patch_embedding(img)
          x = self.dropout(x)
          x = self.transformer(x)
          x = x.mean(dim = 1) if self.pool == 'mean' else x[:, 0]
          x = self.to_latent(x)
          return self.mlp_head(x)
</code></pre>
<p>Then, wrap with transformer_auto_wrap_policy.</p>
<pre><code class="lang-auto">from deep_vit import Residual, PatchEmbedding
import functools

transformer_auto_wrapper_policy = functools.partial(
        transformer_auto_wrap_policy,
        transformer_layer_cls={
            PatchEmbedding,
            Residual,  # &lt; ---- Your Transformer layer class
        },
    )

sharded_model = FSDP(
        model,
        auto_wrap_policy=transformer_auto_wrapper_policy,
        mixed_precision=mp_policy,
        sharding_strategy=model_sharding_strategy,
        device_id=torch.cuda.current_device(),  # streaming init
    )
</code></pre>
<p>Do the above steps collect?<br>
If so, FSDP shards PatchEmbedding across GPUs, right? (I’m not sure whether this is the right thing or not.)</p>

yf225: <p>cc. FSDP1 experts <a class="mention" href="/u/weifengpy">@weifengpy</a> <a class="mention" href="/u/agu">@agu</a></p>

agu: <p>I am not sure if I followed completely. For <code>self.to_patch_embedding</code>, if you wrap it with FSDP and run <code>self.to_patch_embedding(...)</code> (its forward), then it should work.</p>

Yu_I2: <p>I created a new account as I lost the previous password.</p>
<p>When I wrap (transformer wrap) the <code>self.to_patch_embedding</code> following the tutorial, I get this error:</p>
<blockquote>
<p>RuntimeError: Output 0 of ViewBackward0 is a view and its base or another view of its base has been modified inplace. This view<br>
is the output of a function that returns multiple views. Such functions do not allow the output views to be modified inplace.<br>
You should replace the inplace operation by an out-of-place one.</p>
</blockquote>
<p>The below codes are coming from the link (<a href="https://github.com/lessw2020/transformer_central/blob/main/transformer_wrapping_tutorial/transformer_wrapper_tutorial.ipynb" class="inline-onebox" rel="noopener nofollow ugc">transformer_central/transformer_wrapping_tutorial/transformer_wrapper_tutorial.ipynb at main · lessw2020/transformer_central · GitHub</a>)</p>
<pre><code class="lang-auto">class DeepViT(nn.Module):
  def __init__(self, *, image_size, ...):
          super().__init__()
          self.to_patch_embedding = nn.Sequential(
                      Rearrange('b c (h p1) (w p2) -&gt; b (h w) (p1 p2 c)', p1 = patch_size, p2 = patch_size),
                      nn.Linear(patch_dim, dim),
                  )
</code></pre>
<p>Then,</p>
<pre><code class="lang-auto">from deep_vit import Residual
import functools

transformer_auto_wrapper_policy = functools.partial(
        transformer_auto_wrap_policy,
        transformer_layer_cls={
            Residual,  # &lt; ---- Your Transformer layer class
        },
    )

sharded_model = FSDP(
        model,
        auto_wrap_policy=transformer_auto_wrapper_policy,
        mixed_precision=mp_policy,
        sharding_strategy=model_sharding_strategy,
        device_id=torch.cuda.current_device(),  # streaming init
    )
</code></pre>
<p>So, I think somehow FSDP cannot wrap sub-modules <img src="https://discuss.pytorch.org/images/emoji/apple/thinking.png?v=12" title=":thinking:" class="emoji" alt=":thinking:" loading="lazy" width="20" height="20"></p>

agu: <p>I am still not following exactly. How did you update the <code>auto_wrap_policy</code>? Did you include <code>nn.Sequential</code> as one of the classes?</p>

cobarnes_adobe: <p>I ran into the same error. I was able to resolve it by defining a new empty subclass of <code>nn.Sequential</code> e.g. <code>MySequential</code>, instantiating the submodule using that new subclass, and adding an <code>isinstance(module, MySequential)</code> check to my existing custom auto wrap policy such that the new class is wrapped (I did not use the <code>transformer_auto_wrapper_policy</code>).</p>
<p>I feel the PyTorch <a href="https://pytorch.org/docs/stable/fsdp.html" rel="noopener nofollow ugc">FSDP documentation</a> is misleading when it says: “FSDP does not support running the forward pass of a submodule that is contained in an FSDP instance.” Apparently, this sentence meant to say: “FSDP does not support running the forward pass of a submodule that is contained in an FSDP instance, unless the submodule is itself wrapped by FSDP.”</p>