sharded_ddp with auto_wrap is secretly a no-op
Author: justheuristic
Pull request: no
Labels: name: bug

### System Info

```shell
transformers 26e5e129b43760138aed2dfc1cc3c75b481a95e6
torch 1.11.0
fairscale 0.4.6
```


### Who can help?

@sgugger 

### Information

- [X] The official example scripts
- [ ] My own modified scripts

### Tasks

- [x] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

The problem occurs in any scenario where you use Trainer with `--sharded_ddp` and use auto_wrap.

Trainer initializes sharded_ddp ([trainer.py:L1183](https://github.com/huggingface/transformers/blob/2f59ad1/src/transformers/trainer.py#L1183-L1184)), here's how trainer calls auto_wrap
```python
    if ShardedDDPOption.AUTO_WRAP in self.args.sharded_ddp:
        model = auto_wrap(model)
```
Then, auto_wrap checks if you are .in_autowrap_context ([auto_wrap.py:L224](https://github.com/facebookresearch/fairscale/blob/3b8f445/fairscale/nn/wrap/auto_wrap.py#L224-L229)), which is [False by default](https://github.com/facebookresearch/fairscale/blob/3b8f445/fairscale/nn/wrap/auto_wrap.py#L238)
```python
    if ConfigAutoWrap.in_autowrap_context:  # <-- False
        wrapped_module, remainder = ConfigAutoWrap.recursive_wrap(
            module, auto_wrap_policy=auto_wrap_policy, module_is_root=True, **kwargs
        )
        return wrapped_module
    return module  # <-- so you get this
```

As a result, auto_wrap does nothing and the model ends up not being wrapped. In this setting, ShardedDDP is just an inferior version of regular DDP with no memory savings. The worst part is that __it doesn't give you any clue that something is wrong__. It just silently does nothing.


__Example scenario:__
1. install the latest transformers and fairscale (see versions above)
3. go to [language-modelling](https://github.com/huggingface/transformers/tree/main/examples/pytorch/language-modeling) example and follow the tutorial for "GPT-2/GPT and causal language modeling"
4. get to the part where you [run_clm.py](https://github.com/huggingface/transformers/blob/main/examples/pytorch/language-modeling/run_clm.py)
5. add the following parameters: `--fp16 --sharded_ddp="zero_dp_3 auto_wrap"`




### Expected behavior


```shell
I would expect these options to result in optimizer and/or parameters being sharded and result in reduced GPU memory usage, but that does not happen.

(A) if you have time, it would be nice to make auto_wrap actually wrap the model

(B) if you don't, perhaps it would be simpler to add an assert


from fairscale.whatever import WrappedClassNameHere
assert any(isinstance(module, WrappedClassNameHere) for module in model.modules()), "running sharded_ddp, but there are no modules wrapped for sharding."
```
```


Comments:
We don't really support fairscale sharded ddp anymore. You should switch to the newer native FSDP integration which we support (doc is [here](https://huggingface.co/docs/transformers/main_classes/trainer#pytorch-fully-sharded-data-parallel)).

Thanks for the response.
I only mean to say, that it would be great to add some kind of warning, saying that sharded_ddp is not supported - in code and/or the above tutorial.
Otherwise, the code pretends to work, but does nothing which confuse some users.

p.s. as of right now, FSDP does not support --fp16 and does not work with stable pytorch, so we're trying deepspeed as a workaround.




Which tutorial advertises fairscale sharded DDP? We haven't really recommended it (we do recommend DeepSpeed which is more mature at this stage).

> (doc is [here](https://huggingface.co/docs/transformers/main_classes/trainer#pytorch-fully-sharded-data-parallel)).

In the link you suggested, if you scroll up, it says:

> Trainer Integrations

> The [Trainer](https://huggingface.co/docs/transformers/v4.19.2/en/main_classes/trainer#transformers.Trainer) has been extended to support libraries that may dramatically improve your training time and fit much bigger models.

> Currently it __supports__ third party solutions, [DeepSpeed](https://github.com/microsoft/DeepSpeed) __and [FairScale](https://github.com/facebookresearch/fairscale/)__, which implement parts of the paper [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models, by Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, Yuxiong He](https://arxiv.org/abs/1910.02054).

And then, just before the FSDP section,
> FairScale

> By integrating [FairScale](https://github.com/facebookresearch/fairscale/) the __[Trainer](https://huggingface.co/docs/transformers/v4.19.2/en/main_classes/trainer#transformers.Trainer) provides support__ for the following features from [the ZeRO paper](https://arxiv.org/abs/1910.02054): [list of features]

One may misinterpret this to mean that Trainer supports sharded_ddp in the same way it supports FSDP - or at least, that was how i misinterpreted it.

That said, i fully understand that this issue is low priority.

Thanks a lot for the pointers! I'll have a look at that doc and update it :-)