Update sharding type assertions to support COLUMN_WISE sharding (#3344)

Summary:
Pull Request resolved: #3344

astitled

Reviewed By: Raahul46

Differential Revision: D81543306

fbshipit-source-id: 5160a94ab5842f134a459bb2d76920d940c13e40

Author: jd7-tr

torchrec/distributed/embedding.py

@@ -133,9 +133,10 @@ def get_device_from_parameter_sharding(
     if len(set(device_type_list)) == 1:
         return device_type_list[0]
     else:
-        assert (
-            ps.sharding_type == "row_wise"
-        ), "Only row_wise sharding supports sharding across multiple device types for a table"
+        assert ps.sharding_type in [
+            ShardingType.ROW_WISE.value,
+            ShardingType.COLUMN_WISE.value,
+        ], "Only row_wise or column_wise sharding supports sharding across multiple device types for a table"
         return device_type_list
 
 

torchrec/distributed/quant_embedding.py

@@ -166,9 +166,10 @@ def get_device_from_parameter_sharding(
     if len(set(device_type_list)) == 1:
         return device_type_list[0]
     else:
-        assert (
-            ps.sharding_type == "row_wise"
-        ), "Only row_wise sharding supports sharding across multiple device types for a table"
+        assert ps.sharding_type in [
+            ShardingType.ROW_WISE.value,
+            ShardingType.COLUMN_WISE.value,
+        ], "Only row_wise or column_wise sharding supports sharding across multiple device types for a table"
         return device_type_list