Validate sharding plan correctness: Shard to rank assignment (#3495)

Summary:
Pull Request resolved: #3495

- validate correct rank assignment
    - checks for None ranks
    - checks for ranks in correct rank range
    - check consistency in Manifold planner sharding plan rank assignment

Reviewed By: aporialiao

Differential Revision: D85878244

fbshipit-source-id: 09847323717609ece5d6480c5cf595f8a5fc4260

Author: isururanawaka

torchrec/distributed/planner/planners.py

@@ -118,6 +118,45 @@ def to_sharding_plan(
     return ShardingPlan(plan)
 
 
+def validate_rank_assignment(sharding_plan: ShardingPlan, topology: Topology) -> None:
+    """
+    Validates that all shards in the given sharding plan have valid rank assignments.
+
+    This function iterates through each module and parameter in the provided sharding plan,
+    checking that each shard's placement has a valid rank (i.e., not None, not negative, and
+    less than the topology's world size). If any shard fails these checks, a PlannerError is raised.
+
+    Args:
+        sharding_plan (ShardingPlan): The sharding plan to validate.
+        topology (Topology): The topology containing world size information.
+
+    Raises:
+        PlannerError: If any shard has an invalid rank assignment or if a sharding spec is missing.
+    """
+    for module_name, module_plan in sharding_plan.plan.items():
+        # pyre-ignore
+        for param_name, param_plan in module_plan.items():
+            if param_plan.sharding_spec is not None:
+                for shard in param_plan.sharding_spec.shards:
+                    if shard.placement.rank() is None or shard.placement.rank() < 0:
+                        msg = f"Rank is not assigned for shard {shard}"
+                        logging.error(msg)
+                        raise PlannerError(
+                            error_type=PlannerErrorType.INVALID_RANK_ASSIGNMENT,
+                            message=msg,
+                        )
+                    if shard.placement.rank() >= topology.world_size:
+                        msg = f"Shard {shard} has rank {shard.placement.rank()} which is greater than world size {dist.get_world_size()}."
+                        logging.error(msg)
+                        raise PlannerError(
+                            error_type=PlannerErrorType.INVALID_RANK_ASSIGNMENT,
+                            message=msg,
+                        )
+            else:
+                msg = f"Sharding spec not found for {module_name}.{param_name}"
+                logging.warning(msg)
+
+
 def extract_plan(
     search_space: List[ShardingOption],
     loaded_sharding_options: Dict[int, ShardingOption],
@@ -634,6 +673,8 @@ def plan(
                         timeout_seconds=self._timeout_seconds,
                     ),
                 )
+
+            validate_rank_assignment(sharding_plan, self._topology)
             return sharding_plan
         else:
             global_storage_capacity = reduce(
@@ -997,6 +1038,7 @@ def plan(
                 sharding_plan = to_sharding_plan(
                     best_plan, self._topology_groups[group]
                 )
+                validate_rank_assignment(sharding_plan, self._topology_groups[group])
                 best_plans.append(sharding_plan)
 
                 end_time = perf_counter()
@@ -1090,4 +1132,5 @@ def plan(
                         + last_planner_error_info,
                     )
 
-        return _merge_plans(best_plans)
+        sharding_plan = _merge_plans(best_plans)
+        return sharding_plan

torchrec/distributed/planner/tests/test_planners.py

@@ -78,6 +78,28 @@ def setUp(self) -> None:
         )
         self.planner = EmbeddingShardingPlanner(topology=self.topology)
 
+    def test_tw_rank_assignment(self) -> None:
+        tables = [
+            EmbeddingBagConfig(
+                num_embeddings=100,
+                embedding_dim=64,
+                name="table_" + str(i),
+                feature_names=["feature_" + str(i)],
+            )
+            for i in range(4)
+        ]
+        model = TestSparseNN(tables=tables, sparse_device=torch.device("meta"))
+        sharding_plan = self.planner.plan(module=model, sharders=[TWSharder()])
+        ranks = [
+            cast(List[int], param_shard.ranks)
+            for param_shard in cast(
+                EmbeddingModuleShardingPlan, sharding_plan.plan["sparse.ebc"]
+            ).values()
+        ]
+        for rank_list in ranks:
+            for rank in rank_list:
+                self.assertTrue(0 <= rank <= 1, f"Rank {rank} not in [0,1]")
+
     def test_tw_solution(self) -> None:
         tables = [
             EmbeddingBagConfig(

torchrec/distributed/planner/types.py

@@ -799,6 +799,7 @@ class PlannerErrorType(Enum):
     OTHER = "other"
     PLANNER_INPUT_CONTEXT_MISMATCH = "planner_input_context_mismatch"
     PLAN_LOADING_FAILED = "plan_loading_failed"
+    INVALID_RANK_ASSIGNMENT = "invalid_rank_assignment"
 
 
 class PlannerError(Exception):