Repository: facebookresearch/fairscale

Issue:
[SDP] Bucket Failed to be sent error in FairScale 0.3.1
## 🐛 Bug

When upgrading from 0.3.0 to 0.3.1 on the Lightning side our bug report script fails with sharded:

You will need to install lightning (`pip install pytorch-lightning`).

```python
import os

import torch
from torch.utils.data import Dataset

from pl_examples import cli_lightning_logo
from pytorch_lightning import LightningModule, Trainer


class RandomDataset(Dataset):
    """
    >>> RandomDataset(size=10, length=20)  # doctest: +ELLIPSIS
    <...bug_report_model.RandomDataset object at ...>
    """

    def __init__(self, size, length):
        self.len = length
        self.data = torch.randn(length, size)

    def __getitem__(self, index):
        return self.data[index]

    def __len__(self):
        return self.len


class BoringModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(32, 2)

    def forward(self, x):
        return self.layer(x)

    def loss(self, batch, prediction):
        # An arbitrary loss to have a loss that updates the model weights during `Trainer.fit` calls
        return torch.nn.functional.mse_loss(prediction, torch.ones_like(prediction))

    def step(self, x):
        x = self.layer(x)
        out = torch.nn.functional.mse_loss(x, torch.ones_like(x))
        return out

    def training_step(self, batch, batch_idx):
        output = self.layer(batch)
        loss = self.loss(batch, output)
        return {"loss": loss}

    def validation_step(self, batch, batch_idx):
        output = self.layer(batch)
        loss = self.loss(batch, output)
        return {"x": loss}

    def test_step(self, batch, batch_idx):
        output = self.layer(batch)
        loss = self.loss(batch, output)
        return {"y": loss}

    def configure_optimizers(self):
        optimizer = torch.optim.SGD(self.layer.parameters(), lr=0.1)
        lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1)
        return [optimizer], [lr_scheduler]


if __name__ == '__main__':
    # fake data
    train_data = torch.utils.data.DataLoader(RandomDataset(32, 64))
    val_data = torch.utils.data.DataLoader(RandomDataset(32, 64))
    test_data = torch.utils.data.DataLoader(RandomDataset(32, 64))

    # model
    model = BoringModel()
    trainer = Trainer(
        default_root_dir=os.getcwd(),
        limit_train_batches=1,
        limit_val_batches=1,
        max_epochs=1,
        weights_summary=None,
        gpus=1,
        accelerator='ddp_sharded',

    )
    trainer.fit(model, train_data, val_data)
```

```
  File "/home/sean/pytorch-lightning/pytorch_lightning/trainer/trainer.py", line 589, in run_train
    self.train_loop.run_training_epoch()
  File "/home/sean/pytorch-lightning/pytorch_lightning/trainer/training_loop.py", line 491, in run_training_epoch
    batch_output = self.run_training_batch(batch, batch_idx, dataloader_idx)
  File "/home/sean/pytorch-lightning/pytorch_lightning/trainer/training_loop.py", line 650, in run_training_batch
    self.optimizer_step(optimizer, opt_idx, batch_idx, train_step_and_backward_closure)
  File "/home/sean/pytorch-lightning/pytorch_lightning/trainer/training_loop.py", line 425, in optimizer_step
    model_ref.optimizer_step(
  File "/home/sean/pytorch-lightning/pytorch_lightning/core/lightning.py", line 1397, in optimizer_step
    optimizer.step(closure=optimizer_closure)
  File "/home/sean/pytorch-lightning/pytorch_lightning/core/optimizer.py", line 214, in step
    self.__optimizer_step(*args, closure=closure, profiler_name=profiler_name, **kwargs)
  File "/home/sean/pytorch-lightning/pytorch_lightning/core/optimizer.py", line 134, in __optimizer_step
    trainer.accelerator.optimizer_step(optimizer, self._optimizer_idx, lambda_closure=closure, **kwargs)
  File "/home/sean/pytorch-lightning/pytorch_lightning/accelerators/accelerator.py", line 289, in optimizer_step
    self.run_optimizer_step(optimizer, opt_idx, lambda_closure, **kwargs)
  File "/home/sean/pytorch-lightning/pytorch_lightning/accelerators/accelerator.py", line 296, in run_optimizer_step
    self.training_type_plugin.optimizer_step(optimizer, lambda_closure=lambda_closure, **kwargs)
  File "/home/sean/pytorch-lightning/pytorch_lightning/plugins/training_type/training_type_plugin.py", line 171, in optimizer_step
    optimizer.step(closure=lambda_closure, **kwargs)
  File "/home/sean/miniconda3/lib/python3.8/site-packages/torch/optim/optimizer.py", line 89, in wrapper
    return func(*args, **kwargs)
  File "/home/sean/miniconda3/lib/python3.8/site-packages/fairscale/optim/oss.py", line 224, in step
    loss = self.optim.step(closure=closure, **kwargs)  # type: ignore
  File "/home/sean/miniconda3/lib/python3.8/site-packages/torch/optim/optimizer.py", line 89, in wrapper
    return func(*args, **kwargs)
  File "/home/sean/miniconda3/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/home/sean/miniconda3/lib/python3.8/site-packages/torch/optim/sgd.py", line 87, in step
    loss = closure()
  File "/home/sean/pytorch-lightning/pytorch_lightning/trainer/training_loop.py", line 644, in train_step_and_backward_closure
    result = self.training_step_and_backward(
  File "/home/sean/pytorch-lightning/pytorch_lightning/trainer/training_loop.py", line 738, in training_step_and_backward
    result = self.training_step(split_batch, batch_idx, opt_idx, hiddens)
  File "/home/sean/pytorch-lightning/pytorch_lightning/trainer/training_loop.py", line 292, in training_step
    training_step_output = self.trainer.accelerator.training_step(args)
  File "/home/sean/pytorch-lightning/pytorch_lightning/accelerators/accelerator.py", line 168, in training_step
    return self.training_type_plugin.training_step(*args)
  File "/home/sean/pytorch-lightning/pytorch_lightning/plugins/training_type/ddp.py", line 298, in training_step
    return self.model(*args, **kwargs)
  File "/home/sean/miniconda3/lib/python3.8/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/sean/miniconda3/lib/python3.8/site-packages/fairscale/nn/data_parallel/sharded_ddp.py", line 209, in forward
    self._clear_counters()
  File "/home/sean/miniconda3/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/home/sean/miniconda3/lib/python3.8/site-packages/fairscale/nn/data_parallel/sharded_ddp.py", line 367, in _clear_counters
    assert (
AssertionError: A bucket failed to be sent, cannot continue as results would be wrong. You can try de-activating ShardedDDP buckets -set `reduce_buffer_size` to 0-Please submit a GitHub issue, this should not happen
```

