Back out "Add new argument for sharding type" (#3022)

Summary:
Pull Request resolved: #3022

Rolling back the new `sharding_type` param. Other sharding types do not seem to work properly with the current `shard_quant_model` logic. Instead we are making changes to allow user to provide their own sharding plan for other sharding types (e.g. column-wise sharding).

Original commit changeset: bebefd06f38b

Original Phabricator Diff: D73540534

Reviewed By: aporialiao

Differential Revision: D75727458

fbshipit-source-id: 622962c0025cf5f2103e1f163bdf8747feb5393f

Author: jd7-tr

torchrec/inference/modules.py

@@ -499,7 +499,6 @@ def shard_quant_model(
     device_memory_size: Optional[int] = None,
     constraints: Optional[Dict[str, ParameterConstraints]] = None,
     ddr_cap: Optional[int] = None,
-    sharding_type: ShardingType = ShardingType.TABLE_WISE,
 ) -> Tuple[torch.nn.Module, ShardingPlan]:
     """
     Shard a quantized TorchRec model, used for generating the most optimal model for inference and
@@ -535,10 +534,6 @@ def shard_quant_model(
         quant_model = quantize_inference_model(module)
         sharded_model, _ = shard_quant_model(quant_model)
     """
-    # TODO(T220572301): remove after new sharding types are validated.
-    assert (
-        sharding_type == ShardingType.TABLE_WISE
-    ), "Only table-wise sharding is supported now."
 
     if constraints is None:
         table_fqns = []
@@ -557,7 +552,7 @@ def shard_quant_model(
         constraints = {}
         for name in table_fqns:
             constraints[name] = ParameterConstraints(
-                sharding_types=[sharding_type.value],
+                sharding_types=[ShardingType.TABLE_WISE.value],
                 compute_kernels=[EmbeddingComputeKernel.QUANT.value],
             )