Repository: facebookresearch/fairscale

Issue:
Example of auto_wrap is not working as expected
## ❓ Questions and Help

In this doc (https://fairscale.readthedocs.io/en/stable/api/nn/fsdp.html), there is an following example.
```python
import torch
from fairscale.nn.auto_wrap import enable_wrap, auto_wrap
from fairscale.nn.data_parallel import FullyShardedDataParallel as FSDP
fsdp_params = dict(wrapper_cls=FSDP, mixed_precision=True, flatten_parameters=True)
with enable_wrap(**fsdp_params):
    # Wraps layer in FSDP by default if within context
    self.l1 = wrap(torch.nn.Linear(5, 5))
    assert isinstance(self.l1, FSDP)
    # Separately Wraps children modules with more than 1e8 params
    large_tfmr = torch.nn.Transformer(d_model=2048, encoder_layers=12, decoder_layers=12)
    self.l2 = auto_wrap(large_tfmr, min_num_params=1e8)
    assert isinstance(self.l2, FSDP)
```

However, there is some syntax issue, such as no definition of wrap function when it wraps l1. Then, i modified it as the following
```python
     import torch
      from fairscale.nn.wrap import wrap, enable_wrap, auto_wrap
      from fairscale.nn.data_parallel import FullyShardedDataParallel as FSDP
      fsdp_params = dict(wrapper_cls=FSDP, mixed_precision=True, flatten_parameters=True)
      with enable_wrap(**fsdp_params):
          l1 = wrap(torch.nn.Linear(5, 5))
          assert isinstance(l1, FSDP)
          # Wraps layer in FSDP by default if within context
          # Separately Wraps children modules with more than 1e8 params
          large_tfmr = torch.nn.Transformer(d_model=2048, num_encoder_layers=12,
                                            num_decoder_layers=12)
          l2 = auto_wrap(large_tfmr)
          assert isinstance(l2, FSDP)

```

However, l2 is not FSDP. 

pytorch: 1.7.1
fairscale: 0.3.7



Discussion:
Thanks for the report. The example were indeed incorrect. It was meant as pseudo-code. But with the following code, it works now:

```
import torch
from fairscale.nn.wrap import wrap, enable_wrap, auto_wrap
from fairscale.nn.data_parallel import FullyShardedDataParallel as FSDP
from fairscale.utils.testing import dist_init, teardown, rmf

result = dist_init(0, 1, "/tmp/t1", "/tmp/t2")
assert result
fsdp_params = dict(wrapper_cls=FSDP, mixed_precision=True, flatten_parameters=True)
with enable_wrap(**fsdp_params):
    l1 = wrap(torch.nn.Linear(5, 5))
    assert isinstance(l1, FSDP)
    # Wraps layer in FSDP by default if within context
    # Separately Wraps children modules with more than 1e8 params
    large_tfmr = torch.nn.Transformer(d_model=2048, num_encoder_layers=12,
                                      num_decoder_layers=12)
    l2 = auto_wrap(large_tfmr)
    assert isinstance(l2.encoder, FSDP)
    assert isinstance(l2.decoder, FSDP)
teardown()
rmf("/tmp/t1")
rmf("/tmp/t2")
```

Basically, the encoder and decoder modules are wrapped in `l2`, not `l2` itself.

I was using pytorch 1.8.1 and fairscale master.


Resolving patch:
diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -93,10 +93,11 @@ class FullyShardedDataParallel(nn.Module):
     .. _`Xu et al.`: https://arxiv.org/abs/2004.13336
     .. _DeepSpeed: https://www.deepspeed.ai/
 
-    Usage::
+    Pseudo-code usage::
 
         import torch
         from fairscale.nn.data_parallel import FullyShardedDataParallel as FSDP
+
         torch.cuda.set_device(device_id)
         sharded_module = FSDP(my_module)
         optim = torch.optim.Adam(sharded_module.parameters(), lr=0.0001)
@@ -112,17 +113,27 @@ class FullyShardedDataParallel(nn.Module):
     across the forward pass. For example::
 
         import torch
-        from fairscale.nn.auto_wrap import enable_wrap, auto_wrap, wrap
+        from fairscale.nn.wrap import wrap, enable_wrap, auto_wrap
         from fairscale.nn.data_parallel import FullyShardedDataParallel as FSDP
+        from fairscale.utils.testing import dist_init, teardown, rmf
+
+        result = dist_init(0, 1, "/tmp/t1", "/tmp/t2")
+        assert result
         fsdp_params = dict(wrapper_cls=FSDP, mixed_precision=True, flatten_parameters=True)
         with enable_wrap(**fsdp_params):
+            l1 = wrap(torch.nn.Linear(5, 5))
+            assert isinstance(l1, FSDP)
             # Wraps layer in FSDP by default if within context
-            self.l1 = wrap(torch.nn.Linear(5, 5))
-            assert isinstance(self.l1, FSDP)
             # Separately Wraps children modules with more than 1e8 params
-            large_tfmr = torch.nn.Transformer(d_model=2048, encoder_layers=12, decoder_layers=12)
-            self.l2 = auto_wrap(large_tfmr, min_num_params=1e8)
-            assert isinstance(self.l2, FSDP)
+            large_tfmr = torch.nn.Transformer(d_model=2048, num_encoder_layers=12,
+                                              num_decoder_layers=12)
+            l2 = auto_wrap(large_tfmr)
+            assert isinstance(l2.encoder, FSDP)
+            assert isinstance(l2.decoder, FSDP)
+            print(l2)  # You can print the model to examine FSDP wrapping.
+        teardown()
+        rmf("/tmp/t1")
+        rmf("/tmp/t2")
 
     .. warning::
 


Test patch: