Repository: facebookresearch/fairscale

Issue:
Simple example of use `FullyShardedDataParallel`
## ❓ Questions and Help

I may be doing something completely wrong, but I'm struggling to just get one forward pass through `FullyShardedDataParallel` without error. I've created a simple script:

```python
import os
from fairscale.nn.data_parallel import FullyShardedDataParallel as FullyShardedDDP
from transformers import BertModel
import torch
import torch.nn as nn
def main():
    torch.distributed.init_process_group(backend="nccl")
    # device = torch.device(int(os.environ.get("LOCAL_RANK", -1)))
    model = nn.Linear(200, 200).cuda()
    model = FullyShardedDDP(model)
    fake_inputs = torch.randn(32, 200).to(torch.device("cuda"))
    outputs = model(fake_inputs)

if __name__ == "__main__":
    main()
```
launched via
```bash
python -m torch.distributed.launch --nproc_per_node 2 --use_env path_to_script
```
and it fails with the following stack trace:
```
*****************************************
Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. 
*****************************************
Traceback (most recent call last):
  File "test_zero.py", line 20, in <module>
Traceback (most recent call last):
  File "test_zero.py", line 20, in <module>
    main()
  File "test_zero.py", line 13, in main
    model = FullyShardedDDP(model.cuda())
  File "/home/sgugger/git/fairscale/fairscale/nn/data_parallel/fully_sharded_data_parallel.py", line 161, in __init__
    main()
  File "test_zero.py", line 13, in main
    validate_process_group(compute_device, self.process_group)
  File "/home/sgugger/git/fairscale/fairscale/utils/parallel.py", line 41, in validate_process_group
    model = FullyShardedDDP(model.cuda())
  File "/home/sgugger/git/fairscale/fairscale/nn/data_parallel/fully_sharded_data_parallel.py", line 161, in __init__
    dist.all_gather(output, input_tensor, group=process_group)
  File "/home/sgugger/.pyenv/versions/base/lib/python3.7/site-packages/torch/distributed/distributed_c10d.py", line 1434, in all_gather
    validate_process_group(compute_device, self.process_group)
  File "/home/sgugger/git/fairscale/fairscale/utils/parallel.py", line 41, in validate_process_group
    dist.all_gather(output, input_tensor, group=process_group)
  File "/home/sgugger/.pyenv/versions/base/lib/python3.7/site-packages/torch/distributed/distributed_c10d.py", line 1434, in all_gather
    work = group.allgather([tensor_list], [tensor])
RuntimeError: NCCL error in: /pytorch/torch/lib/c10d/ProcessGroupNCCL.cpp:784, invalid usage, NCCL version 2.7.8
    work = group.allgather([tensor_list], [tensor])
RuntimeError: NCCL error in: /pytorch/torch/lib/c10d/ProcessGroupNCCL.cpp:784, invalid usage, NCCL version 2.7.8
Traceback (most recent call last):
  File "/home/sgugger/.pyenv/versions/3.7.9/lib/python3.7/runpy.py", line 193, in _run_module_as_main
    "__main__", mod_spec)
  File "/home/sgugger/.pyenv/versions/3.7.9/lib/python3.7/runpy.py", line 85, in _run_code
    exec(code, run_globals)
  File "/home/sgugger/.pyenv/versions/base/lib/python3.7/site-packages/torch/distributed/launch.py", line 260, in <module>
    main()
  File "/home/sgugger/.pyenv/versions/base/lib/python3.7/site-packages/torch/distributed/launch.py", line 256, in main
    cmd=cmd)
subprocess.CalledProcessError: Command '['/home/sgugger/.pyenv/versions/base/bin/python', '-u', 'test_zero.py']' returned non-zero exit status 1.
```

I'm using the `.cuda()` following what is done in the tests. Otherwise using the device (commented out in the script) and `.to(device)` also fails, with a different error:
```
*****************************************
Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. 
*****************************************
Traceback (most recent call last):
  File "test_zero.py", line 20, in <module>
    main()
  File "test_zero.py", line 16, in main
    outputs = model(fake_inputs)
  File "/home/sgugger/.pyenv/versions/base/lib/python3.7/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/sgugger/git/fairscale/fairscale/nn/data_parallel/fully_sharded_data_parallel.py", line 612, in forward
    self._rebuild_full_params()
  File "/home/sgugger/.pyenv/versions/base/lib/python3.7/site-packages/torch/autograd/grad_mode.py", line 26, in decorate_context
    return func(*args, **kwargs)
  File "/home/sgugger/git/fairscale/fairscale/nn/data_parallel/fully_sharded_data_parallel.py", line 822, in _rebuild_full_params
    dist.all_gather(chunks, p.data, group=self.process_group)
  File "/home/sgugger/.pyenv/versions/base/lib/python3.7/site-packages/torch/distributed/distributed_c10d.py", line 1434, in all_gather
    work = group.allgather([tensor_list], [tensor])
RuntimeError: Corresponding input/output tensors to scatter/gather must all reside on the same device
```



