Repository: lightly-ai/lightly

Issue:
Bug in `GatherLayer.backward`
Hi,

We've been implementing a model at [cellarium-ml](https://github.com/cellarium-ai/cellarium-ml) using your `NTXentLoss`. Comparing the model training with a single GPU and two GPUs we noticed that they do not match. By investigating it we found an apparent bug in the `GatherLayer.backward` where gradients are not sum-reduced over GPUs. Here is our fixed version (https://github.com/cellarium-ai/cellarium-ml/blob/main/cellarium/ml/distributed/gather.py#L17-L21):

```py
    @staticmethod
    def backward(ctx, *grads) -> torch.Tensor:
        grad_out = grads[dist.get_rank()].contiguous()
        dist.all_reduce(grad_out, op=dist.ReduceOp.SUM)
        return grad_out
```

and the [test](https://github.com/cellarium-ai/cellarium-ml/blob/main/tests/distributed/test_gather.py) we wrote. Would you agree that this is indeed a bug? I would be happy to contribute a PR with the fix.


Discussion:
Hi, thanks a lot for reporting the issue!

There seems to be indeed something wrong with our `GatherLayer` implementation. I was able to reproduce the problem and will look more into it in the coming days. A PR with the fix would be very welcome!

Btw. there might be something off with your test. The comments in the code indicate:
```
    # loss_0 = w_0 + w_1
    # loss_1 = w_0 + w_1

    # dloss/dw_0 = dloss_0/dw_0 + dloss_1/dw_0 = 2
    # dloss/dw_1 = dloss_0/dw_1 + dloss_1/dw_1 = 2
```

But AFAIK `loss = w_0 + w_1` and not `loss = loss_0 + loss_1` which means that `dloss/dw_0 = dw_0/dw_0 + dw_1/dw_0 = 1` and not `2`. 
Hi, thanks for responding!

> But AFAIK loss = w_0 + w_1 and not loss = loss_0 + loss_1 which means that dloss/dw_0 = dw_0/dw_0 + dw_1/dw_0 = 1 and not 2.

The point I was trying to make in this test is that `w_0` contributes to both `loss_0` and `loss_1`. Therefore, we need to accumulate gradients from both losses. Similarly, `w_1` contributes to both `loss_0` and `loss_1`. And since `w_0` and `w_1` are equivalent to a `w` in a single GPU case, their gradients should be the same (`w_0.grad == w_1.grad == w.grad == 2`).

Resolving patch:
diff --git a/lightly/utils/dist.py b/lightly/utils/dist.py
--- a/lightly/utils/dist.py
+++ b/lightly/utils/dist.py
@@ -1,29 +1,29 @@
-from typing import Optional, Tuple
+from typing import Any, Callable, Optional, Tuple, TypeVar
 
 import torch
 import torch.distributed as dist
+from torch.autograd.function import FunctionCtx
 
 
 class GatherLayer(torch.autograd.Function):
     """Gather tensors from all processes, supporting backward propagation.
 
     This code was taken and adapted from here:
-    https://github.com/Spijkervet/SimCLR
+    https://github.com/vturrisi/solo-learn/blob/b69b4bd27472593919956d9ac58902a301537a4d/solo/utils/misc.py#L187
 
     """
 
     @staticmethod
-    def forward(ctx, input: torch.Tensor) -> Tuple[torch.Tensor, ...]:
-        ctx.save_for_backward(input)
+    def forward(ctx, input: torch.Tensor) -> Tuple[torch.Tensor, ...]:  # type: ignore
         output = [torch.empty_like(input) for _ in range(dist.get_world_size())]
         dist.all_gather(output, input)
         return tuple(output)
 
     @staticmethod
-    def backward(ctx, *grads: torch.Tensor) -> torch.Tensor:
-        (input,) = ctx.saved_tensors
-        grad_out = torch.empty_like(input)
-        grad_out[:] = grads[dist.get_rank()]
+    def backward(ctx, *grads) -> torch.Tensor:  # type: ignore
+        all_gradients = torch.stack(grads)
+        dist.all_reduce(all_gradients)
+        grad_out = all_gradients[dist.get_rank()]
         return grad_out
 
 
@@ -39,7 +39,7 @@ def world_size() -> int:
 
 def gather(input: torch.Tensor) -> Tuple[torch.Tensor]:
     """Gathers this tensor from all processes. Supports backprop."""
-    return GatherLayer.apply(input)
+    return GatherLayer.apply(input)  # type: ignore[no-any-return]
 
 
 def eye_rank(n: int, device: Optional[torch.device] = None) -> torch.Tensor:
@@ -70,7 +70,10 @@ def eye_rank(n: int, device: Optional[torch.device] = None) -> torch.Tensor:
     return diag_mask
 
 
-def rank_zero_only(fn):
+R = TypeVar("R")
+
+
+def rank_zero_only(fn: Callable[..., R]) -> Callable[..., Optional[R]]:
     """Decorator that only runs the function on the process with rank 0.
 
     Example:
@@ -79,17 +82,17 @@ def rank_zero_only(fn):
         >>>     print(message)
         >>>
         >>> print_rank_zero("Hello from rank 0!")
-
     """
 
-    def wrapped(*args, **kwargs):
+    def wrapped(*args: Any, **kwargs: Any) -> Optional[R]:
         if rank() == 0:
             return fn(*args, **kwargs)
+        return None
 
     return wrapped
 
 
 @rank_zero_only
-def print_rank_zero(*args, **kwargs) -> None:
+def print_rank_zero(*args: Any, **kwargs: Any) -> None:  # type: ignore[misc]
     """Equivalent to print, but only runs on the process with rank 0."""
     print(*args, **kwargs)


Test patch:
diff --git a/tests/utils/test_dist__gather.py b/tests/utils/test_dist__gather.py
new file mode 100644
--- /dev/null
+++ b/tests/utils/test_dist__gather.py
@@ -0,0 +1,112 @@
+from typing import Any, Generator, Tuple
+
+import pytest
+import pytorch_lightning as pl
+import torch
+import torch.distributed
+from pytorch_lightning import LightningModule, Trainer
+from pytorch_lightning.strategies.ddp import DDPStrategy
+from torch import Tensor
+from torch.nn import Linear, MSELoss
+from torch.optim import SGD
+from torch.utils.data import TensorDataset
+
+from lightly.loss.ntx_ent_loss import NTXentLoss
+from lightly.utils import dist as lightly_dist
+
+"""
+WARNING: 
+Using a DDPStrategy causes the file to be executed once per device.
+Thus this test needs to be in a separate file.
+"""
+
+
+class Model(LightningModule):
+    def __init__(self, gather: bool):
+        super().__init__()
+        self.model = Linear(5, 2, bias=False)
+        self.model.weight.data = torch.Tensor(
+            [
+                [0.1, 0.2, 0.3, 0.4, 0.5],
+                [0.6, 0.7, 0.8, 0.9, 1.0],
+            ]
+        )
+        self.criterion = NTXentLoss(gather_distributed=gather)
+        self.gather = gather
+
+    def training_step(self, batch: Tuple[Tensor, Tensor], batch_idx: int) -> Tensor:
+        x = batch[0]
+        y = batch[1]
+        x = self.model(x).flatten(start_dim=1)
+        y = y.flatten(start_dim=1)
+        loss: Tensor = self.criterion(x, y)
+        return loss
+
+    def configure_optimizers(self) -> Any:
+        return SGD(self.parameters(), lr=0.01)
+
+
+@pytest.fixture
+def close_torch_distributed() -> Generator[None, None, None]:
+    yield None
+    torch.distributed.destroy_process_group()
+
+
+class TestGatherLayer:
+    def test(self, close_torch_distributed: None) -> None:
+        """
+        Tests that the gather layer works as expected.
+
+        Because using a DDPStrategy causes the whole script to be executed multiple
+        times, running a proper test is difficult. The approach used here:
+
+        1. This test was run once with n_devices=1 and gather=False. The resulting
+        parameters after 10 epochs are saved as expected_params__10_epochs__no_gather.
+
+        2. This test is now run with n_devices=2 and gather=True. The resulting
+        parameters are asserted to be the same ad with n_devices=1 and gather=False.
+
+        Note that the results would not be the same for n_devices=1 and n_devices=2 if
+        there was:
+        - Any randomness in the transform, as the order the samples are processed and
+          thus the random seed differ.
+        - Any batch normalization in the model, as the batch size is split between
+          the devices and there is no gather layer for the batch normalization.
+        """
+        n_devices = 2
+        n_samples = 8
+        batch_size = int(n_samples / n_devices)
+        gather = n_devices > 1
+
+        pl.seed_everything(0, workers=True)
+
+        xs = torch.arange(n_samples * 2 * 5).reshape(n_samples, 2, 5).float()
+        ys = torch.arange(n_samples * 2 * 2).reshape(n_samples, 2, 2).float()
+        dataset = TensorDataset(xs, ys)
+
+        model = Model(gather=gather)
+
+        dataloader = torch.utils.data.DataLoader(
+            dataset,
+            batch_size=batch_size,
+            shuffle=True,
+            drop_last=True,
+            num_workers=0,
+        )
+
+        trainer = Trainer(
+            devices=n_devices,
+            accelerator="cpu",
+            strategy=DDPStrategy(find_unused_parameters=False),
+            max_epochs=10,
+        )
+        trainer.fit(model=model, train_dataloaders=dataloader)
+
+        params = list(model.parameters())[0]
+        expected_params__10_epochs__no_gather = torch.Tensor(
+            [
+                [0.1076, 0.2076, 0.3077, 0.4077, 0.5078],
+                [0.5976, 0.6973, 0.7971, 0.8969, 0.9967],
+            ]
+        )
+        assert torch.allclose(params, expected_params__10_epochs__no_gather, rtol=1e-2)