add embedding config options in benchmark (#3430)

Summary:
Pull Request resolved: #3430

# context
* add "additional tables" option in test_tables so that the benchmark config can play with different flavors of tables.
* unweighted tables would add two more tables (0_0 and 0_1), while the weighted tables remain the same. And in case there's need for a third embedding module (tables), 2_1 and 2_2 are added there.
```
  additional_tables:
    - - name: additional_tables_0_0
        embedding_dim: 128
        num_embeddings: 100_000
        feature_names: ["additional_0_0"]
      - name: additional_tables_0_1
        embedding_dim: 128
        num_embeddings: 100_000
        feature_names: ["additional_0_1"]
    - []
    - - name: additional_tables_2_1
        embedding_dim: 128
        num_embeddings: 100_000
        feature_names: ["additional_2_1"]
```

Reviewed By: spmex

Differential Revision: D83881145

fbshipit-source-id: c40fdad26bfe9345550908cb67c00204da2a8cbf

Author: TroyGarden

torchrec/distributed/benchmark/README.md

@@ -2,15 +2,13 @@
 ## usage
 - internal:
 ```
-hash=$(hg whereami | cut -c 1-10)
 buck2 run @fbcode//mode/opt fbcode//torchrec/distributed/benchmark:benchmark_train_pipeline -- \
     --yaml_config=fbcode/torchrec/distributed/benchmark/yaml/sparse_data_dist_base.yml \
-    --profile_name=sparse_data_dist_base_${hash:-$USER} # overrides the yaml config
+    --profile_name=sparse_data_dist_base_$(hg whereami | cut -c 1-10 || echo $USER) # overrides the yaml config
 ```
 - oss:
 ```
-hash=`git rev-parse --short HEAD`
 python -m torchrec.distributed.benchmark.benchmark_train_pipeline \
     --yaml_config=fbcode/torchrec/distributed/benchmark/yaml/sparse_data_dist_base.yml \
-    --profile_name=sparse_data_dist_base_${hash:-$USER} # overrides the yaml config
+    --profile_name=sparse_data_dist_base_$(git rev-parse --short HEAD || echo $USER) # overrides the yaml config
 ```

torchrec/distributed/benchmark/benchmark_train_pipeline.py

@@ -39,7 +39,6 @@
     CPUMemoryStats,
     generate_planner,
     generate_sharded_model_and_optimizer,
-    generate_tables,
     GPUMemoryStats,
 )
 from torchrec.distributed.comm import get_local_size
@@ -52,6 +51,7 @@
 )
 from torchrec.distributed.test_utils.test_input import ModelInput
 from torchrec.distributed.test_utils.test_model import TestOverArchLarge
+from torchrec.distributed.test_utils.test_tables import EmbeddingTablesConfig
 from torchrec.distributed.train_pipeline import TrainPipeline
 from torchrec.distributed.types import ShardingType
 from torchrec.modules.embedding_configs import EmbeddingBagConfig
@@ -116,28 +116,6 @@ class RunOptions:
     export_stacks: bool = False
 
 
