Removed redundant default value initialization in benchmarking script

torchrec/distributed/benchmark/benchmark_pipeline_utils.py

@@ -7,6 +7,15 @@
 
 # pyre-strict
 
+"""
+Utilities for benchmarking training pipelines with different model configurations.
+
+Adding New Model Support:
+    1. Create config class inheriting from BaseModelConfig with generate_model() method
+    2. Add the model to model_configs dict in create_model_config()
+    3. Add model-specific params to ModelSelectionConfig and create_model_config's arguments in benchmark_train_pipeline.py
+"""
+
 import copy
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, fields
@@ -119,8 +128,8 @@ def generate_model(
 class TestTowerSparseNNConfig(BaseModelConfig):
     """Configuration for TestTowerSparseNN model."""
 
-    embedding_groups: Optional[Dict[str, List[str]]] = None
-    feature_processor_modules: Optional[Dict[str, torch.nn.Module]] = None
+    embedding_groups: Optional[Dict[str, List[str]]]
+    feature_processor_modules: Optional[Dict[str, torch.nn.Module]]
 
     def generate_model(
         self,
@@ -143,8 +152,8 @@ def generate_model(
 class TestTowerCollectionSparseNNConfig(BaseModelConfig):
     """Configuration for TestTowerCollectionSparseNN model."""
 
-    embedding_groups: Optional[Dict[str, List[str]]] = None
-    feature_processor_modules: Optional[Dict[str, torch.nn.Module]] = None
+    embedding_groups: Optional[Dict[str, List[str]]]
+    feature_processor_modules: Optional[Dict[str, torch.nn.Module]]
 
     def generate_model(
         self,

torchrec/distributed/benchmark/benchmark_train_pipeline.py

@@ -7,7 +7,18 @@
 
 # pyre-strict
 
-#!/usr/bin/env python3
+"""
+Example usage:
+
+Buck2 (internal):
+    buck2 run @fbcode//mode/opt fbcode//torchrec/distributed/benchmark:benchmark_train_pipeline -- --world_size=2 --pipeline=sparse --batch_size=10
+
+OSS (external):
+    python -m torchrec.distributed.benchmark.benchmark_train_pipeline --world_size=4 --pipeline=sparse --batch_size=10
+
+Adding New Model Support:
+    See benchmark_pipeline_utils.py for step-by-step instructions.
+"""
 
 from dataclasses import dataclass, field
 from typing import Dict, List, Optional, Type, Union