Discussion:
Ah, you need to set the default cuda device depending on the rank, then it works:
```python
# (...)
device = torch.device(int(os.environ.get("LOCAL_RANK", -1)))
torch.cuda.set_device(device)
# (...)
```

It seems this is also needed for `DistributedDataParallel`, right? Just want to make sure FSDP isn't less robust than PyTorch DDP in this respect...
Indeed! To answer, you don't need to, no: for instance this script

```python
import os

import torch
import torch.nn as nn


def main():
    torch.distributed.init_process_group(backend="nccl")
    local_rank = int(os.environ.get("LOCAL_RANK", -1))
    device = torch.device(local_rank)
    model = nn.Linear(200, 200).to(device)
    model =  torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank],  output_device=local_rank)
    fake_inputs = torch.randn(32, 200).to(device)
    outputs = model(fake_inputs)


if __name__ == "__main__":
    main()
```

will work.

Second question now is: does `FullyShardedDataParallel` support Embeddings layer? Changing the script to the following:
```python
import os
from fairscale.nn.data_parallel import FullyShardedDataParallel as FullyShardedDDP

from transformers import BertModel
import torch
import torch.nn as nn


def main():
    torch.distributed.init_process_group(backend="nccl")
    local_rank = int(os.environ.get("LOCAL_RANK", -1))
    device = torch.device(local_rank)
    torch.cuda.set_device(device)
    model = BertModel.from_pretrained("bert-base-cased").to(device)
    model = FullyShardedDDP(model)
    fake_inputs = torch.randint(0, 1000, (32, 200)).long().to(device)
    outputs = model(fake_inputs)


if __name__ == "__main__":
    main()
```

fails with the following stack trace:

```
*****************************************
Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed.
*****************************************
Traceback (most recent call last):
  File "test_zero.py", line 21, in <module>
Traceback (most recent call last):
  File "test_zero.py", line 21, in <module>
    main()
  File "test_zero.py", line 17, in main
    outputs = model(fake_inputs)
  File "/home/sgugger/.pyenv/versions/base/lib/python3.7/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    main()
  File "test_zero.py", line 17, in main
    outputs = model(fake_inputs)
  File "/home/sgugger/.pyenv/versions/base/lib/python3.7/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/sgugger/git/fairscale/fairscale/nn/data_parallel/fully_sharded_data_parallel.py", line 618, in forward
    result = self.forward(*input, **kwargs)
  File "/home/sgugger/git/fairscale/fairscale/nn/data_parallel/fully_sharded_data_parallel.py", line 618, in forward
    outputs = self.module(*args, **kwargs)
  File "/home/sgugger/.pyenv/versions/base/lib/python3.7/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    outputs = self.module(*args, **kwargs)
  File "/home/sgugger/.pyenv/versions/base/lib/python3.7/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/sgugger/git/fairscale/fairscale/nn/misc/flatten_params_wrapper.py", line 153, in forward
    result = self.forward(*input, **kwargs)
  File "/home/sgugger/git/fairscale/fairscale/nn/misc/flatten_params_wrapper.py", line 153, in forward
    return self.module(*inputs, **kwinputs)
  File "/home/sgugger/.pyenv/versions/base/lib/python3.7/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    return self.module(*inputs, **kwinputs)
  File "/home/sgugger/.pyenv/versions/base/lib/python3.7/site-packages/torch/nn/modules/module.py", line 727, in _call_impl      
    result = self.forward(*input, **kwargs) 
  File "/home/sgugger/git/transformers/src/transformers/models/bert/modeling_bert.py", line 964, in forward
    result = self.forward(*input, **kwargs)
  File "/home/sgugger/git/transformers/src/transformers/models/bert/modeling_bert.py", line 964, in forward 
    past_key_values_length=past_key_values_length,
  File "/home/sgugger/.pyenv/versions/base/lib/python3.7/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    past_key_values_length=past_key_values_length,
  File "/home/sgugger/.pyenv/versions/base/lib/python3.7/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/sgugger/git/transformers/src/transformers/models/bert/modeling_bert.py", line 205, in forward
    result = self.forward(*input, **kwargs)
  File "/home/sgugger/git/transformers/src/transformers/models/bert/modeling_bert.py", line 205, in forward
    position_embeddings = self.position_embeddings(position_ids)
  File "/home/sgugger/.pyenv/versions/base/lib/python3.7/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    position_embeddings = self.position_embeddings(position_ids)
  File "/home/sgugger/.pyenv/versions/base/lib/python3.7/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/sgugger/.pyenv/versions/base/lib/python3.7/site-packages/torch/nn/modules/sparse.py", line 126, in forward  result = self.forward(*input, **kwargs)
  File "/home/sgugger/.pyenv/versions/base/lib/python3.7/site-packages/torch/nn/modules/sparse.py", line 126, in forward  self.norm_type, self.scale_grad_by_freq, self.sparse)
  File "/home/sgugger/.pyenv/versions/base/lib/python3.7/site-packages/torch/nn/functional.py", line 1852, in embedding
    self.norm_type, self.scale_grad_by_freq, self.sparse)
  File "/home/sgugger/.pyenv/versions/base/lib/python3.7/site-packages/torch/nn/functional.py", line 1852, in embedding
       return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
	   
RuntimeErrorRuntimeError: : Expected tensor for argument #1 'indices' to have scalar type Long; but got torch.cuda.FloatTensor instead (while checking arguments for embedding)Expected tensor for argument #1 'indices' to have scalar type Long; but got torch.cuda.FloatTensor instead (while checking arguments for embedding)                                                                                                                                                                                                                                                                                                                             Traceback (most recent call last):                                                                                                                                                                                     File "/home/sgugger/.pyenv/versions/3.7.9/lib/python3.7/runpy.py", line 193, in _run_module_as_main                                                                                                                    "__main__", mod_spec)                                                                                                                                                                                              File "/home/sgugger/.pyenv/versions/3.7.9/lib/python3.7/runpy.py", line 85, in _run_code                                                                                                                               exec(code, run_globals)                                                                                                                                                                                            File "/home/sgugger/.pyenv/versions/base/lib/python3.7/site-packages/torch/distributed/launch.py", line 260, in <module>                                                                                               main()                                                                                                                                                                                                             File "/home/sgugger/.pyenv/versions/base/lib/python3.7/site-packages/torch/distributed/launch.py", line 256, in main                                                                                                   cmd=cmd)                                                                                                                                                                                                         


subprocess.CalledProcessError: Command '['/home/sgugger/.pyenv/versions/base/bin/python', '-u', 'test_zero.py']' returned non-zero exit status 1.  
```

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
@@ -931,11 +931,15 @@ def cast_inputs_to_fp16(*args: Any, **kwargs: Any) -> Tuple[Any, Any]:
 def cast_buffers_(
     module: nn.Module, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None
 ) -> None:
