Change util to return a set of bucket metadata per shard (#2920)

Summary:
Pull Request resolved: #2920

Based on discussions, we are modifying the util to give a master set of metadata about buckets for the given shards.

We will stack a follow up diff for ShardedEC to return this data.

Reviewed By: emlin

Differential Revision: D73522460

fbshipit-source-id: 89648ef3531756c9a6d43422877c12454bd1ffcb

Author: kausv

torchrec/distributed/tests/test_utils.py

@@ -29,12 +29,13 @@
     ModuleSharder,
     MultiPassPrefetchConfig,
     ParameterSharding,
+    ShardingBucketMetadata,
     ShardMetadata,
 )
 from torchrec.distributed.utils import (
     add_params_from_parameter_sharding,
     convert_to_fbgemm_types,
-    get_bucket_offsets_from_shard_metadata,
+    get_bucket_metadata_from_shard_metadata,
     get_unsharded_module_names,
     merge_fused_params,
 )
@@ -472,15 +473,15 @@ def test_convert_to_fbgemm_types(self) -> None:
         self.assertFalse(isinstance(per_table_fused_params["output_dtype"], DataType))
 
 
-class TestBucketOffsets(unittest.TestCase):
-    def test_bucket_offsets(self) -> None:
+class TestBucketMetadata(unittest.TestCase):
+    def test_bucket_metadata(self) -> None:
         # Given no shards
-        # When we get bucket offsets from get_bucket_offsets_from_shard_metadata
+        # When we get bucket metadata from get_bucket_metadata_from_shard_metadata
         # Then an error should be raised
         self.assertRaisesRegex(
             AssertionError,
             "Shards cannot be empty",
-            get_bucket_offsets_from_shard_metadata,
+            get_bucket_metadata_from_shard_metadata,
             [],
             num_buckets=4,
         )
@@ -490,11 +491,13 @@ def test_bucket_offsets(self) -> None:
             ShardMetadata(shard_offsets=[0], shard_sizes=[5], placement="rank:0/cuda:0")
         ]
 
-        # When we get bucket offsets from get_bucket_offsets_from_shard_metadata
-        bucket_offsets = get_bucket_offsets_from_shard_metadata(shards, num_buckets=5)
+        # When we get bucket offsets from get_bucket_metadata_from_shard_metadata
+        bucket_metadata = get_bucket_metadata_from_shard_metadata(shards, num_buckets=5)
         # Then we should get 1 offset with value 0
-        expected_offsets = [0]
-        self.assertEqual(bucket_offsets, expected_offsets)
+        expected_metadata = ShardingBucketMetadata(
+            num_buckets_per_shard=[5], bucket_offsets_per_shard=[0], bucket_size=1
+        )
+        self.assertEqual(bucket_metadata, expected_metadata)
 
         # Given 2 shards of size 5 and 4 buckets
         shards = [
@@ -506,12 +509,12 @@ def test_bucket_offsets(self) -> None:
             ),
         ]
 
-        # When we get bucket offsets from get_bucket_offsets_from_shard_metadata
+        # When we get bucket offsets from get_bucket_metadata_from_shard_metadata
         # Then an error should be raised
         self.assertRaisesRegex(
             AssertionError,
             "Table size '10' must be divisible by num_buckets '4'",
-            get_bucket_offsets_from_shard_metadata,
+            get_bucket_metadata_from_shard_metadata,
             shards,
             num_buckets=4,
         )
@@ -526,12 +529,12 @@ def test_bucket_offsets(self) -> None:
             ),
         ]
 
-        # When we get bucket offsets from get_bucket_offsets_from_shard_metadata
+        # When we get bucket offsets from get_bucket_metadata_from_shard_metadata
         # Then an error should be raised
         self.assertRaisesRegex(
             AssertionError,
             "Table size '4' must be divisible by num_buckets '5'",
-            get_bucket_offsets_from_shard_metadata,
+            get_bucket_metadata_from_shard_metadata,
             shards,
             num_buckets=5,
         )
@@ -546,12 +549,12 @@ def test_bucket_offsets(self) -> None:
             ),
         ]
 
-        # When we get bucket offsets from get_bucket_offsets_from_shard_metadata
+        # When we get bucket offsets from get_bucket_metadata_from_shard_metadata
         # Then an error should be raised
         self.assertRaisesRegex(
             AssertionError,
             r"Shard shard_offsets\[1\] '5' is not 0. Table should be only row-wise sharded for bucketization",
-            get_bucket_offsets_from_shard_metadata,
+            get_bucket_metadata_from_shard_metadata,
             shards,
             num_buckets=2,
         )
@@ -566,17 +569,17 @@ def test_bucket_offsets(self) -> None:
             ),
         ]
 
-        # When we get bucket offsets from get_bucket_offsets_from_shard_metadata
+        # When we get bucket offsets from get_bucket_metadata_from_shard_metadata
         # Then an error should be raised
         self.assertRaisesRegex(
             AssertionError,
             r"Shard size\[0\] '10' is not divisible by bucket size '4'",
-            get_bucket_offsets_from_shard_metadata,
+            get_bucket_metadata_from_shard_metadata,
             shards,
             num_buckets=5,
         )
 
-        # Given 2 shards of size 20 and 4 buckets
+        # Given 2 shards of size 20 and 10 buckets
         shards = [
             ShardMetadata(
                 shard_offsets=[0], shard_sizes=[20], placement="rank:0/cuda:0"
@@ -585,13 +588,20 @@ def test_bucket_offsets(self) -> None:
                 shard_offsets=[20], shard_sizes=[20], placement="rank:0/cuda:0"
             ),
         ]
