integrate torchcomms

.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

.pyre_configuration

@@ -13,6 +13,9 @@
     },
     {
       "site-package": "parameterized"
+    },
+    {
+      "site-package": "torchcomms"
     }
   ]
 }

torchft/__init__.py

@@ -4,6 +4,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+from torchft.comms import TorchCommGloo, TorchCommNCCL
 from torchft.data import DistributedSampler
 from torchft.ddp import DistributedDataParallel
 from torchft.manager import Manager
@@ -31,4 +32,6 @@
     "ProcessGroupBabyNCCL",
     "ProcessGroupBabyXCCL",
     "ProcessGroupGloo",
+    "TorchCommNCCL",
+    "TorchCommGloo",
 )