Repository: Lightning-AI/torchmetrics

Issue:
List Metric synchronization fails in corner case
## 🐛 Bug

Empty list metrics bypass a call to `torch.barrier()` and cause deadlocks or other strange synchronization issues.

### To Reproduce

Define the `ListMetric` class as follows:

```py
import torch
import torchmetrics

class ListMetric(torchmetrics.Metric):
     full_state_update: bool = False
     
     def __init__(self):
         super().__init__(sync_on_compute=False)
         
        self.add_state("example_state", default=[], dist_reduce_fx="cat")
        
    def update(x: torch.Tensor):
        # in this example, x is a (2,)-shaped tensor and the internal state of this metric is a list of tensors of this shape
        self.example_state.append(x)
        
    def compute(x: torch.Tensor):
        return self.example_state

    def sync(
        self,
        dist_sync_fn: Optional[Callable] = None,
        process_group: Optional[Any] = None,
        should_sync: bool = True,
        distributed_available: Optional[Callable] = None,
    ) -> None:
        super().sync(
            dist_sync_fn=dist_sync_fn,
            process_group=process_group,
            should_sync=should_sync,
            distributed_available=distributed_available,
        )

        tensor_and_empty = (
            isinstance(self.example_state, torch.Tensor)
            and torch.numel(self.example_state) == 0
        )
        list_and_empty = (
            isinstance(self.example_state, List) and len(self.example_state) == 0
        )
        if tensor_and_empty or list_and_empty:
            self.example_state = []
            return
        """
        Torchmetrics has a strange quirk:
        Depending on whether it's being run in a distributed setting, the type of 'self.resolutions' may differ.
        The `dim_zero_cat()` function fixes this.

        see https://lightning.ai/docs/torchmetrics/v1.3.1/pages/implement.html "working with list states"
        """
        self.example_state = torchmetrics.utilities.dim_zero_cat(self.example_state)
        self.example_state = list(torch.split(self.example_state, 2))
```

In your training loop, do this with 2 or more GPUs:

```py
def on_validation_epoch_end(self, outputs):
    lm = ListMetric()
    if self.local_rank == 0:
        lm.update(torch.zeros((2,)))
    lm.sync()  # everything but local_rank == 0 will skip the barrier inside this call to `sync()` and then the script will deadlock
```

### Expected behavior

Empty list metrics properly sync (waiting on a barrier) with non-empty list metrics on different GPUs

### Environment

Linux, torchmetrics 1.3.0.post0 from pip install

### Additional context

It seems like this is caused by the line `if reduction_fn == dim_zero_cat and isinstance(input_dict[attr], list) and len(input_dict[attr]) > 1:` in `Metric._sync_dist()`, which causes `return [function(x, *args, **kwargs) for x in data]` in `lightning_utilities.core.apply_func.py` to execute. Since `function` contains the blocking call to wait on the barrier and the list is empty, no blocking ever occurs for GPUs whose metric states are an empty list.



Discussion:
Hi! thanks for your contribution!, great first issue!

Resolving patch:
diff --git a/src/torchmetrics/metric.py b/src/torchmetrics/metric.py
--- a/src/torchmetrics/metric.py
+++ b/src/torchmetrics/metric.py
@@ -38,6 +38,7 @@
 )
 from torchmetrics.utilities.distributed import gather_all_tensors
 from torchmetrics.utilities.exceptions import TorchMetricsUserError
+from torchmetrics.utilities.imports import _TORCH_GREATER_EQUAL_2_1
 from torchmetrics.utilities.plot import _AX_TYPE, _PLOT_OUT_TYPE, plot_single_or_multi_val
 from torchmetrics.utilities.prints import rank_zero_warn
 
@@ -438,6 +439,15 @@ def _sync_dist(self, dist_sync_fn: Callable = gather_all_tensors, process_group:
             if reduction_fn == dim_zero_cat and isinstance(input_dict[attr], list) and len(input_dict[attr]) > 1:
                 input_dict[attr] = [dim_zero_cat(input_dict[attr])]
 
+            # cornor case in distributed settings where a rank have not received any data, create empty to concatenate
+            if (
+                _TORCH_GREATER_EQUAL_2_1
+                and reduction_fn == dim_zero_cat
+                and isinstance(input_dict[attr], list)
+                and len(input_dict[attr]) == 0
+            ):
+                input_dict[attr] = [torch.tensor([], device=self.device, dtype=self.dtype)]
+
         output_dict = apply_to_collection(
             input_dict,
             Tensor,


Test patch:
diff --git a/tests/unittests/bases/test_ddp.py b/tests/unittests/bases/test_ddp.py
--- a/tests/unittests/bases/test_ddp.py
+++ b/tests/unittests/bases/test_ddp.py
@@ -22,6 +22,7 @@
 from torchmetrics import Metric
 from torchmetrics.utilities.distributed import gather_all_tensors
 from torchmetrics.utilities.exceptions import TorchMetricsUserError
+from torchmetrics.utilities.imports import _TORCH_GREATER_EQUAL_2_1
 
 from unittests import NUM_PROCESSES
 from unittests._helpers import seed_all
@@ -269,11 +270,28 @@ def test_sync_on_compute(sync_on_compute, test_func):
 def _test_sync_with_empty_lists(rank):
     dummy = DummyListMetric()
     val = dummy.compute()
-    assert val == []
+    assert torch.allclose(val, tensor([]))
 
 
 @pytest.mark.DDP()
+@pytest.mark.skipif(not _TORCH_GREATER_EQUAL_2_1, reason="test only works on newer torch versions")
 @pytest.mark.skipif(sys.platform == "win32", reason="DDP not available on windows")
 def test_sync_with_empty_lists():
     """Test that synchronization of states can be enabled and disabled for compute."""
     pytest.pool.map(_test_sync_with_empty_lists, range(NUM_PROCESSES))
+
+
+def _test_sync_with_unequal_size_lists(rank):
+    """Test that synchronization of list states work even when some ranks have not received any data yet."""
+    dummy = DummyListMetric()
+    if rank == 0:
+        dummy.update(torch.zeros(2))
+    assert torch.all(dummy.compute() == tensor([0.0, 0.0]))
+
+
+@pytest.mark.DDP()
+@pytest.mark.skipif(not _TORCH_GREATER_EQUAL_2_1, reason="test only works on newer torch versions")
+@pytest.mark.skipif(sys.platform == "win32", reason="DDP not available on windows")
+def test_sync_with_unequal_size_lists():
+    """Test that synchronization of states can be enabled and disabled for compute."""
+    pytest.pool.map(_test_sync_with_unequal_size_lists, range(NUM_PROCESSES))