diff --git a/.github/workflows/pyre.yml b/.github/workflows/pyre.yml index be99ff62c..ed25404e7 100644 --- a/.github/workflows/pyre.yml +++ b/.github/workflows/pyre.yml @@ -19,7 +19,7 @@ jobs: uses: actions/checkout@v2 - name: Install dependencies run: > - conda install --yes pytorch cpuonly -c pytorch-nightly && + pip install torch --index-url https://download.pytorch.org/whl/nightly/cpu && pip install fbgemm-gpu --index-url https://download.pytorch.org/whl/nightly/cpu && pip install -r requirements.txt && pip install pyre-check-nightly==$(cat .pyre_configuration | grep version | awk '{print $2}' | sed 's/\"//g') diff --git a/.github/workflows/unittest_ci.yml b/.github/workflows/unittest_ci.yml index 8865acee4..0b3fe0638 100644 --- a/.github/workflows/unittest_ci.yml +++ b/.github/workflows/unittest_ci.yml @@ -73,13 +73,15 @@ jobs: conda info python --version conda run -n build_binary python --version - conda install -n build_binary \ - --yes \ - pytorch pytorch-cuda=11.8 -c pytorch-nightly -c nvidia + conda run -n build_binary \ + pip install torch --index-url https://download.pytorch.org/whl/nightly/${{ matrix.cuda-tag }} + conda run -n build_binary \ + python -c "import torch" + echo "torch succeeded" conda run -n build_binary \ python -c "import torch.distributed" conda run -n build_binary \ - pip install fbgemm-gpu --index-url https://download.pytorch.org/whl/nightly/cu118 + pip install fbgemm-gpu --index-url https://download.pytorch.org/whl/nightly/${{ matrix.cuda-tag }} conda run -n build_binary \ python -c "import fbgemm_gpu" echo "fbgemm_gpu succeeded" diff --git a/.github/workflows/unittest_ci_cpu.yml b/.github/workflows/unittest_ci_cpu.yml index 2861029e1..1efe64178 100644 --- a/.github/workflows/unittest_ci_cpu.yml +++ b/.github/workflows/unittest_ci_cpu.yml @@ -45,9 +45,11 @@ jobs: conda info python --version conda run -n build_binary python --version - conda install -n build_binary \ - --yes \ - pytorch cpuonly -c pytorch-nightly + conda run -n build_binary \ + pip install torch --index-url https://download.pytorch.org/whl/nightly/cpu + conda run -n build_binary \ + python -c "import torch" + echo "torch succeeded" conda run -n build_binary \ python -c "import torch.distributed" conda run -n build_binary \ diff --git a/torchrec/distributed/model_parallel.py b/torchrec/distributed/model_parallel.py index 0f60362b6..5cbd2429b 100644 --- a/torchrec/distributed/model_parallel.py +++ b/torchrec/distributed/model_parallel.py @@ -746,6 +746,7 @@ def sync(self, include_optimizer_state: bool = True) -> None: all_weights = [ w for emb_kernel in self._modules_to_sync + # pyre-fixme[29]: `Union[Module, Tensor]` is not a function. for w in emb_kernel.split_embedding_weights() ] handle = self._replica_pg.allreduce_coalesced(all_weights, opts=opts) @@ -755,6 +756,7 @@ def sync(self, include_optimizer_state: bool = True) -> None: # Sync accumulated square of grad of local optimizer shards optim_list = [] for emb_kernel in self._modules_to_sync: + # pyre-fixme[29]: `Union[Module, Tensor]` is not a function. all_optimizer_states = emb_kernel.get_optimizer_state() momentum1 = [optim["sum"] for optim in all_optimizer_states] optim_list.extend(momentum1) @@ -864,6 +866,8 @@ def _find_sharded_modules( if isinstance(module, SplitTableBatchedEmbeddingBagsCodegen): sharded_modules.append(module) if hasattr(module, "_lookups"): + # pyre-fixme[29]: `Union[(self: Tensor) -> Any, Module, Tensor]` is + # not a function. for lookup in module._lookups: _find_sharded_modules(lookup) return