Introduce base collective and main subclasses

src/lightning_lite/plugins/__init__.py

@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
 from lightning_lite.plugins.environments.cluster_environment import ClusterEnvironment
 from lightning_lite.plugins.io.checkpoint_plugin import CheckpointIO
 from lightning_lite.plugins.io.torch_plugin import TorchCheckpointIO
@@ -28,10 +27,10 @@
     "CheckpointIO",
     "TorchCheckpointIO",
     "XLACheckpointIO",
+    "Precision",
     "DeepSpeedPrecision",
     "DoublePrecision",
     "NativeMixedPrecision",
-    "Precision",
     "TPUPrecision",
     "TPUBf16Precision",
 ]

src/lightning_lite/plugins/collectives/__init__.py

@@ -0,0 +1,9 @@
+from lightning_lite.plugins.collectives.collective import Collective
+from lightning_lite.plugins.collectives.single_device_collective import SingleDeviceCollective
+from lightning_lite.plugins.collectives.torch_collective import TorchCollective
+
+__all__ = [
+    "Collective",
+    "TorchCollective",
+    "SingleDeviceCollective",
+]

src/lightning_lite/plugins/collectives/collective.py

@@ -0,0 +1,146 @@
+from abc import ABC, abstractmethod
+from typing import Any, List, Optional
+
+import torch
+from typing_extensions import Self
+
+from lightning_lite.utilities.types import CollectibleGroup
+
+
+class Collective(ABC):
+    def __init__(self, instantiate_group: bool = False, **group_kwargs: Any) -> None:
+        self._group_kwargs = group_kwargs
+        self._group: Optional[CollectibleGroup] = None
+        if instantiate_group:
+            self.create_group()
+
+    def create_group(self, **kwargs: Any) -> Self:  # type: ignore[valid-type]
+        if self._group is not None:
+            raise RuntimeError(f"{type(self).__name__} already owns a group.")
+        self._group_kwargs.update(kwargs)
+        self._group = self.init_group(**self._group_kwargs)
+        return self
+
+    @property
+    def group(self) -> CollectibleGroup:
+        if self._group is None:
+            raise RuntimeError(
+                f"{type(self).__name__} does not own a group. HINT: try `collective.create_group().group`"
+            )
+        return self._group
+
+    @property
+    @abstractmethod
+    def rank(self) -> int:
+        pass
+
+    @property
+    @abstractmethod
+    def world_size(self) -> int:
+        pass
+
+    @staticmethod
+    @abstractmethod
+    def init_group(
+        **kwargs: Any,
+    ) -> CollectibleGroup:
+        pass
+
+    def teardown(self) -> None:
+        if self._group is None:
+            raise RuntimeError(f"{type(self).__name__} does not own a group to destroy.")
+        self.destroy_group(self._group)
+        self._group = None
+
+    @staticmethod
+    @abstractmethod
+    def destroy_group(group: CollectibleGroup) -> None:
+        pass
+
+    @staticmethod
+    @abstractmethod
+    def _convert_to_native_op(op: str) -> Any:
+        pass
+
+    @abstractmethod
+    def send(self, tensor: torch.Tensor, dst: int, tag: Optional[int] = 0) -> None:
+        pass
+
+    @abstractmethod
+    def recv(self, tensor: torch.Tensor, src: Optional[int] = None, tag: Optional[int] = 0) -> torch.Tensor:
+        pass
+
+    @abstractmethod
+    def broadcast(
+        self,
+        tensor: torch.Tensor,
+        src: int,
+    ) -> torch.Tensor:
+        pass
+
+    @abstractmethod
+    def all_reduce(
+        self,
+        tensor: torch.Tensor,
+        op: str,
+    ) -> torch.Tensor:
+        pass
+
+    @abstractmethod
+    def reduce(
+        self,
+        tensor: torch.Tensor,
+        dst: int,
+        op: str,
+    ) -> torch.Tensor:
+        pass
+
+    @abstractmethod
+    def all_gather(
+        self,
+        tensor_list: List[torch.Tensor],
+        tensor: torch.Tensor,
+    ) -> List[torch.Tensor]:
+        pass
+
+    @abstractmethod
+    def gather(
+        self,
+        tensor: torch.Tensor,
+        gather_list: Optional[List[torch.Tensor]] = None,
+        dst: int = 0,
+    ) -> Optional[List[torch.Tensor]]:
+        pass
+
+    @abstractmethod
+    def scatter(
+        self,
+        tensor: torch.Tensor,
+        scatter_list: Optional[List[torch.Tensor]] = None,
+        src: int = 0,
+    ) -> torch.Tensor:
+        pass
+
+    @abstractmethod
+    def reduce_scatter(
+        self,
+        output: torch.Tensor,
+        input_list: List[torch.Tensor],
+        op: str,
+    ) -> torch.Tensor:
+        pass
+
+    @abstractmethod
+    def all_to_all(
+        self,
+        output_tensor_list: List[torch.Tensor],
+        input_tensor_list: List[torch.Tensor],
+    ) -> List[torch.Tensor]:
+        pass
+
+    @abstractmethod
+    def barrier(
+        self,
+        device_ids: Optional[List[int]] = None,
+    ) -> None:
+        pass

