fix(fabric): raise on CPU tensor passed to all_reduce in non-CPU setup (#21530)

Author: GudeJunge

src/lightning/fabric/fabric.py

@@ -701,8 +701,22 @@ def all_reduce(
             A tensor of the same shape as the input with values reduced pointwise across processes. The same is
             applied to tensors in a collection if a collection is given as input.
 
+        Raises:
+            RuntimeError: If a CPU tensor is passed while running on a non-CPU device. Move the tensor to
+                ``fabric.device`` before calling ``all_reduce``.
+
         """
         self._validate_launched()
+        if self.device.type != "cpu":
+
+            def _validate_tensor_device(tensor: Tensor) -> None:
+                if tensor.device.type == "cpu":
+                    raise RuntimeError(
+                        "`Fabric.all_reduce` received a CPU tensor while running on a non-CPU device. Move the tensor"
+                        " to fabric.device before calling all_reduce."
+                    )
+
+            apply_to_collection(data, Tensor, _validate_tensor_device)
         group = group if group is not None else torch.distributed.group.WORLD
         data = convert_to_tensors(data, device=self.device)
         return apply_to_collection(data, Tensor, self._strategy.all_reduce, group=group, reduce_op=reduce_op)

tests/tests_fabric/test_fabric.py

@@ -1214,6 +1214,17 @@ def test_all_reduce():
     fabric._strategy.all_reduce.assert_has_calls([call(torch.tensor(4), **defaults), call(torch.tensor(5), **defaults)])
 
 
+def test_all_reduce_raises_for_cpu_tensor_on_non_cpu_root_device():
+    """Test that `Fabric.all_reduce()` raises an error when a CPU tensor is passed while the root device is not CPU."""
+    fabric = Fabric()
+    fabric._strategy = Mock(root_device=torch.device("cuda", 0))
+    fabric._launched = True
+
+    tensor_cpu = torch.tensor(1.0, device="cpu")
+    with pytest.raises(RuntimeError, match=r"Move the tensor to fabric.device before calling all_reduce"):
+        fabric.all_reduce(tensor_cpu)
+
+
 def test_rank_zero_first(monkeypatch):
     """Test that rank 0 completes first before all other processes can execute under `.rank_zero_first()`."""