I haven't been able to track the root cause yet so there is always the possibility it's on the lightning side, but can help to track it down! cc @blefaudeux 


Discussion:
Oh, the buckets flipped to being default on, but that's with an extra callback which should always make sure that they are sent (it even works fine with imgpt for instance, and there's a "unused params" unit test for that also). I'll check that out ASAP, should not have happened 
I'm unable to repro outside of lightning, even with a single GPU as it seems to be the case here. Testing the lightning case right now
ok, I can repro indeed, thanks for the script @SeanNaren 
limited to 1 gpu it seems, so probably not affecting a lot of people. Still, a bit perplexing

so it looks like, on one gpu at least, autograd never touches anything (either no params are trainable or they are not used). Since ShardedDDP relies on autograd to trigger the reduce calls, the assert rightfully calls quit. @ananthsub would you happen to know how the lightning shardedddp plugin works with a single gpu ?

Resolving patch:
diff --git a/fairscale/nn/data_parallel/sharded_ddp.py b/fairscale/nn/data_parallel/sharded_ddp.py
--- a/fairscale/nn/data_parallel/sharded_ddp.py
+++ b/fairscale/nn/data_parallel/sharded_ddp.py
@@ -160,6 +160,11 @@ def __init__(
         # - setup buckets and tensor views
         model_size = sum([p.numel() for p in self.module.parameters()])
         self.buffer_max_size = min(reduce_buffer_size, model_size)
+
+        if dist.get_world_size(self.process_group) == 1:
+            self.buffer_max_size = 0
+            logging.info("Training is not really distributed, single rank. Deactivating buckets")
+
         logging.info(
             "ShardedDDP bucket size: {:.2f}M parameters, model size {:.2f}M parameters".format(
                 self.buffer_max_size / 2 ** 20, model_size / 2 ** 20


Test patch:
diff --git a/tests/nn/data_parallel/test_sharded_ddp_features.py b/tests/nn/data_parallel/test_sharded_ddp_features.py
--- a/tests/nn/data_parallel/test_sharded_ddp_features.py
+++ b/tests/nn/data_parallel/test_sharded_ddp_features.py
@@ -408,9 +408,9 @@ def closure():
 
 @skip_if_no_cuda
 @skip_if_single_gpu
-def test_gpt2():
-    # Check that the ShardedDDP wrapper accepts tuple(tensors) as inputs
-    world_size = 2
+@pytest.mark.parametrize("world_size", [1, 2])
+def test_gpt2(world_size):
+    # Check that having trainable unused params is fine
     backend = "gloo"
     temp_file_name = tempfile.mkstemp()[1]
     device = "cuda"