suppress errors in torchrec (#3016)

Summary: Pull Request resolved: #3016

Differential Revision: D75570790

fbshipit-source-id: c3200326178f2872a45dada80b7f4ec551321729

Author: generatedunixname89002005307016

torchrec/datasets/utils.py

@@ -43,6 +43,7 @@ def to(self, device: torch.device, non_blocking: bool = False) -> "Batch":
 
     def record_stream(self, stream: torch.Stream) -> None:
         self.dense_features.record_stream(stream)
+        # pyre-fixme[6]: For 1st argument expected `Stream` but got `Stream`.
         self.sparse_features.record_stream(stream)
         self.labels.record_stream(stream)
 

torchrec/distributed/batched_embedding_kernel.py

@@ -840,6 +840,8 @@ def _gen_named_parameters_by_table_ssd_pmt(
     for table_config, pmt in zip(config.embedding_tables, pmts):
         table_name = table_config.name
         emb_table = pmt
+        # pyre-fixme[6]: For 1st argument expected `Tensor` but got
+        #  `Union[PartiallyMaterializedTensor, Tensor]`.
         weight: nn.Parameter = nn.Parameter(emb_table)
         # pyre-ignore
         weight._in_backward_optimizers = [EmptyFusedOptimizer()]
@@ -1229,6 +1231,10 @@ def split_embedding_weights(self, no_snapshot: bool = True) -> Tuple[
         Optional[List[torch.Tensor]],
         Optional[List[torch.Tensor]],
     ]:
+        # pyre-fixme[7]: Expected `Tuple[List[PartiallyMaterializedTensor],
+        #  Optional[List[Tensor]], Optional[List[Tensor]]]` but got
+        #  `Tuple[Union[List[PartiallyMaterializedTensor], List[Tensor]],
+        #  Optional[List[Tensor]], Optional[List[Tensor]]]`.
         return self.emb_module.split_embedding_weights(no_snapshot)
 
 
@@ -2027,6 +2033,10 @@ def split_embedding_weights(self, no_snapshot: bool = True) -> Tuple[
         Optional[List[torch.Tensor]],
         Optional[List[torch.Tensor]],
     ]:
+        # pyre-fixme[7]: Expected `Tuple[List[PartiallyMaterializedTensor],
+        #  Optional[List[Tensor]], Optional[List[Tensor]]]` but got
+        #  `Tuple[Union[List[PartiallyMaterializedTensor], List[Tensor]],
+        #  Optional[List[Tensor]], Optional[List[Tensor]]]`.
         return self.emb_module.split_embedding_weights(no_snapshot)
 
 

torchrec/distributed/composable/tests/test_fused_optim_nccl.py

@@ -41,6 +41,8 @@ def _test_sharded_fused_optimizer_state_dict(
             ebc = EmbeddingBagCollection(tables=tables, device=torch.device("meta"))
             apply_optimizer_in_backward(
                 RowWiseAdagrad,
+                # pyre-fixme[6]: For 2nd argument expected `Iterable[Parameter]` but
+                #  got `Iterable[Union[Tensor, Module]]`.
                 [
                     ebc.embedding_bags["table_0"].weight,
                     ebc.embedding_bags["table_1"].weight,
@@ -49,6 +51,8 @@ def _test_sharded_fused_optimizer_state_dict(
             )
             apply_optimizer_in_backward(
                 PartialRowWiseAdam,
+                # pyre-fixme[6]: For 2nd argument expected `Iterable[Parameter]` but
+                #  got `Iterable[Union[Tensor, Module]]`.
                 [
                     ebc.embedding_bags["table_2"].weight,
                     ebc.embedding_bags["table_3"].weight,

torchrec/distributed/embedding.py

@@ -315,6 +315,7 @@ def record_stream(self, stream: torch.Stream) -> None:
         for ctx in self.sharding_contexts:
             ctx.record_stream(stream)
         for f in self.input_features:
+            # pyre-fixme[6]: For 1st argument expected `Stream` but got `Stream`.
             f.record_stream(stream)
         for r in self.reverse_indices:
             r.record_stream(stream)
@@ -892,6 +893,10 @@ def _initialize_torch_state(self) -> None:  # noqa
                     _model_parallel_name_to_compute_kernel[table_name]
                     != EmbeddingComputeKernel.DENSE.value
                 ):
+                    # pyre-fixme[16]: `Module` has no attribute
+                    #  `_in_backward_optimizers`.
+                    # pyre-fixme[16]: `Tensor` has no attribute
+                    #  `_in_backward_optimizers`.
                     self.embeddings[table_name].weight._in_backward_optimizers = [
                         EmptyFusedOptimizer()
                     ]
@@ -1110,6 +1115,8 @@ def reset_parameters(self) -> None:
             if sharding_type == ShardingType.DATA_PARALLEL.value:
                 pg = self._env.process_group
                 with torch.no_grad():
+                    # pyre-fixme[6]: For 1st argument expected `Tensor` but got
+                    #  `TypeUnion[Module, Tensor]`.
                     dist.broadcast(param.data, src=0, group=pg)
 
     def _generate_permute_indices_per_feature(

torchrec/distributed/embedding_lookup.py

@@ -134,6 +134,7 @@ def _load_state_dict(
                     dst_param.detach().copy_(src_param)
                 unexpected_keys.remove(key)
             else:
+                # pyre-fixme[22]: The cast is redundant.
                 missing_keys.append(cast(str, key))
     return missing_keys, unexpected_keys
 
@@ -278,13 +279,17 @@ def prefetch(
                             SSDTableBatchedEmbeddingBags,
                         ),
                     )
+                    # pyre-fixme[16]: Item `Tensor` of `Tensor | Module` has no
+                    #  attribute `prefetch_pipeline`.
                     and not emb_op.emb_module.prefetch_pipeline
                 ):
                     logging.error(
                         f"Invalid setting on {type(emb_op.emb_module)} modules. prefetch_pipeline must be set to True.\n"
                         "If you don’t turn on prefetch_pipeline, cache locations might be wrong in backward and can cause wrong results.\n"
                     )
                 if hasattr(emb_op.emb_module, "prefetch"):
+                    # pyre-fixme[16]: Item `Tensor` of `Tensor | Module` has no
+                    #  attribute `prefetch`.
                     emb_op.emb_module.prefetch(
                         indices=features.values(),
                         offsets=features.offsets(),
@@ -317,6 +322,7 @@ def state_dict(
             destination._metadata = OrderedDict()
 
         for emb_module in self._emb_modules:
+            # pyre-fixme[19]: Expected 0 positional arguments.
             emb_module.state_dict(destination, prefix, keep_vars)
 
         return destination
@@ -363,6 +369,7 @@ def named_parameters_by_table(
             for (
                 table_name,
                 tbe_slice,
+                # pyre-fixme[29]: `Union[Module, Tensor]` is not a function.
             ) in embedding_kernel.named_parameters_by_table():
                 yield (table_name, tbe_slice)
 
@@ -389,10 +396,12 @@ def get_named_split_embedding_weights_snapshot(
 
     def flush(self) -> None:
         for emb_module in self._emb_modules:
+            # pyre-fixme[29]: `Union[Module, Tensor]` is not a function.
             emb_module.flush()
 
     def purge(self) -> None:
         for emb_module in self._emb_modules:
+            # pyre-fixme[29]: `Union[Module, Tensor]` is not a function.
             emb_module.purge()
 
 
@@ -521,6 +530,8 @@ def _need_prefetch(config: GroupedEmbeddingConfig) -> bool:
                 self._feature_splits,
             )
             for emb_op, features in zip(self._emb_modules, features_by_group):
+                # pyre-fixme[6]: For 1st argument expected `GroupedEmbeddingConfig`
+                #  but got `Union[Module, Tensor]`.
                 if not _need_prefetch(emb_op.config):
                     continue
                 if (
@@ -531,13 +542,17 @@ def _need_prefetch(config: GroupedEmbeddingConfig) -> bool:
                             SSDTableBatchedEmbeddingBags,
                         ),
                     )
+                    # pyre-fixme[16]: Item `Tensor` of `Tensor | Module` has no
+                    #  attribute `prefetch_pipeline`.
                     and not emb_op.emb_module.prefetch_pipeline
                 ):
                     logging.error(
                         f"Invalid setting on {type(emb_op.emb_module)} modules. prefetch_pipeline must be set to True.\n"
                         "If you don't turn on prefetch_pipeline, cache locations might be wrong in backward and can cause wrong results.\n"
                     )
                 if hasattr(emb_op.emb_module, "prefetch"):
+                    # pyre-fixme[16]: Item `Tensor` of `Tensor | Module` has no
+                    #  attribute `prefetch`.
                     emb_op.emb_module.prefetch(
                         indices=features.values(),
                         offsets=features.offsets(),
@@ -633,6 +648,7 @@ def state_dict(
             destination._metadata = OrderedDict()
 
         for emb_module in self._emb_modules:
+            # pyre-fixme[19]: Expected 0 positional arguments.
             emb_module.state_dict(destination, prefix, keep_vars)
 
         return destination
@@ -679,6 +695,7 @@ def named_parameters_by_table(
             for (
                 table_name,
                 tbe_slice,
+                # pyre-fixme[29]: `Union[Module, Tensor]` is not a function.
             ) in embedding_kernel.named_parameters_by_table():
                 yield (table_name, tbe_slice)
 
@@ -703,10 +720,12 @@ def get_named_split_embedding_weights_snapshot(
 
     def flush(self) -> None:
         for emb_module in self._emb_modules:
+            # pyre-fixme[29]: `Union[Module, Tensor]` is not a function.
             emb_module.flush()
 
     def purge(self) -> None:
         for emb_module in self._emb_modules:
+            # pyre-fixme[29]: `Union[Module, Tensor]` is not a function.
             emb_module.purge()
 
 
@@ -813,6 +832,7 @@ def state_dict(
             destination._metadata = OrderedDict()
 
         for emb_module in self._emb_modules:
+            # pyre-fixme[19]: Expected 0 positional arguments.
             emb_module.state_dict(destination, prefix, keep_vars)
 
         return destination
@@ -980,6 +1000,7 @@ def state_dict(
             destination._metadata = OrderedDict()
 
         for emb_module in self._emb_modules:
+            # pyre-fixme[19]: Expected 0 positional arguments.
             emb_module.state_dict(destination, prefix, keep_vars)
 
         return destination

torchrec/distributed/embedding_tower_sharding.py

@@ -501,6 +501,8 @@ def __init__(
         # create mapping of logical towers to physical towers
         tables_per_lt: List[Set[str]] = []
         for tower in module.towers:
+            # pyre-fixme[6]: For 1st argument expected `EmbeddingTower` but got
+            #  `Module`.
             lt_tables = set(tower_sharder.shardable_parameters(tower).keys())
             tables_per_lt.append(lt_tables)
             # check the tables in a logical tower are on same physical tower
@@ -546,6 +548,8 @@ def __init__(
             self._kjt_num_features_per_pt.append(len(kjt_names))
             self._wkjt_num_features_per_pt.append(len(wkjt_names))
 
+        # pyre-fixme[9]: local_towers has type `List[Tuple[str, EmbeddingTower]]`;
+        #  used as `List[Tuple[str, Module]]`.
         local_towers: List[Tuple[str, EmbeddingTower]] = [
             (str(i), tower)
             for i, tower in enumerate(module.towers)
@@ -603,6 +607,8 @@ def __init__(
         # Setup output dists for quantized comms
         output_dists = nn.ModuleList()
         for embedding in self.embeddings.values():
+            # pyre-fixme[6]: For 1st argument expected `Iterable[Module]` but got
+            #  `Union[Module, Tensor]`.
             output_dists.extend(embedding._output_dists)
         self._output_dists: nn.ModuleList = output_dists
 
@@ -780,10 +786,13 @@ def state_dict(
             # pyre-ignore [16]
             destination._metadata = OrderedDict()
         for i, embedding in self.embeddings.items():
+            # pyre-fixme[19]: Expected 0 positional arguments.
             embedding.state_dict(
                 destination, prefix + f"towers.{i}.embedding.", keep_vars
             )
         for i, interaction in self.interactions.items():
+            # pyre-fixme[16]: Item `Tensor` of `Tensor | Module` has no attribute
+            #  `state_dict`.
             interaction.module.state_dict(
                 destination, prefix + f"towers.{i}.interaction.", keep_vars
             )
@@ -792,6 +801,9 @@ def state_dict(
     @property
     def fused_optimizer(self) -> KeyedOptimizer:
         return CombinedOptimizer(
+            # pyre-fixme[6]: For 1st argument expected `List[Union[Tuple[str,
+            #  KeyedOptimizer], KeyedOptimizer]]` but got `List[Tuple[str,
+            #  Union[Module, Tensor]]]`.
             [
                 (f"towers.{tower_index}.embedding", embedding.fused_optimizer)
                 for tower_index, embedding in self.embeddings.items()
@@ -809,6 +821,8 @@ def named_parameters(
             )
         for i, interaction in self.interactions.items():
             yield from (
+                # pyre-fixme[16]: Item `Tensor` of `Tensor | Module` has no
+                #  attribute `named_parameters`.
                 interaction.module.named_parameters(
                     append_prefix(prefix, f"towers.{i}.interaction"), recurse
                 )
@@ -825,6 +839,8 @@ def named_buffers(
             )
         for i, interaction in self.interactions.items():
             yield from (
+                # pyre-fixme[16]: Item `Tensor` of `Tensor | Module` has no
+                #  attribute `named_buffers`.
                 interaction.module.named_buffers(
                     append_prefix(prefix, f"towers.{i}.interaction"), recurse
                 )
@@ -833,13 +849,16 @@ def named_buffers(
     def sharded_parameter_names(self, prefix: str = "") -> Iterator[str]:
         for i, embedding in self.embeddings.items():
             yield from (
+                # pyre-fixme[29]: `Union[Module, Tensor]` is not a function.
                 embedding.sharded_parameter_names(
                     append_prefix(prefix, f"towers.{i}.embedding")
                 )
             )
         for i, interaction in self.interactions.items():
             yield from (
                 key
+                # pyre-fixme[16]: Item `Tensor` of `Tensor | Module` has no
+                #  attribute `named_parameters`.
                 for key, _ in interaction.module.named_parameters(
                     append_prefix(prefix, f"towers.{i}.interaction")
                 )
@@ -993,6 +1012,8 @@ def shardable_parameters(
 
         named_parameters: Dict[str, nn.Parameter] = {}
         for tower in module.towers:
+            # pyre-fixme[6]: For 1st argument expected `EmbeddingTower` but got
+            #  `Module`.
             named_parameters.update(self._tower_sharder.shardable_parameters(tower))
         return named_parameters
 

torchrec/distributed/embedding_types.py

@@ -135,6 +135,8 @@ def __iter__(self) -> Iterator[KeyedJaggedTensor]:
         return iter(self.features)
 
     @torch.jit._drop
+    # pyre-fixme[14]: `record_stream` overrides method defined in `Multistreamable`
+    #  inconsistently.
     def record_stream(self, stream: torch.cuda.streams.Stream) -> None:
         for feature in self.features:
             feature.record_stream(stream)
@@ -160,6 +162,7 @@ class InputDistOutputs(Multistreamable):
 
     def record_stream(self, stream: torch.Stream) -> None:
         for feature in self.features:
+            # pyre-fixme[6]: For 1st argument expected `Stream` but got `Stream`.
             feature.record_stream(stream)
         if self.unbucketize_permute_tensor is not None:
             self.unbucketize_permute_tensor.record_stream(stream)
@@ -189,6 +192,7 @@ def __iter__(self) -> Iterator[KJTList]:
     @torch.jit._drop
     def record_stream(self, stream: torch.Stream) -> None:
         for feature in self.features_list:
+            # pyre-fixme[6]: For 1st argument expected `Stream` but got `Stream`.
             feature.record_stream(stream)
 
     @torch.jit._drop

torchrec/distributed/embeddingbag.py

@@ -972,6 +972,10 @@ def _initialize_torch_state(self, skip_registering: bool = False) -> None:  # no
                     _model_parallel_name_to_compute_kernel[table_name]
                     != EmbeddingComputeKernel.DENSE.value
                 ):
+                    # pyre-fixme[16]: `Module` has no attribute
+                    #  `_in_backward_optimizers`.
+                    # pyre-fixme[16]: `Tensor` has no attribute
+                    #  `_in_backward_optimizers`.
                     self.embedding_bags[table_name].weight._in_backward_optimizers = [
                         EmptyFusedOptimizer()
                     ]
@@ -1137,6 +1141,8 @@ def reset_parameters(self) -> None:
             if sharding_type == ShardingType.DATA_PARALLEL.value:
                 pg = self._env.process_group
                 with torch.no_grad():
+                    # pyre-fixme[6]: For 1st argument expected `Tensor` but got
+                    #  `Union[Module, Tensor]`.
                     dist.broadcast(param.data, src=0, group=pg)
 
     def _create_input_dist(

torchrec/distributed/fp_embeddingbag.py

@@ -83,6 +83,8 @@ def __init__(
             self._is_collection = True
         else:
             self._feature_processors = torch.nn.ModuleDict(
+                # pyre-fixme[29]: `Union[(self: ModuleDict) -> ItemsView[str,
+                #  Module], Module, Tensor]` is not a function.
                 module._feature_processors.items()
             )
             self._is_collection = False

torchrec/distributed/mc_embedding.py

@@ -60,6 +60,7 @@ def record_stream(self, stream: torch.Stream) -> None:
                     continue
                 value.record_stream(stream)
         if self.remapped_kjt is not None:
+            # pyre-fixme[6]: For 1st argument expected `Stream` but got `Stream`.
             self.remapped_kjt.record_stream(stream)
 
 

torchrec/distributed/mc_embeddingbag.py

@@ -45,6 +45,7 @@ def record_stream(self, stream: torch.Stream) -> None:
                     continue
                 value.record_stream(stream)
         if self.remapped_kjt is not None:
+            # pyre-fixme[6]: For 1st argument expected `Stream` but got `Stream`.
             self.remapped_kjt.record_stream(stream)