2/n Consolidate collective functions - collective base and subclasses

Author: four4fish

CHANGELOG.md

@@ -116,6 +116,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added `inference_mode` for evaluation and prediction ([8813](https://github.com/PyTorchLightning/pytorch-lightning/pull/8813))
 
 
+- Add collective base class and subclasses ([#9414](https://github.com/PyTorchLightning/pytorch-lightning/pull/9414))
+
+
 ### Changed
 
 - Parsing of the `gpus` Trainer argument has changed: `gpus="n"` (str) no longer selects the GPU index n and instead selects the first n devices. ([#8770](https://github.com/PyTorchLightning/pytorch-lightning/pull/8770))

pytorch_lightning/plugins/collective/horovod_collective.py

@@ -32,7 +32,7 @@
 
 
 class HorovodCollective(Collective):
-    """Base class for collective functions for training type plugins."""
+    """Collective interface for Horovod training type plugins."""
 
     def __init__(
         self,

pytorch_lightning/plugins/collective/single_device_collective.py

@@ -19,7 +19,7 @@
 
 
 class SingleNodeCollective(Collective):
-    """Base class for collective functions for training type plugins."""
+    """Collective interface for single device training type plugins."""
 
     def barrier(self, name: Optional[str] = None, *args, **kwargs) -> None:
         """Forces all possibly joined processes to wait for each other."""

pytorch_lightning/plugins/collective/tpu_collective.py

@@ -29,7 +29,7 @@
 
 
 class TPUCollective(Collective):
-    """Base class for collective functions for training type plugins."""
+    """Collective interface for TPU and TPUSpawning training type plugins."""
 
     def barrier(self, name: Optional[str] = None) -> None:
         if self.is_distributed: