add sparse_data_dist_base.yml config for pipeline benchmark (meta-pytorch#3429)

Summary:
Pull Request resolved: meta-pytorch#3429

# context
* yaml config support was added in previous diff
* here we add the basic commands and usages in README
* also add typical benchmark use cases in the yaml config

Reviewed By: spmex

Differential Revision: D83849104

fbshipit-source-id: bf1328680bf0eefa7b0e05e0703324c75a5502fa

Author: TroyGarden

torchrec/distributed/benchmark/README.md

@@ -0,0 +1,16 @@
+# TorchRec Benchmark
+## 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
+```
+- 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
+```

torchrec/distributed/benchmark/benchmark_train_pipeline.py

@@ -78,6 +78,7 @@ class RunOptions:
             Default is "kjt" (KeyedJaggedTensor).
         profile (str): Directory to save profiling results. If empty, profiling is disabled.
             Default is "" (disabled).
+        profile_name (str): Name of the profiling file. Default is pipeline classname.
         planner_type (str): Type of sharding planner to use. Options are:
             - "embedding": EmbeddingShardingPlanner (default)
             - "hetero": HeteroEmbeddingShardingPlanner
@@ -100,6 +101,7 @@ class RunOptions:
     compute_kernel: EmbeddingComputeKernel = EmbeddingComputeKernel.FUSED
     input_type: str = "kjt"
     profile: str = ""
+    profile_name: str = ""
     planner_type: str = "embedding"
     pooling_factors: Optional[List[float]] = None
     num_poolings: Optional[List[float]] = None
@@ -406,7 +408,11 @@ def _func_to_benchmark(
         pipeline.progress(iter(bench_inputs))
 
         result = benchmark_func(
-            name=type(pipeline).__name__,
+            name=(
+                type(pipeline).__name__
+                if run_option.profile_name == ""
+                else run_option.profile_name
+            ),
             bench_inputs=bench_inputs,  # pyre-ignore
             prof_inputs=bench_inputs,  # pyre-ignore
             num_benchmarks=5,

torchrec/distributed/benchmark/yaml/sparse_data_dist_base.yml

@@ -0,0 +1,15 @@
+# this is a very basic sparse data dist config
+# runs on 2 ranks, showing traces with reasonable workloads
+RunOptions:
+  world_size: 2
+  num_batches: 10
+  sharding_type: table_wise
+  profile: "."
+  profile_name: "sparse_data_dist_base"
+  # export_stacks: True # enable this to export stack traces
+PipelineConfig:
+  pipeline: "sparse"
+EmbeddingTablesConfig:
+  num_unweighted_features: 100
+  num_weighted_features: 100
+  embedding_feature_dim: 128