Remove unused imports (pytorch#2635)

Summary:

Remove unused imports

Reviewed By: PaulZhang12

Differential Revision: D67111062

torchrec/distributed/benchmark/benchmark_train.py

@@ -10,7 +10,6 @@
 #!/usr/bin/env python3
 
 import argparse
-import copy
 import logging
 import os
 import time

torchrec/distributed/embedding_dim_bucketer.py

@@ -10,10 +10,7 @@
 from enum import Enum, unique
 from typing import Dict, List
 
-from torchrec.distributed.embedding_types import (
-    EmbeddingComputeKernel,
-    ShardedEmbeddingTable,
-)
+from torchrec.distributed.embedding_types import ShardedEmbeddingTable
 from torchrec.modules.embedding_configs import DATA_TYPE_NUM_BITS, DataType
 
 

torchrec/distributed/fused_params.py

@@ -7,7 +7,7 @@
 
 # pyre-strict
 
-from typing import Any, Dict, Iterable, List, Optional
+from typing import Any, Dict, Iterable, Optional
 
 import torch
 
@@ -55,7 +55,7 @@ def is_fused_param_register_tbe(fused_params: Optional[Dict[str, Any]]) -> bool:
 
 
 def get_fused_param_tbe_row_alignment(
-    fused_params: Optional[Dict[str, Any]]
+    fused_params: Optional[Dict[str, Any]],
 ) -> Optional[int]:
     if fused_params is None or FUSED_PARAM_TBE_ROW_ALIGNMENT not in fused_params:
         return None
@@ -64,7 +64,7 @@ def get_fused_param_tbe_row_alignment(
 
 
 def fused_param_bounds_check_mode(
-    fused_params: Optional[Dict[str, Any]]
+    fused_params: Optional[Dict[str, Any]],
 ) -> Optional[BoundsCheckMode]:
     if fused_params is None or FUSED_PARAM_BOUNDS_CHECK_MODE not in fused_params:
         return None
@@ -73,7 +73,7 @@ def fused_param_bounds_check_mode(
 
 
 def fused_param_lengths_to_offsets_lookup(
-    fused_params: Optional[Dict[str, Any]]
+    fused_params: Optional[Dict[str, Any]],
 ) -> bool:
     if (
         fused_params is None
@@ -85,7 +85,7 @@ def fused_param_lengths_to_offsets_lookup(
 
 
 def is_fused_param_quant_state_dict_split_scale_bias(
-    fused_params: Optional[Dict[str, Any]]
+    fused_params: Optional[Dict[str, Any]],
 ) -> bool:
     return (
         fused_params
@@ -95,7 +95,7 @@ def is_fused_param_quant_state_dict_split_scale_bias(
 
 
 def tbe_fused_params(
-    fused_params: Optional[Dict[str, Any]]
+    fused_params: Optional[Dict[str, Any]],
 ) -> Optional[Dict[str, Any]]:
     if not fused_params:
         return None

torchrec/distributed/shard.py

@@ -29,7 +29,6 @@
 )
 from torchrec.distributed.utils import init_parameters
 from torchrec.modules.utils import reset_module_states_post_sharding
-from torchrec.types import CacheMixin
 
 
 def _join_module_path(path: str, name: str) -> str:

torchrec/distributed/sharding/cw_sharding.py

@@ -32,7 +32,6 @@
     DTensorMetadata,
     EmbeddingComputeKernel,
     InputDistOutputs,
-    KJTList,
     ShardedEmbeddingTable,
 )
 from torchrec.distributed.sharding.tw_sharding import (

torchrec/distributed/sharding/rw_sequence_sharding.py

@@ -17,7 +17,6 @@
 )
 from torchrec.distributed.embedding_lookup import (
     GroupedEmbeddingsLookup,
-    InferCPUGroupedEmbeddingsLookup,
     InferGroupedEmbeddingsLookup,
 )
 from torchrec.distributed.embedding_sharding import (

torchrec/distributed/tensor_sharding.py

@@ -11,7 +11,7 @@
 
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from typing import cast, Iterable, List, Optional, Tuple
+from typing import Iterable, List, Optional, Tuple
 
 import torch
 from torch import distributed as dist

torchrec/distributed/test_utils/test_sharding.py

@@ -41,7 +41,6 @@
 )
 from torchrec.distributed.types import (
     EmbeddingModuleShardingPlan,
-    EnumerableShardingSpec,
     ModuleSharder,
     ShardedTensor,
     ShardingEnv,
@@ -288,7 +287,6 @@ def sharding_single_rank_test(
     world_size_2D: Optional[int] = None,
     node_group_size: Optional[int] = None,
 ) -> None:
-
     with MultiProcessContext(rank, world_size, backend, local_size) as ctx:
         # Generate model & inputs.
         (global_model, inputs) = gen_model_and_input(

torchrec/distributed/tests/test_infer_shardings.py

@@ -37,10 +37,8 @@
 )
 from torchrec.distributed.quant_embedding import QuantEmbeddingCollectionSharder
 from torchrec.distributed.quant_embeddingbag import (
-    QuantEmbeddingBagCollection,
     QuantEmbeddingBagCollectionSharder,
     QuantFeatureProcessedEmbeddingBagCollectionSharder,
-    ShardedQuantEmbeddingBagCollection,
 )
 from torchrec.distributed.quant_state import sharded_tbes_weights_spec, WeightSpec
 from torchrec.distributed.shard import _shard_modules

torchrec/distributed/tests/test_quant_model_parallel.py

@@ -8,7 +8,7 @@
 # pyre-strict
 
 import unittest
-from typing import Any, cast, Dict, List, Optional, Tuple
+from typing import cast, Dict, Optional, Tuple
 
 import hypothesis.strategies as st
 import torch

torchrec/distributed/tests/test_sequence_model_parallel_single_rank.py

@@ -9,26 +9,21 @@
 
 
 import unittest
-from typing import cast, Dict, List, Optional, OrderedDict, Tuple
+from typing import cast, OrderedDict
 
 import hypothesis.strategies as st
 import torch
 from hypothesis import given, settings, Verbosity
-from torch import distributed as dist, nn
-from torchrec import distributed as trec_dist
-from torchrec.distributed import DistributedModelParallel
+from torch import nn
 from torchrec.distributed.embedding_types import EmbeddingComputeKernel
-from torchrec.distributed.model_parallel import get_default_sharders
-from torchrec.distributed.planner import EmbeddingShardingPlanner, Topology
-from torchrec.distributed.test_utils.test_model import ModelInput
 from torchrec.distributed.test_utils.test_model_parallel_base import (
     ModelParallelSingleRankBase,
 )
 from torchrec.distributed.tests.test_sequence_model import (
     TestEmbeddingCollectionSharder,
     TestSequenceSparseNN,
 )
-from torchrec.distributed.types import ModuleSharder, ShardingEnv, ShardingType
+from torchrec.distributed.types import ModuleSharder, ShardingType
 from torchrec.modules.embedding_configs import DataType, EmbeddingConfig
 
 

torchrec/distributed/train_pipeline/train_pipelines.py

@@ -11,7 +11,6 @@
 import contextlib
 import logging
 from collections import deque
-from contextlib import contextmanager
 from dataclasses import dataclass
 from typing import (
     Any,
@@ -31,7 +30,6 @@
 )
 
 import torch
-import torchrec.distributed.comm_ops
 from torch.autograd.profiler import record_function
 from torchrec.distributed.dist_data import KJTAllToAllTensorsAwaitable
 from torchrec.distributed.model_parallel import ShardedModule
@@ -760,7 +758,6 @@ def _grad_swap(self) -> None:
             param.grad = grad
 
     def _init_embedding_streams(self) -> None:
-
         for _ in self._pipelined_modules:
             self._embedding_streams.append(
                 (torch.get_device_module(self._device).Stream(priority=0))

torchrec/inference/client.py

@@ -11,7 +11,6 @@
 import grpc
 import predictor_pb2, predictor_pb2_grpc
 import torch
-from torch.utils.data import DataLoader
 from torchrec.datasets.criteo import DEFAULT_CAT_NAMES, DEFAULT_INT_NAMES
 from torchrec.datasets.random import RandomRecDataset
 from torchrec.datasets.utils import Batch

torchrec/inference/tests/test_inference.py

@@ -14,7 +14,6 @@
 
 import torch
 from fbgemm_gpu.split_embedding_configs import SparseType
-from torch.fx import symbolic_trace
 from torchrec import PoolingType
 from torchrec.datasets.criteo import DEFAULT_CAT_NAMES, DEFAULT_INT_NAMES
 from torchrec.distributed.fused_params import (

torchrec/ir/types.py

@@ -10,7 +10,7 @@
 #!/usr/bin/env python3
 
 import abc
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional
 
 import torch