integrate torchcomms (#290)

Summary: Pull Request resolved: #290

Differential Revision: D86343575

Author: tushar00jain

.github/workflows/docs.yaml

@@ -25,7 +25,7 @@ jobs:
 
           sudo apt-get install -y protobuf-compiler
 
-          pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128
+          pip install --pre torch torchvision torchaudio torchcomms --index-url https://download.pytorch.org/whl/nightly/cu128
           pip install .[dev] -v
 
           pip install -r docs/requirements.txt

.github/workflows/lint.yaml

@@ -23,7 +23,7 @@ jobs:
 
           sudo apt-get install -y protobuf-compiler
 
-          pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128
+          pip install --pre torch torchvision torchaudio torchcomms --index-url https://download.pytorch.org/whl/nightly/cu128
           pip install .[dev] -v
 
           # install recent version of Rust via rustup

.github/workflows/unittest-mac.yaml

@@ -12,7 +12,7 @@ jobs:
     steps:
       - name: Checkout
         uses: actions/checkout@v4
-        
+
       - name: Setup miniconda
         uses: pytorch/test-infra/.github/actions/setup-miniconda@main
         with:
@@ -39,7 +39,7 @@ jobs:
 
           python -m pip install --upgrade pip
 
-          pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cpu
+          pip install --pre torch torchvision torchaudio torchcomms --index-url https://download.pytorch.org/whl/nightly/cpu
 
           pip install -e .[dev] -v
 

.github/workflows/unittest.yaml

@@ -41,10 +41,10 @@ jobs:
 
         # Optionally install torch nightly, pulls latest CUDA from pip otherwise
         if [ "${{ matrix.torch-version }}" = "nightly" ]; then
-          pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128
+          pip install --pre torch torchvision torchaudio torchcomms --index-url https://download.pytorch.org/whl/nightly/cu128
         fi
         if [ "${{ matrix.torch-version }}" = "test" ]; then
-          pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/test/cu128
+          pip install --pre torch torchvision torchaudio torchcomms --index-url https://download.pytorch.org/whl/test/cu128
         fi
 
         # Install dependencies

torchft/__init__.py

@@ -16,6 +16,7 @@
     ProcessGroupNCCL,
     ProcessGroupXCCL,
 )
+from torchft.torchcomms import TorchCommGloo, TorchCommNCCL
 
 setup_logger("torchft_quorums")
 setup_logger("torchft_commits")
@@ -31,4 +32,6 @@
     "ProcessGroupBabyNCCL",
     "ProcessGroupBabyXCCL",
     "ProcessGroupGloo",
+    "TorchCommNCCL",
+    "TorchCommGloo",
 )

torchft/manager.py

@@ -51,6 +51,7 @@
 
 import torch
 import torch.distributed as dist
+import torchcomms
 from torch.distributed import ReduceOp, TCPStore
 from torch.distributed.distributed_c10d import AllreduceOptions, ReduceOp, Work
 
@@ -63,6 +64,7 @@
 
 if TYPE_CHECKING:
     from torchft.process_group import ProcessGroup
+    from torchft.torchcomms import TorchComm
 
 IS_TRITON_AVAILABLE = True
 try:
@@ -163,7 +165,7 @@ class Manager:
 
     def __init__(
         self,
-        pg: "ProcessGroup",
+        pg: Union["ProcessGroup", "TorchComm"],
         load_state_dict: Optional[Callable[[T], None]],
         state_dict: Optional[Callable[[], T]],
         min_replica_size: int,
@@ -188,6 +190,7 @@ def __init__(
     ) -> None:
         """
         Args:
+            pg: process group or torchcomms wrapper to use for communication.
             load_state_dict: function to load the state dict when recovering
             state_dict: function to save the state dict with recovering
             min_replica_size: minimum number of replicas on each step
@@ -456,7 +459,9 @@ def allreduce(
         try:
             # Run the allreduce async and save the work object so we can wait on
             # it later.
+            # TODO: Support quantization with torchcomms
             if should_quantize and IS_TRITON_AVAILABLE:
+                assert isinstance(self._pg, ProcessGroup)
                 work = allreduce_quantized(
                     [tensor],
                     pg_reduce_op,
@@ -465,9 +470,21 @@ def allreduce(
                     torch.accelerator.current_stream(),
                 )
             else:
-                opts = AllreduceOptions()
-                opts.reduceOp = pg_reduce_op
-                work = self._pg.allreduce([tensor], opts)
+                # Check if we're using torchcomms or ProcessGroup
+                if isinstance(self._pg, TorchComm):
+                    # Convert PyTorch ReduceOp to torchcomms ReduceOp
+                    if pg_reduce_op == ReduceOp.SUM:
+                        tc_op = torchcomms.ReduceOp.SUM
+                    elif pg_reduce_op == ReduceOp.AVG:
+                        tc_op = torchcomms.ReduceOp.AVG
+                    else:
+                        raise AssertionError("unsupported reduce op")
+
+                    work = self._pg.allreduce(tensor, tc_op)
+                else:
+                    opts = AllreduceOptions()
+                    opts.reduceOp = pg_reduce_op
+                    work = self._pg.allreduce([tensor], opts)
 
             # schedule grad normalization as a continuation
             # on the Future

torchft/process_group.py

@@ -1331,6 +1331,7 @@ class ManagedProcessGroup(ProcessGroupWrapper):
     """
 
     def __init__(self, manager: "Manager") -> None:
+        assert isinstance(manager._pg, ProcessGroup)
         super().__init__(pg=manager._pg)
 
         self._manager = manager
@@ -1350,6 +1351,7 @@ def size(self) -> int:
         return self._manager.num_participants()
 
     def getBackendName(self) -> str:
+        assert isinstance(self._manager._pg, ProcessGroup)
         return self._manager._pg.getBackendName()