rowwise for feature processors (#3606)

Summary:

In this diff we introduce row based sharding (TWRW, RW, GRID) type support for feature processors.  Previously, feature processors did not support row based sharding since feature processors are data parallel. This means by splitting up the input for row based shards the accessed feature processor weights were in correct. In column/data sharding based approaches, the data is duplicated ensuring the correct weight is accessed across ranks.

The indices/buckets are calculated post input split/distribution, to make it compatible with row based sharding we calculate this pre input split/distribution. This couples the train pipeline and feature processors. For each feature, we preprocess the input and place the calculated indices in KJT.weights, this propagates the indices correctly and indexs into the right weight to use for the final step in the feature processing.

This applies in both pipelined and non pipelined situations - the input modification is done either at the pipelined forward call or in the input dist of the FPEBC. This is determined by the pipelining flag set through rewrite_model in train pipeline.

**Previous versions of this diff were reverted as this change applied to all feature processors regardless of row wise sharding applied which surfaced errors that are not captured in usual E2E and unit tests. We now gate the change in two ways: 1) row based shardings must be specified by users to be applied for FP sharding and 2) pre processing input in pipeline will ONLY happen when row based sharding is present. This way FP sharding without row based shardings applied will go through the original forward path.**

Differential Revision: D88093763

Author: iamzainhuda

torchrec/distributed/fp_embeddingbag.py

@@ -8,11 +8,22 @@
 # pyre-strict
 
 from functools import partial
-from typing import Any, Dict, Iterator, List, Optional, Type, Union
+from typing import (
+    Any,
+    Dict,
+    Iterator,
+    List,
+    Mapping,
+    Optional,
+    Tuple,
+    Type,
+    TypeVar,
+    Union,
+)
 
+import pyjk as justknobs
 import torch
 from torch import nn
-
 from torchrec.distributed.embedding_types import (
     BaseEmbeddingSharder,
     KJTList,
@@ -31,14 +42,20 @@
     ShardingEnv,
     ShardingType,
 )
-from torchrec.distributed.utils import append_prefix, init_parameters
+from torchrec.distributed.utils import (
+    append_prefix,
+    init_parameters,
+    modify_input_for_feature_processor,
+)
 from torchrec.modules.feature_processor_ import FeatureProcessorsCollection
 from torchrec.modules.fp_embedding_modules import (
     apply_feature_processors_to_kjt,
     FeatureProcessedEmbeddingBagCollection,
 )
 from torchrec.sparse.jagged_tensor import KeyedJaggedTensor, KeyedTensor
 
+_T = TypeVar("_T")
+
 
 def param_dp_sync(kt: KeyedTensor, no_op_tensor: torch.Tensor) -> KeyedTensor:
     kt._values.add_(no_op_tensor)
@@ -74,6 +91,16 @@ def __init__(
             )
         )
 
+        self._row_wise_sharded: bool = False
+        for param_sharding in table_name_to_parameter_sharding.values():
+            if param_sharding.sharding_type in [
+                ShardingType.ROW_WISE.value,
+                ShardingType.TABLE_ROW_WISE.value,
+                ShardingType.GRID_SHARD.value,
+            ]:
+                self._row_wise_sharded = True
+                break
+
         self._lookups: List[nn.Module] = self._embedding_bag_collection._lookups
 
         self._is_collection: bool = False