-@dataclass
-class EmbeddingTablesConfig:
-    """
-    Configuration for embedding tables.
-
-    This class defines the parameters for generating embedding tables with both weighted
-    and unweighted features.
-
-    Args:
-        num_unweighted_features (int): Number of unweighted features to generate.
-            Default is 100.
-        num_weighted_features (int): Number of weighted features to generate.
-            Default is 100.
-        embedding_feature_dim (int): Dimension of the embedding vectors.
-            Default is 128.
-    """
-
-    num_unweighted_features: int = 100
-    num_weighted_features: int = 100
-    embedding_feature_dim: int = 128
-
-
 @dataclass
 class PipelineConfig:
     """
@@ -206,11 +184,7 @@ def main(
     pipeline_config: PipelineConfig,
     model_config: Optional[BaseModelConfig] = None,
 ) -> None:
-    tables, weighted_tables = generate_tables(
-        num_unweighted_features=table_config.num_unweighted_features,
-        num_weighted_features=table_config.num_weighted_features,
-        embedding_feature_dim=table_config.embedding_feature_dim,
-    )
+    tables, weighted_tables, *_ = table_config.generate_tables()
 
     if model_config is None:
         model_config = create_model_config(
@@ -256,11 +230,7 @@ def run_pipeline(
     model_config: BaseModelConfig,
 ) -> BenchmarkResult:
 
-    tables, weighted_tables = generate_tables(
-        num_unweighted_features=table_config.num_unweighted_features,
-        num_weighted_features=table_config.num_weighted_features,
-        embedding_feature_dim=table_config.embedding_feature_dim,
-    )
+    tables, weighted_tables, *_ = table_config.generate_tables()
 
     benchmark_res_per_rank = run_multi_process_func(
         func=runner,

torchrec/distributed/benchmark/benchmark_utils.py

@@ -381,55 +381,6 @@ def set_embedding_config(
     return embedding_configs, pooling_configs
 
 
-def generate_tables(
-    num_unweighted_features: int = 100,
-    num_weighted_features: int = 100,
-    embedding_feature_dim: int = 128,
-) -> Tuple[
-    List[EmbeddingBagConfig],
-    List[EmbeddingBagConfig],
-]:
-    """
-    Generate embedding bag configurations for both unweighted and weighted features.
-
-    This function creates two lists of EmbeddingBagConfig objects:
-    1. Unweighted tables: Named as "table_{i}" with feature names "feature_{i}"
-    2. Weighted tables: Named as "weighted_table_{i}" with feature names "weighted_feature_{i}"
-
-    For both types, the number of embeddings scales with the feature index,
-    calculated as max(i + 1, 100) * 1000.
-
-    Args:
-        num_unweighted_features (int): Number of unweighted features to generate.
-        num_weighted_features (int): Number of weighted features to generate.
-        embedding_feature_dim (int): Dimension of the embedding vectors.
-
-    Returns:
-        Tuple[List[EmbeddingBagConfig], List[EmbeddingBagConfig]]: A tuple containing
-        two lists - the first for unweighted embedding tables and the second for
-        weighted embedding tables.
-    """
-    tables = [
-        EmbeddingBagConfig(
-            num_embeddings=max(i + 1, 100) * 1000,
-            embedding_dim=embedding_feature_dim,
-            name="table_" + str(i),
-            feature_names=["feature_" + str(i)],
-        )
-        for i in range(num_unweighted_features)
-    ]
-    weighted_tables = [
-        EmbeddingBagConfig(
-            num_embeddings=max(i + 1, 100) * 1000,
-            embedding_dim=embedding_feature_dim,
-            name="weighted_table_" + str(i),
-            feature_names=["weighted_feature_" + str(i)],
-        )
-        for i in range(num_weighted_features)
-    ]
-    return tables, weighted_tables
-
-
 def generate_planner(
     planner_type: str,
     topology: Topology,

torchrec/distributed/benchmark/yaml/sparse_data_dist_base.yml

@@ -13,3 +13,17 @@ EmbeddingTablesConfig:
   num_unweighted_features: 100
   num_weighted_features: 100
   embedding_feature_dim: 128
+  additional_tables:
+    - - name: additional_tables_0_0
+        embedding_dim: 128
+        num_embeddings: 100_000
+        feature_names: ["additional_0_0"]
+      - name: additional_tables_0_1
+        embedding_dim: 128
+        num_embeddings: 100_000
+        feature_names: ["additional_0_1"]
+    - []
+    - - name: additional_tables_2_1
+        embedding_dim: 128
+        num_embeddings: 100_000
+        feature_names: ["additional_2_1"]

torchrec/distributed/test_utils/test_tables.py

@@ -0,0 +1,97 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-strict
+
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional
+
+from torchrec.modules.embedding_configs import EmbeddingBagConfig
+
+
+@dataclass
+class EmbeddingTablesConfig:
+    """
+    Configuration for generating embedding tables for test and benchmark
+
+    This class defines the parameters for generating embedding tables with both weighted
+    and unweighted features.
+
+    Args:
+        num_unweighted_features (int): Number of unweighted features to generate.
+            Default is 100.
+        num_weighted_features (int): Number of weighted features to generate.
+            Default is 100.
+        embedding_feature_dim (int): Dimension of the embedding vectors.
+            Default is 128.
+        additional_tables (List[List[Dict[str, Any]]]): Additional tables to include in the configuration.
+            Default is an empty list.
+    """
+
+    num_unweighted_features: int = 100
+    num_weighted_features: int = 100
+    embedding_feature_dim: int = 128
+    additional_tables: List[List[Dict[str, Any]]] = field(default_factory=list)
+
+    def generate_tables(
+        self,
+    ) -> List[List[EmbeddingBagConfig]]:
+        """
+        Generate embedding bag configurations for both unweighted and weighted features.
+
+        This function creates two lists of EmbeddingBagConfig objects:
+        1. Unweighted tables: Named as "table_{i}" with feature names "feature_{i}"
+        2. Weighted tables: Named as "weighted_table_{i}" with feature names "weighted_feature_{i}"
+
+        For both types, the number of embeddings scales with the feature index,
+        calculated as max(i + 1, 100) * 1000.
+
+        Args:
+            num_unweighted_features (int): Number of unweighted features to generate.
+            num_weighted_features (int): Number of weighted features to generate.
+            embedding_feature_dim (int): Dimension of the embedding vectors.
+
+        Returns:
+            Tuple[List[EmbeddingBagConfig], List[EmbeddingBagConfig]]: A tuple containing
+            two lists - the first for unweighted embedding tables and the second for
+            weighted embedding tables.
+        """
+        unweighted_tables = [
+            EmbeddingBagConfig(
+                num_embeddings=max(i + 1, 100) * 1000,
+                embedding_dim=self.embedding_feature_dim,
+                name="table_" + str(i),
+                feature_names=["feature_" + str(i)],
+            )
+            for i in range(self.num_unweighted_features)
+        ]
+        weighted_tables = [
+            EmbeddingBagConfig(
+                num_embeddings=max(i + 1, 100) * 1000,
+                embedding_dim=self.embedding_feature_dim,
+                name="weighted_table_" + str(i),
+                feature_names=["weighted_feature_" + str(i)],
+            )
+            for i in range(self.num_weighted_features)
+        ]
+        tables_list = []
+        for idx, adts in enumerate(self.additional_tables):
+            if idx == 0:
+                tables = unweighted_tables
+            elif idx == 1:
+                tables = weighted_tables
+            else:
+                tables = []
+            for adt in adts:
+                tables.append(EmbeddingBagConfig(**adt))
+
+        if len(tables_list) == 0:
+            tables_list.append(unweighted_tables)
+            tables_list.append(weighted_tables)
+        elif len(tables_list) == 1:
+            tables_list.append(weighted_tables)
+        return tables_list