Repository: facebookresearch/fairscale

Issue:
The Argument broadcast_fp16 is not working on OSS
## 🐛 Bug

The Argument `broadcast_fp16` is not working on `OSS`.

## Command



## To Reproduce

Steps to reproduce the behavior:

1. Add one print line in fairscale's source code [here](https://github.com/facebookresearch/fairscale/blob/31e3645371e1462c0fb119a7a54642e327d0d4bf/fairscale/optim/oss.py#L603)
```
            for device in self.buckets.keys():
                for dst_rank, bucket in self.buckets[device].items():
                    print(bucket.buffer.dtype, bucket.buffer.shape)
                    work_handles.append(
```
2. Run a test script as below:
```
import sys
import pdb
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP

from fairscale.optim.oss import OSS
from fairscale.nn.data_parallel import ShardedDataParallel as ShardedDDP

CUDA_DEVICE = [0, 1]
HIDDEN_SIZE = 8

class ToyModel(torch.nn.Module):
    def __init__(self):
        super(ToyModel, self).__init__()
        self.net1 = torch.nn.Linear(HIDDEN_SIZE, HIDDEN_SIZE)
        self.relu = torch.nn.ReLU()
        self.net2 = torch.nn.Linear(HIDDEN_SIZE, 5)

    def forward(self, x):
        return self.net2(self.relu(self.net1(x)))


def main(rank, world_size):

    dist.init_process_group(
            backend='nccl',
            init_method="tcp://localhost:29501",
            rank=rank,
            world_size=world_size)
    device = CUDA_DEVICE[rank]
    torch.manual_seed(rank + 1007)
    torch.cuda.set_device(device)
    torch.cuda.manual_seed(rank + 1007)
    print('Rank: %d/%d  Device: %s' % (rank, world_size, device))


    model = ToyModel().to(device)
    loss_fn = torch.nn.MSELoss()
    # ddp_model = DDP(model, device_ids=[device])
    # optimizer = torch.optim.SGD(ddp_model.parameters(), lr=0.1)
    optimizer = OSS(params=model.parameters(), optim=torch.optim.SGD, lr=0.1, broadcast_fp16=True)
    ddp_model = ShardedDDP(model, optimizer)

    inputs = torch.randn(2, HIDDEN_SIZE).to(device)
    print('inputs', rank, inputs)

    optimizer.zero_grad()
    outputs = ddp_model(inputs)
    labels = torch.randn(2, 5).to(device)
    loss_fn(outputs, labels).backward()
    optimizer.step()
    

    dist.destroy_process_group()


if __name__ == "__main__":
    WORLD_SIZE = len(CUDA_DEVICE)
    mp.spawn(
        main,
        args=(WORLD_SIZE,),
        nprocs=WORLD_SIZE,
        join=True,
    )

```

And the output is:
```
...
torch.float32 torch.Size([64])
torch.float32 torch.Size([53])
torch.float32 torch.Size([64])
torch.float32 torch.Size([53])
```

Note that we have already set the argument `broadcast_fp16` to be true, but fairscale still boardcasts float32 tensor.

## Reason

The Tensor.to [(code here)](https://github.com/facebookresearch/fairscale/blob/31e3645371e1462c0fb119a7a54642e327d0d4bf/fairscale/nn/misc/param_bucket.py#L35) method is not in-place operation.

## Posible Solution
Change [fairscale/nn/misc/param_bucket.py:35](https://github.com/facebookresearch/fairscale/blob/31e3645371e1462c0fb119a7a54642e327d0d4bf/fairscale/nn/misc/param_bucket.py#L35) to 
```
self.buffer = self.buffer.to(device, dtype, non_blocking)
```


Discussion:
Thank you for opening the bug! We might want to return a new self.buffer post move [here](https://github.com/facebookresearch/fairscale/blob/31e3645371e1462c0fb119a7a54642e327d0d4bf/fairscale/optim/oss.py#L596) where we are converting the buffers to FP16. I also did not see a test where we check for the buffer type conversion which would be hard to do unless you had a way to not convert back to FP32 post broadcast (so that you could check). 

@blefaudeux  WDYT? You have the most context here. I also remember you fixed a bug where we were only in FP16 throughout. I wonder how that was possible if we were never broadcasting to FP16. 

Oh ! Thanks for the finding and the detailed bug report, it must have been broken during a refactoring (the buckets used to be in the same file if I remember correctly), it's not tested indeed :( I'll put up a fix + a test around 9:30 PT, sorry about that

Resolving patch:
diff --git a/fairscale/nn/misc/param_bucket.py b/fairscale/nn/misc/param_bucket.py
--- a/fairscale/nn/misc/param_bucket.py
+++ b/fairscale/nn/misc/param_bucket.py
@@ -32,7 +32,7 @@ def to(  # type: ignore
         Move the underlying buffer
         """
         assert self.buffer is not None, "Cannot move a collapsed bucket, please rebuild it"
-        self.buffer.to(device, dtype, non_blocking)
+        self.buffer = self.buffer.to(device, dtype, non_blocking)
 
 
 class ParamBucket(Bucket):
@@ -98,6 +98,8 @@ def _reattach_params(self) -> None:
 
         self._fill = 0
         for p in self._params:
+            if p.dtype != self.buffer.dtype:
+                p.data = p.data.to(self.buffer.dtype)
             self._add_param_as_view(p, keep_existing_value=False)
 
 


Test patch:
diff --git a/tests/nn/misc/test_param_bucket.py b/tests/nn/misc/test_param_bucket.py
--- a/tests/nn/misc/test_param_bucket.py
+++ b/tests/nn/misc/test_param_bucket.py
@@ -48,7 +48,10 @@ def test_type_change():
 
     # Move the bucket to fp16 and back
     bucket.to(dtype=torch.float16, device=param.device)
+    assert bucket.buffer.dtype == torch.float16
+
     bucket.to(dtype=torch.float32, device=param.device, keep_param_alignment=True)
+    assert bucket.buffer.dtype == torch.float32
 
     # Same with the reference tensor
     param_.to(dtype=torch.float16)