@@ -96,6 +123,12 @@ def __init__(
     def input_dist(
         self, ctx: EmbeddingBagCollectionContext, features: KeyedJaggedTensor
     ) -> Awaitable[Awaitable[KJTList]]:
+        if justknobs.check("pytorch/torchrec:enable_rw_feature_processor"):
+            if not self.is_pipelined and self._row_wise_sharded:
+                # transform input to support row based sharding when not pipelined
+                modify_input_for_feature_processor(
+                    features, self._feature_processors, self._is_collection
+                )
         return self._embedding_bag_collection.input_dist(ctx, features)
 
     def apply_feature_processors_to_kjt_list(self, dist_input: KJTList) -> KJTList:
@@ -166,6 +199,21 @@ def sharded_parameter_names(self, prefix: str = "") -> Iterator[str]:
     def _initialize_torch_state(self, skip_registering: bool = False) -> None:  # noqa
         self._embedding_bag_collection._initialize_torch_state(skip_registering)
 
+    def preprocess_input(
+        self, args: List[_T], kwargs: Mapping[str, _T]
+    ) -> Tuple[List[_T], Mapping[str, _T]]:
+        if self._row_wise_sharded:
+            # only if we are row wise sharded we will pre process the KJT
+            # and call the associated preprocess methods in feature processor
+            for x in args + list(kwargs.values()):
+                if isinstance(x, KeyedJaggedTensor):
+                    modify_input_for_feature_processor(
+                        features=x,
+                        feature_processors=self._feature_processors,
+                        is_collection=self._is_collection,
+                    )
+        return args, kwargs
+
 
 class FeatureProcessedEmbeddingBagCollectionSharder(
     BaseEmbeddingSharder[FeatureProcessedEmbeddingBagCollection]
@@ -236,4 +284,13 @@ def sharding_types(self, compute_device_type: str) -> List[str]:
             ShardingType.TABLE_COLUMN_WISE.value,
         ]
 
+        if justknobs.check("pytorch/torchrec:enable_rw_feature_processor"):
+            types.extend(
+                [
+                    ShardingType.TABLE_ROW_WISE.value,
+                    ShardingType.ROW_WISE.value,
+                    ShardingType.GRID_SHARD.value,
+                ]
+            )
+
         return types

torchrec/distributed/planner/enumerators.py

@@ -54,6 +54,15 @@
     EmbeddingComputeKernel.DRAM_VIRTUAL_TABLE,
 }
 
