support shared bucketized features for id score list unified embedding table (#2592)

Summary:
Pull Request resolved: #2592

with shared embedding table with multiple features in UE, we need to do sanity check on whether all id score list features in a certain UE group are of same bucketized type.

This diff extended sparse embedding bag preprocessing to enable this check with backward compatibility with single feature scenarios.

Reviewed By: yusuo

Differential Revision: D66314015

fbshipit-source-id: 9ace9d84d545461e5d470db2bb62bfa91eab9b22

Author: Fei Yu

torchrec/distributed/sharding_plan.py

@@ -836,9 +836,10 @@ def construct_module_sharding_plan(
         module, sharder.module_type
     ), f"Incorrect sharder for module type {type(module)}"
     shardable_parameters = sharder.shardable_parameters(module)
-    assert (
-        shardable_parameters.keys() == per_param_sharding.keys()
-    ), "per_param_sharding_config doesn't match the shardable parameters of the module"
+    assert shardable_parameters.keys() == per_param_sharding.keys(), (
+        "per_param_sharding_config doesn't match the shardable parameters of the module,"
+        f"got {list(shardable_parameters.keys())} != {list(per_param_sharding.keys())}"
+    )
 
     local_size = local_size or get_local_size()
     world_size = world_size or dist.get_world_size()