src/lightning_lite/plugins/collectives/single_device_collective.py

@@ -0,0 +1,110 @@
+from typing import Any, List, Optional
+
+import torch
+
+from lightning_lite.plugins.collectives.collective import Collective
+from lightning_lite.utilities.types import CollectibleGroup
+
+
+class SingleDeviceCollective(Collective):
+    @property
+    def rank(self) -> int:
+        return 0
+
+    @property
+    def world_size(self) -> int:
+        return 1
+
+    @staticmethod
+    def init_group(
+        **kwargs: Any,
+    ) -> CollectibleGroup:
+        return object()  # type: ignore[return-value]
+
+    @staticmethod
+    def _convert_to_native_op(op: str) -> str:
+        return op
+
+    @staticmethod
+    def destroy_group(group: CollectibleGroup) -> None:
+        pass
+
+    def send(self, *_: Any, **__: Any) -> None:
+        pass
+
+    def recv(self, tensor: torch.Tensor, *_: Any, **__: Any) -> torch.Tensor:
+        return tensor
+
+    def broadcast(
+        self,
+        tensor: torch.Tensor,
+        *_: Any,
+        **__: Any,
+    ) -> torch.Tensor:
+        return tensor
+
+    def all_reduce(
+        self,
+        tensor: torch.Tensor,
+        *_: Any,
+        **__: Any,
+    ) -> torch.Tensor:
+        return tensor
+
+    def reduce(
+        self,
+        tensor: torch.Tensor,
+        *_: Any,
+        **__: Any,
+    ) -> torch.Tensor:
+        return tensor
+
+    def all_gather(
+        self,
+        tensor_list: List[torch.Tensor],
+        tensor: torch.Tensor,
+        **__: Any,
+    ) -> List[torch.Tensor]:
+        return [tensor]
+
+    def gather(
+        self,
+        tensor: torch.Tensor,
+        *_: Any,
+        **__: Any,
+    ) -> Optional[List[torch.Tensor]]:
+        return [tensor]
+
+    def scatter(  # type: ignore[override]
+        self,
+        tensor: torch.Tensor,
+        scatter_list: List[torch.Tensor],  # it doesn't make sense to have a None here for a single device
+        *_: Any,
+        **__: Any,
+    ) -> torch.Tensor:
+        return scatter_list[0]
+
+    def reduce_scatter(
+        self,
+        output: torch.Tensor,
+        input_list: List[torch.Tensor],
+        *_: Any,
+        **__: Any,
+    ) -> torch.Tensor:
+        return input_list[0]
+
+    def all_to_all(
+        self,
+        output_tensor_list: List[torch.Tensor],
+        input_tensor_list: List[torch.Tensor],
+        *_: Any,
+        **__: Any,
+    ) -> List[torch.Tensor]:
+        return input_tensor_list
+
+    def barrier(
+        self,
+        *_: Any,
+        **__: Any,
+    ) -> None:
+        pass