+# sharding types that require explicit user specification for feature-processed modules
+# row wise sharding uses a different pipelined configuration for feature processing
+# to guard against areas that aren't well tested, user must specify row wise sharding
+GUARDED_SHARDING_TYPES_FOR_FP_MODULES: Set[str] = {
+    ShardingType.ROW_WISE.value,
+    ShardingType.TABLE_ROW_WISE.value,
+    ShardingType.GRID_SHARD.value,
+}
+
 
 class EmbeddingEnumerator(Enumerator):
     """
@@ -186,7 +195,7 @@ def enumerate(
                 sharding_options_per_table: List[ShardingOption] = []
 
                 for sharding_type in self._filter_sharding_types(
-                    name, sharder.sharding_types(self._compute_device)
+                    name, sharder.sharding_types(self._compute_device), sharder_key
                 ):
                     for compute_kernel in self._filter_compute_kernels(
                         name,
@@ -308,18 +317,37 @@ def populate_estimates(self, sharding_options: List[ShardingOption]) -> None:
             estimator.estimate(sharding_options, self._sharder_map)
 
     def _filter_sharding_types(
-        self, name: str, allowed_sharding_types: List[str]
+        self, name: str, allowed_sharding_types: List[str], sharder_key: str = ""
     ) -> List[str]:
         # GRID_SHARD is only supported if specified by user in parameter constraints
+        # ROW based shardings only supported for FP modules if specified by user in parameter constraints
         if not self._constraints or not self._constraints.get(name):
-            return [
+            filtered = [
                 t for t in allowed_sharding_types if t != ShardingType.GRID_SHARD.value
             ]
+            # For feature-processed modules, row-wise sharding types require explicit
+            # user specification due to potential issues with position weighted FPs
+            if "FeatureProcessedEmbeddingBagCollection" in sharder_key:
+                filtered = [
+                    t
+                    for t in filtered
+                    if t not in GUARDED_SHARDING_TYPES_FOR_FP_MODULES
+                ]
+            return filtered
         constraints: ParameterConstraints = self._constraints[name]
         if not constraints.sharding_types:
-            return [
+            filtered = [
                 t for t in allowed_sharding_types if t != ShardingType.GRID_SHARD.value
             ]
+            # For feature-processed modules, row-wise sharding types require explicit
+            # user specification due to potential issues with position weighted FPs
+            if "FeatureProcessedEmbeddingBagCollection" in sharder_key:
+                filtered = [
+                    t
+                    for t in filtered
+                    if t not in GUARDED_SHARDING_TYPES_FOR_FP_MODULES
+                ]
+            return filtered
         constrained_sharding_types: List[str] = constraints.sharding_types
 
         filtered_sharding_types = list(

torchrec/distributed/planner/tests/test_enumerators.py

@@ -1115,3 +1115,76 @@ def test_throw_ex_no_sharding_option_for_table(self) -> None:
             "Module: torchrec.modules.embedding_modules.EmbeddingBagCollection, sharder: CWSharder, compute device: cuda. "
             "To debug, search above for warning logs about no available sharding types/compute kernels for table: table_1",
         )
+
+    def test_filter_sharding_types_fp_ebc_no_constraints(self) -> None:
+        """Test that row-wise sharding types are filtered out for FP modules without constraints."""
+        enumerator = EmbeddingEnumerator(
+            topology=MagicMock(),
+            batch_size=MagicMock(),
+            constraints=None,
+        )
+
+        # Test with FeatureProcessedEmbeddingBagCollection sharder key
+        # Without constraints, row-wise types should be filtered out
+        all_sharding_types = [
+            ShardingType.DATA_PARALLEL.value,
+            ShardingType.TABLE_WISE.value,
+            ShardingType.ROW_WISE.value,
+            ShardingType.TABLE_ROW_WISE.value,
+            ShardingType.COLUMN_WISE.value,
+            ShardingType.GRID_SHARD.value,
+        ]
+
+        allowed_sharding_types = enumerator._filter_sharding_types(
+            "table_0",
+            all_sharding_types,
+            "torchrec.modules.fp_embedding_modules.FeatureProcessedEmbeddingBagCollection",
+        )
+
+        # ROW_WISE, TABLE_ROW_WISE, and GRID_SHARD should be filtered out
+        self.assertEqual(
+            set(allowed_sharding_types),
+            {
+                ShardingType.DATA_PARALLEL.value,
+                ShardingType.TABLE_WISE.value,
+                ShardingType.COLUMN_WISE.value,
+            },
+        )
+
+    def test_filter_sharding_types_fp_ebc_with_row_wise_constraint(self) -> None:
+        """Test that row-wise sharding types are allowed for FP modules with explicit constraints."""
+        constraint = ParameterConstraints(
+            sharding_types=[
+                ShardingType.ROW_WISE.value,
+                ShardingType.TABLE_WISE.value,
+            ],
+        )
+        constraints = {"table_0": constraint}
+        enumerator = EmbeddingEnumerator(
+            topology=MagicMock(),
+            batch_size=MagicMock(),
+            constraints=constraints,
+        )
+
+        all_sharding_types = [
+            ShardingType.DATA_PARALLEL.value,
+            ShardingType.TABLE_WISE.value,
+            ShardingType.ROW_WISE.value,
+            ShardingType.TABLE_ROW_WISE.value,
+            ShardingType.COLUMN_WISE.value,
+        ]
+
+        # With explicit constraint specifying ROW_WISE, it should be allowed
+        allowed_sharding_types = enumerator._filter_sharding_types(
+            "table_0",
+            all_sharding_types,
+            "torchrec.modules.fp_embedding_modules.FeatureProcessedEmbeddingBagCollection",
+        )
+
+        self.assertEqual(
+            set(allowed_sharding_types),
+            {
+                ShardingType.ROW_WISE.value,
+                ShardingType.TABLE_WISE.value,
+            },
+        )