Repository: pytorch
Issue: [Discussion][FSDP] Enhancements to auto_wrap_policy
Labels: oncall: distributed,triaged,module: fsdp

### 🚀 The feature, motivation and pitch

For meta device support, the following API has issues:

```
class Model:
    def __init__(self):
        self.l1 = MyModule(wrap(Linear), Linear)
        # more wrap() calls

with enable_wrap(wrapper_cls=FSDP):
    model = wrap(deferred_init(Model()))
```

for details, see: https://github.com/cbalioglu/torchdistx/issues/3

To summarize the issue, when recursively nesting FSDP instances and composing that with deferred_init, tensors in `apply()` function can be created in fake tensor context resulting in issues. In general, it feels like using `auto_wrap_policy` with meta device modules is more idiomatic, i.e. the user simply does the following:

```
model = FSDP(deferred_init(Model), policy=my_policy)
```

The issue is sometimes it is convenient for users to decide whether to wrap a submodule with FSDP or not as they are building their module / defining layers, i.e. by adding something like `wrap()` which should annotate whether a module should be wrapped or not. This allows users to conveniently wrap / unwrap parts of their model without having to express all of this in a `policy` function which might requiring careful care of the policy code.

As a result, it could be useful to enhance `wrap()` function to do the following:

```
def wrap(module):
    if in_wrap_ctx:
        # proceed with regular logic
    else:
         module._should_wrap = True # annotate it should be wrapped later
```

And add a policy which does the following:

```
def wrap_if_annotated():
""""Returns `True` if `module` was annotated by `wrap`."""
    return getattr(module, '_should_wrap, False)
```

Finally, we would also want to enhance `recursive_wrap` feature which currently exits recursive logic if a parent module does not satisfy the policy, so there is no way for a child to be FSDP if the parent is not. This can actually be taken care of by the policy writer themselves, as they can recurse down the module tree and check child modules appropriately. Although in this case if the policy returns `True`, the parent module would be wrapped even if that wasn't the intent. 

As an example, in the model in the first code block, using `wrap_if_annotated` policy, MyModule would not be wrapped and as a result the child, which is annotated with wrap, would not be either. To fix this, we can add a `callback` arg in `recursive_wrap` /part of `auto_wrap_policy` that allows users to specify what they want to do when a policy returns `True`. Currently when the policy returns `True` we immediately wrap the module with FSDP, but in this case the user might want to just continue recursing down the module hierarchy, wrapping children that were explicitly marked with `wrap`.

Thus the overall proposal is something like:

1. Define model with `wrap`, which just sets an annotate attr:

```
class Model:
    def __init__(self):
        self.l1 = MyModule(wrap(Linear), Linear)
        # more wrap() calls
```

2. Write policy and callback (can provide these out of the box):
```
def wrap_if_annotated(module):
""""Returns `True` if `module` was annotated by `wrap`."""
    for module in child_modules():
        if getattr(module, '_should_wrap', False): return True
    return False

def on_policy_triggered_callback(module, wrapper_cls, **kwargs):
    if module._should_wrap:
        default_policy_callback()     # default implementation just does wrapper_cls(module, **kwargs)
```

3. Make `fsdp_auto_wrap_policy` a dataclass with the policy and callback. For BC and ease of use we will continue to accept a simple lambda function

4. Use this in FSDP:

```
model = FSDP(deferred_init(Model, *args, **kwargs), fsdp_auto_wrap_policy=AutoWrapPolicy(policy=wrap_if_annotated, callback=on_policy_triggered_callback))
```

Most of this besides changes to recursive wrapping can actually be done without changing FSDP core codebase. Thus the main proposal for now is to make this bc-compatible change to recursive wrapping, and optionally provide these helpful policies to work with deferred_init.







### Alternatives

_No response_

### Additional context

_No response_

cc @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @SciPioneer @H-Huang