Support weighted_bwd_compute_multiplier in sharding estimators

Differential Revision: D53550851

Author: ys97529

torchrec/distributed/planner/constants.py

@@ -34,6 +34,7 @@
 HALF_BLOCK_PENALTY: float = 1.15  # empirical studies
 QUARTER_BLOCK_PENALTY: float = 1.75  # empirical studies
 BWD_COMPUTE_MULTIPLIER: float = 2  # empirical studies
+WEIGHTED_FEATURE_BWD_COMPUTE_MULTIPLIER: float = 1  # empirical studies
 WEIGHTED_KERNEL_MULTIPLIER: float = 1.1  # empirical studies
 DP_ELEMENTWISE_KERNELS_PERF_FACTOR: float = 9.22  # empirical studies
 

torchrec/distributed/planner/shard_estimators.py

@@ -216,6 +216,7 @@ def estimate(
                 intra_host_bw=self._topology.intra_host_bw,
                 inter_host_bw=self._topology.inter_host_bw,
                 bwd_compute_multiplier=self._topology.bwd_compute_multiplier,
+                weighted_feature_bwd_compute_multiplier=self._topology.weighted_feature_bwd_compute_multiplier,
                 is_pooled=sharding_option.is_pooled,
                 is_weighted=is_weighted,
                 is_inference=self._is_inference,
@@ -251,6 +252,7 @@ def perf_func_emb_wall_time(
         intra_host_bw: float,
         inter_host_bw: float,
         bwd_compute_multiplier: float,
+        weighted_feature_bwd_compute_multiplier: float,
         is_pooled: bool,
         is_weighted: bool = False,
         caching_ratio: Optional[float] = None,
@@ -336,6 +338,7 @@ def perf_func_emb_wall_time(
                     inter_host_bw=inter_host_bw,
                     intra_host_bw=intra_host_bw,
                     bwd_compute_multiplier=bwd_compute_multiplier,
+                    weighted_feature_bwd_compute_multiplier=weighted_feature_bwd_compute_multiplier,
                     is_pooled=is_pooled,
                     is_weighted=is_weighted,
                     is_inference=is_inference,
@@ -447,6 +450,7 @@ def _get_tw_sharding_perf(
         inter_host_bw: float,
         intra_host_bw: float,
         bwd_compute_multiplier: float,
+        weighted_feature_bwd_compute_multiplier: float,
         is_pooled: bool,
         is_weighted: bool = False,
         is_inference: bool = False,
@@ -507,6 +511,8 @@ def _get_tw_sharding_perf(
 
         # includes fused optimizers
         bwd_compute = fwd_compute * bwd_compute_multiplier
+        if is_weighted:
+            bwd_compute = bwd_compute * weighted_feature_bwd_compute_multiplier
 
         prefetch_compute = cls._get_expected_cache_prefetch_time(
             ddr_mem_bw, expected_cache_fetches, emb_dim, table_data_type_size

torchrec/distributed/planner/types.py

@@ -25,6 +25,7 @@
     HBM_MEM_BW,
     INTRA_NODE_BANDWIDTH,
     POOLING_FACTOR,
+    WEIGHTED_FEATURE_BWD_COMPUTE_MULTIPLIER,
 )
 from torchrec.distributed.types import (
     BoundsCheckMode,
@@ -186,6 +187,7 @@ def __init__(
         inter_host_bw: float = CROSS_NODE_BANDWIDTH,
         bwd_compute_multiplier: float = BWD_COMPUTE_MULTIPLIER,
         custom_topology_data: Optional[CustomTopologyData] = None,
+        weighted_feature_bwd_compute_multiplier: float = WEIGHTED_FEATURE_BWD_COMPUTE_MULTIPLIER,
     ) -> None:
         """
         Representation of a network of devices in a cluster.
@@ -238,6 +240,9 @@ def __init__(
         self._inter_host_bw = inter_host_bw
         self._bwd_compute_multiplier = bwd_compute_multiplier
         self._custom_topology_data = custom_topology_data
+        self._weighted_feature_bwd_compute_multiplier = (
+            weighted_feature_bwd_compute_multiplier
+        )
 
     @property
     def compute_device(self) -> str:
@@ -275,6 +280,10 @@ def inter_host_bw(self) -> float:
     def bwd_compute_multiplier(self) -> float:
         return self._bwd_compute_multiplier
 
+    @property
+    def weighted_feature_bwd_compute_multiplier(self) -> float:
+        return self._weighted_feature_bwd_compute_multiplier
+
     def __repr__(self) -> str:
         topology_repr: str = f"world_size={self._world_size} \n"
         topology_repr += f"compute_device={self._compute_device}\n"