FullyShardedDataParallel model is not wrapped in FSDP in custom functions

f_g: <p>I’m trying to play with FSDP, however it seems like after wrapping, the model’s custom methods wouldn’t think that the model is wrapped. This is the minimally reproducible snippet:</p>
<pre><code class="lang-auto">from torch import Tensor, nn
import torch.distributed as dist
import os
import torch
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

os.environ['MASTER_ADDR'] = 'localhost'
os.environ['MASTER_PORT'] = '12355'

# initialize the process group
dist_rank = int(os.environ.get("RANK", 0))
world_size = int(os.environ.get("WORLD_SIZE", torch.cuda.device_count()))
dist.init_process_group("nccl", rank=dist_rank, world_size=world_size)

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.layer = nn.Linear(3,3)
    def foo(self):
        print(list(self.parameters()))
        return 1

model = MyModel()
model2 = FSDP(model)
model2.foo()
</code></pre>
<p>here, even though <code>model2.parameters()</code> still has the full parameters, within <code>foo</code> it only printed out an empty list.</p>

rvarm1: <p>I think this is expected, since you’re printing within <code>foo</code> that belongs to <code>MyModel</code> class?</p>

f_g: <p>hmm ok, then I think my question would be on how to access the parameters. I found that before wrapping with FSDP, I could do <code>model.parameters()</code> no problem. After wrapping, <code>model.parameters()</code> would be an empty list, but the FSDP wrapped <code>model2.parameters()</code> still has the <code>FlatParameter</code> content. Is this also excepted? This would make accessing MyModel’s parameters inside <code>foo</code> to empty.</p>
<p>Or should I be accessing parameters differently?<br>
(I edited the snippet in the original question a bit to reflect how I accessed the parameters inside foo)</p>

agu: <ol>
<li>Could you explain more what your desired use case is for accessing the model’s parameters from a custom function? What is the desired behavior?</li>
<li>What version of PyTorch are you using? (e.g. v1.13 or the latest nightly)</li>
</ol>

f_g: <ol>
<li>our models have some custom functions that gets the optimizer for example. But if this is desired behavior (i.e we’re not supposed to access parameters within the model), then I can work around it. Although I do think it’s a little surprising behavior, where after wrapping in FSDP, <code>model.parameters()</code> would turn empty? Also why would that be the case?</li>
<li>v 1.13.0</li>
</ol>

agu: <ol>
<li>Sorry, I did not quite follow. Could you explain in more detail how you are using <code>.parameters()</code> to get the optimizer?</li>
<li>After v1.13, we removed <code>FlattenParamsWrapper</code>. This means that the <code>FlatParameter</code>s are registered to the wrapped module (e.g. the <code>MyModel</code> instance in your example). In that case, if you call <code>model.foo()</code> after <code>model2 = FSDP(model)</code>, you should see the <code>FlatParameter</code>s:</li>
</ol>
<pre><code class="lang-auto">[Parameter containing:
Parameter(FlatParameter([-0.4381, -0.5244, -0.5647, -0.5283, -0.1153, -0.5181],
              requires_grad=True))]
[Parameter containing:
Parameter(FlatParameter([ 0.4810, -0.3448, -0.3134,  0.2578, -0.2005, -0.0755],
              requires_grad=True))]
</code></pre>
<p>However, I ask for more clarification on item 1 because I want to see if there is a more recommended way to achieve what you want. Depending on the fact that the <code>FlatParameter</code>s are registered to the wrapped module might not be robust.</p>

f_g: <ol>
<li>I’d call <code>optim.SGD(self.parameters())</code> for example, and the parameters there is empty.</li>
<li>I’m on <code>1.13.1+cu117</code>, does it have the change? or do you mean it’s only available in PyTorch 2.0?</li>
</ol>

agu: <ol>
<li>I see. In that case, it should work after the change that landed after 1.13 that I mentioned. However, is there any reason why the optimizer cannot be created externally? Alternatively, could you make your method a static method and pass in the model, which can later be FSDP-wrapped?</li>
<li>Since the PR with the change landed after the 1.13 branch cut, that should mean it will only appear in the next release. You could access it from a nightly build if that works for your case.</li>
</ol>