Fix kernel test (pytorch#3026)

Summary: Pull Request resolved: pytorch#3026

Differential Revision: D75574459

Author: kausv

torchrec/distributed/batched_embedding_kernel.py

@@ -303,7 +303,7 @@ def __init__(  # noqa C901
                 sharded_t._local_shards[0].tensor
                 for sharded_t in self._sharded_embedding_weight_ids
             ]
-            if self._sharded_embedding_weight_ids is not None
+            if self._sharded_embedding_weight_ids
             else None
         )
 
@@ -1439,7 +1439,13 @@ def _init_sharded_split_embedding_weights(
         )
         emb_table_config_copy = copy.deepcopy(self._config.embedding_tables)
         for emb_table in emb_table_config_copy:
-            emb_table.local_metadata.placement._device = torch.device("cpu")
+            none_throws(
+                none_throws(
+                    emb_table.local_metadata,
+                    f"local_metadata is None for emb_table: {emb_table.name}",
+                ).placement,
+                "placement is None for local_metadata of emb table: {emb_table.name}",
+            )._device = torch.device("cpu")
 
         pmt_sharded_t_list = create_virtual_sharded_tensors(
             emb_table_config_copy,

torchrec/distributed/embedding_kernel.py

@@ -144,7 +144,7 @@ def get_key_from_embedding_table(embedding_table: ShardedEmbeddingTable) -> str:
         key = get_key_from_embedding_table(embedding_table)
         assert embedding_table.use_virtual_table
 
-        assert embedding_table.global_metadata is not None and pg is not None
+        assert embedding_table.global_metadata is not None
         global_metadata = copy.deepcopy(embedding_table.global_metadata)
         create_virtual_table_global_metadata(global_metadata, my_rank, param)
         key_to_global_metadata[key] = global_metadata

torchrec/distributed/tests/test_embedding_sharding.py

@@ -35,10 +35,11 @@
     ShardedEmbeddingTable,
 )
 from torchrec.distributed.sharding.sequence_sharding import SequenceShardingContext
+from torchrec.distributed.types import ShardedTensorMetadata, ShardMetadata
 from torchrec.modules.embedding_configs import DataType, PoolingType
 from torchrec.sparse.jagged_tensor import KeyedJaggedTensor
 
-WORLD_SIZE = 4
+WORLD_SIZE = 2
 
 
 class TestGetWeightedAverageCacheLoadFactor(unittest.TestCase):
@@ -546,24 +547,16 @@ class TestECBucketMetadata(unittest.TestCase):
         data_type=st.sampled_from([DataType.FP16, DataType.FP32]),
         embedding_dim=st.sampled_from(list(range(160, 320, 40))),
         total_bucket=st.sampled_from([14, 20, 32, 40]),
-        my_rank=st.integers(min_value=0, max_value=WORLD_SIZE),
+        my_rank=st.integers(min_value=0, max_value=WORLD_SIZE - 1),
     )
     @settings(max_examples=10, deadline=10000)
     def test_bucket_metadata_calculation_util(
         self, data_type: DataType, embedding_dim: int, total_bucket: int, my_rank: int
     ) -> None:
-        compute_kernels = [
-            EmbeddingComputeKernel.SSD_VIRTUAL_TABLE,
-            EmbeddingComputeKernel.SSD_VIRTUAL_TABLE,
-            EmbeddingComputeKernel.SSD_VIRTUAL_TABLE,
-            EmbeddingComputeKernel.SSD_VIRTUAL_TABLE,
-        ]
+        compute_kernels = [EmbeddingComputeKernel.SSD_VIRTUAL_TABLE] * WORLD_SIZE
         fused_params_groups = [
             {"cache_load_factor": 0.5},
-            {"cache_load_factor": 0.5},
-            {"cache_load_factor": 0.5},
-            {"cache_load_factor": 0.5},
-        ]
+        ] * WORLD_SIZE
         tables = [
             ShardedEmbeddingTable(
                 name=f"table_{i}",
@@ -579,8 +572,27 @@ def test_bucket_metadata_calculation_util(
                 num_embeddings=10000 * (2 * i + 1),
                 total_num_buckets=total_bucket,
                 use_virtual_table=True,
+                local_metadata=ShardMetadata(
+                    shard_offsets=[i * (10000 * (2 * i + 1) // WORLD_SIZE), 0],
+                    shard_sizes=[10000 * (2 * i + 1) // WORLD_SIZE, embedding_dim],
+                    placement=f"rank:{i}/cuda:{i}",
+                ),
+                global_metadata=ShardedTensorMetadata(
+                    shards_metadata=[
+                        ShardMetadata(
+                            shard_offsets=[j * (10000 * (2 * i + 1) // WORLD_SIZE), 0],
+                            shard_sizes=[
+                                10000 * (2 * i + 1) // WORLD_SIZE,
+                                embedding_dim,
+                            ],
+                            placement=f"rank:{j}/cuda:{j}",
+                        )
+                        for j in range(WORLD_SIZE)
+                    ],
+                    size=torch.Size([10000 * (2 * i + 1), embedding_dim]),
+                ),
             )
-            for i in range(len(compute_kernels))
+            for i in range(WORLD_SIZE)
         ]
 
         # since we don't have access to _group_tables_per_rank