Integrate DLRM to the benchmarking framework (#3168)

Summary:
Pull Request resolved: #3168

Integrate DLRM to the benchmarking framework and set the default arch layer sizes.
The last dense arch layer size should be the same as `embedding_feature_dim`

Reviewed By: aliafzal

Differential Revision: D77896635

fbshipit-source-id: c6557d5e2d2ef2322a48f7daa78d98b6a9607a95

Author: SSYernar

torchrec/distributed/benchmark/benchmark_pipeline_utils.py

@@ -48,7 +48,9 @@
     TrainPipelineSemiSync,
 )
 from torchrec.distributed.types import ModuleSharder, ShardingEnv, ShardingType
+from torchrec.models.dlrm import DLRMWrapper
 from torchrec.modules.embedding_configs import EmbeddingBagConfig
+from torchrec.modules.embedding_modules import EmbeddingBagCollection
 
 
 @dataclass
@@ -203,8 +205,16 @@ def generate_model(
         weighted_tables: List[EmbeddingBagConfig],
         dense_device: torch.device,
     ) -> nn.Module:
-        # TODO: Implement DLRM model generation
-        raise NotImplementedError("DLRM model generation not yet implemented")
+        # DLRM only uses unweighted tables
+        ebc = EmbeddingBagCollection(tables=tables, device=torch.device("meta"))
+
+        return DLRMWrapper(
+            embedding_bag_collection=ebc,
+            dense_in_features=self.num_float_features,
+            dense_arch_layer_sizes=self.dense_arch_layer_sizes,
+            over_arch_layer_sizes=self.over_arch_layer_sizes,
+            dense_device=dense_device,
+        )
 
 
 # pyre-ignore[2]: Missing parameter annotation

torchrec/distributed/benchmark/benchmark_train_pipeline.py

@@ -190,8 +190,8 @@ class ModelSelectionConfig:
     deep_fm_dimension: int = 5
 
     # DLRM specific config
-    dense_arch_layer_sizes: List[int] = field(default_factory=lambda: [20, 10])
-    over_arch_layer_sizes: List[int] = field(default_factory=lambda: [5, 3])
+    dense_arch_layer_sizes: List[int] = field(default_factory=lambda: [20, 128])
+    over_arch_layer_sizes: List[int] = field(default_factory=lambda: [5, 1])
 
 
 @cmd_conf