-        # When we get bucket offsets from get_bucket_offsets_from_shard_metadata
-        bucket_offsets = get_bucket_offsets_from_shard_metadata(
+        # When we get bucket offsets from get_bucket_metadata_from_shard_metadata
+        bucket_metadata = get_bucket_metadata_from_shard_metadata(
             shards,
             num_buckets=10,
         )
-        # Then bucket offsets should be set to [0, 5]
-        self.assertEqual(bucket_offsets, [0, 5])
+        # Then num_buckets_per_shard should be set to [5, 5]
+        self.assertEqual(
+            bucket_metadata,
+            ShardingBucketMetadata(
+                num_buckets_per_shard=[5, 5],
+                bucket_offsets_per_shard=[0, 5],
+                bucket_size=4,
+            ),
+        )
 
         # Given 3 uneven shards of sizes 12, 16 and 20 and 12 buckets
         shards = [
@@ -606,10 +616,17 @@ def test_bucket_offsets(self) -> None:
             ),
         ]
 
-        # When we get bucket offsets from get_bucket_offsets_from_shard_metadata
-        bucket_offsets = get_bucket_offsets_from_shard_metadata(
+        # When we get bucket offsets from get_bucket_metadata_from_shard_metadata
+        bucket_metadata = get_bucket_metadata_from_shard_metadata(
             shards,
             num_buckets=12,
         )
-        # Then bucket offsets should be set to [0, 3, 7]
-        self.assertEqual(bucket_offsets, [0, 3, 7])
+        # Then num_buckets_per_shard should be set to [3, 4, 5]
+        self.assertEqual(
+            bucket_metadata,
+            ShardingBucketMetadata(
+                num_buckets_per_shard=[3, 4, 5],
+                bucket_offsets_per_shard=[0, 3, 7],
+                bucket_size=4,
+            ),
+        )

torchrec/distributed/types.py

@@ -1223,3 +1223,19 @@ class ObjectPoolShardingType(Enum):
 class ObjectPoolShardingPlan(ModuleShardingPlan):
     sharding_type: ObjectPoolShardingType
     inference: bool = False
+
+
+@dataclass
+class ShardingBucketMetadata:
+    """
+    If a table is row-wise sharded with bucketization, this class contains the bucket information for the table.
+
+    Attributes:
+        num_buckets_per_shard (List[int]): Number of buckets in each shard of the table.
+        bucket_offsets_per_shard (List[int]): Index of the first bucket in each shard.
+        bucket_size (int): No. of rows in one bucket.
+    """
+
+    num_buckets_per_shard: List[int]
+    bucket_offsets_per_shard: List[int]
+    bucket_size: int

torchrec/distributed/utils.py

@@ -28,6 +28,7 @@
     EmbeddingEvent,
     ParameterSharding,
     ShardedModule,
+    ShardingBucketMetadata,
     ShardingType,
     ShardMetadata,
 )
@@ -567,32 +568,32 @@ def create_global_tensor_shape_stride_from_metadata(
     return size, (size[1], 1) if size else (torch.Size([0, 0]), (0, 1))
 
 
-def get_bucket_offsets_from_shard_metadata(
+def get_bucket_metadata_from_shard_metadata(
     shards: List[ShardMetadata],
     num_buckets: int,
-) -> List[int]:
+) -> ShardingBucketMetadata:
     """
-    Calculate the bucket offsets from shard metadata.
+    Calculate the bucket metadata from shard metadata.
 
     This function assumes the table is to be row-wise sharded in equal sized buckets across bucket boundaries.
-    It computes the sequential bucket offsets for each shard. It ensures that the table size
-    is divisible by the number of buckets and that each shard size is divisible
-    by the bucket size.
+    It computes the number of buckets per shard and the bucket size.
 
     Args:
-        shards (List[ShardMetadata]): A list of shard metadata objects.
+        shards (List[ShardMetadata]): Shard metadata for all shards of a table.
         num_buckets (int): The number of buckets to divide the table into.
 
     Returns:
-        List[int]: A list of bucket offsets.
+        ShardingBucketMetadata: An object containing the number of buckets per shard and the bucket size.
     """
     assert len(shards) > 0, "Shards cannot be empty"
     table_size = shards[-1].shard_offsets[0] + shards[-1].shard_sizes[0]
     assert (
         table_size % num_buckets == 0
     ), f"Table size '{table_size}' must be divisible by num_buckets '{num_buckets}'"
-    bucket_offsets: List[int] = []
     bucket_size = table_size // num_buckets
+    bucket_metadata: ShardingBucketMetadata = ShardingBucketMetadata(
+        num_buckets_per_shard=[], bucket_offsets_per_shard=[], bucket_size=bucket_size
+    )
     current_bucket_offset = 0
     for shard in shards:
         assert (
@@ -602,7 +603,8 @@ def get_bucket_offsets_from_shard_metadata(
             shard.shard_sizes[0] % bucket_size == 0
         ), f"Shard size[0] '{shard.shard_sizes[0]}' is not divisible by bucket size '{bucket_size}'"
         num_buckets_in_shard = shard.shard_sizes[0] // bucket_size
-        bucket_offsets.append(current_bucket_offset)
+        bucket_metadata.num_buckets_per_shard.append(num_buckets_in_shard)
+        bucket_metadata.bucket_offsets_per_shard.append(current_bucket_offset)
         current_bucket_offset += num_buckets_in_shard
 
-    return bucket_offsets
+    return bucket_metadata