Repository: pytorch
Issue: (fsdp) Support for accessing unsharded parameters for methods other than `forward()`
Labels: triaged,module: fsdp

### 🚀 The feature, motivation and pitch

Certain models require execution methods other than `forward()`. A good example family of models are generation models, for example Huggingface's `transformers.PreTrainedModel`. Here, `PreTrainedModel.forward()` is used for training, but during evaluation (which is often in the validation part of the training loop), users need to call `PreTrainedModel.generate()`.

Once the model is FSDP-wrapped, however, the user cannot use methods other than `forward()` to access unsharded parameters. As reported in https://github.com/pytorch/pytorch/issues/82461 , for Encoder-Decoder models, calling `PreTrainedModel.generate()` with FSDP-wrapped model will result in an error due to a direct access to a module (`_fsdp_wrapped_module.encoder`) that is not FSDP-wrapped. The recommendation in this issue was to individually FSDP-wrap `encoder`, but this is not a feasible solution for the case of T5 models, because embedding parameters should be shared across Encoder and Decoder; this results in errors during training. For this reason, HuggingFace Transformer team recommends a hack https://github.com/huggingface/transformers/issues/21667 which is not to wrap `encoder`, but to run `PreTrainedModel.forward()` (with an arbitrary input) before `PreTrainedModel.generate()`, because `forward()`'s side effect seems to populate unsharded parameters. 

In order to support methods other than `forward()`, we could generalize `FullyShardedDataParallel.forward()` to support any function between `_pre_forward()` and `_post_forward()`, making a single-line change to https://github.com/pytorch/pytorch/blob/main/torch/distributed/fsdp/fully_sharded_data_parallel.py#L784-L806 :

```python
class ExecutableFSDP(FullyShardedDataParallel):
    """
    Adds ``execute()`` to FullyShardedDataParallel, which allows access to unsharded parameters.

    .. note:: This is mostly a copy-paste of https://github.com/pytorch/pytorch/blob/main/torch/distributed/fsdp/fully_sharded_data_parallel.py#L784-L806 
    """

    # noinspection PyNoneFunctionAssignment
    def execute(self, fcn: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
        """
        Allows user to run the given ``fcn``, which would access unsharded parameters of the model.

        .. note:: Mostly, copy & paste from ``FullyShardedDataParallel.forward()``.
        """
        with torch.autograd.profiler.record_function(
            "FullyShardedDataParallel.forward"
        ):
            args, kwargs = _root_pre_forward(self, self, args, kwargs)
            unused = None
            unshard_fn = functools.partial(_pre_forward_unshard, self, self._handles)
            reshard_fn = functools.partial(_post_forward_reshard, self, self._handles)
            args, kwargs = _pre_forward(
                self, self._handles, unshard_fn, self._fsdp_wrapped_module, args, kwargs
            )
            for handle in self._handles:
                p_assert(
                    handle.flat_param.device == self.compute_device,
                    "Expected `FlatParameter` to be on the compute device "
                    f"{self.compute_device} but got {handle.flat_param.device}",
                )
            # !!! this is the only line we change from ``FullyShardedDataParallel.forward()`` !!!
            output = fcn(*args, **kwargs)

            return _post_forward(self, self._handles, reshard_fn, self, unused, output)


class GenerationFSDP(ExecutableFSDP):
    """
    Adds support for ``PreTrainedModel.generate()`` for ``FullyShardedDataParallel`` models.
    """

    def __init__(self, module: PreTrainedModel, *args: Any, **kwargs: Any):
        super().__init__(module, *args, **kwargs)

    def generate(
        self, *args: Any, **kwargs: Any
    ) -> Union[GenerateOutput, torch.LongTensor]:
        """
        Runs ``PreTrainedModel.generate()`` with access to unsharded parameters.
        """

        return self.execute(
            self._fsdp_wrapped_module.generate,
            *args,
            **kwargs,
        )
```

Of course, such a copy-paste approach is bound to fail when the implementation of `FullyShardedDataParallel.forward()` changes. Hence I was hoping this feature could be directly supported in `FullyShardedDataParallel`.

If pytorch team wouldn't plan for supporting such feature, would the team advise users for a different approach?

### Alternatives

Continue users to use the hack to run `forward()` before `generate()`, mentioned in  https://github.com/pytorch/pytorch/issues/82461#issuecomment-1203377083 and recommended in https://github.com/huggingface/transformers/issues/21667

### Additional context

_No response_

cc @zhaojuanmao @mrshenli @rohan-varma @awgu @ezyang @msaroufim @wconstab @bdhirsh @anijain2305