-    """Cast all of module.named_buffers to device, dtype."""
+    """Cast all of module.named_buffers to device and floating point buffers to dtype."""
     # if buffers are already on the right device and/or dtype this is just python loop cost
+    assert dtype in {torch.float32, torch.float16}  # assumes compute_dtype == float16
     for key, buf in module.named_buffers(recurse=False):
         if buf is not None:
-            setattr(module, key, buf.to(dtype=dtype, device=device))
+            buf = buf.to(device=device)
+            if torch.is_floating_point(buf):
+                buf = buf.to(dtype=dtype)
+            setattr(module, key, buf)
 
 
 def free_storage_(data: torch.Tensor) -> None:


Test patch:
diff --git a/tests/nn/data_parallel/test_fsdp.py b/tests/nn/data_parallel/test_fsdp.py
--- a/tests/nn/data_parallel/test_fsdp.py
+++ b/tests/nn/data_parallel/test_fsdp.py
@@ -29,8 +29,6 @@
 # How to use remote-pdb: https://gist.github.com/sshleifer/9d43351957179c13606e015b072927d4
 # All helper functions called by spawn must be either @classmethod, @staticmethod
 
-_BUFFER_NAME = "vocab_bias"
-
 
 class DistributedTest(unittest.TestCase):
     def setUp(self):
@@ -411,8 +409,9 @@ def _load_local_and_train(self, config, rank, group, d_model=16, d_vocab=23):
 
         # Assert that parameters were updated since before training
         unchanged = []
+        buffers = {name for name, _ in model.module.named_buffers()}
         for k in state_1:
-            if (state_before_training[k] == state_after_training[k]).all() and (_BUFFER_NAME not in k):
+            if (state_before_training[k] == state_after_training[k]).all() and (k not in buffers):
                 unchanged.append(k)
         if unchanged:
             raise AssertionError(f"params {unchanged} not changed after training")
@@ -651,7 +650,8 @@ def __init__(self, group, *unused_args, d_vocab=23, d_model=16, **unused_kwargs)
         self.output_proj = nn.Linear(d_model, d_vocab)
         # share the embedding and output projection weights
         self.output_proj.weight = self.embed_tokens.weight
-        self.register_buffer(_BUFFER_NAME, self.embed_tokens.weight.new_ones((d_model,)))
+        self.register_buffer("vocab_bias", self.embed_tokens.weight.new_ones((d_model,)))
+        self.register_buffer("long_buffer", torch.zeros_like(self.vocab_bias, dtype=torch.long))
 
     def get_input(self, device):
         torch.manual_seed(1 + self.rank)  # keep everything deterministic
@@ -661,7 +661,7 @@ def get_input(self, device):
 
     def forward(self, src_ids, tgt_ids):
         src = self.embed_tokens(src_ids)
-        src = src + self.vocab_bias
+        src = src + self.vocab_bias + self.long_buffer.type_as(src)
         tgt = self.embed_tokens(tgt_ids)
         x = self.transformer(src, tgt)
         return self.output_proj(x)