wip partial sharding

Signed-off-by: greg-kwasniewski1 <[email protected]>

Author: greg-kwasniewski1

setup.py

@@ -215,7 +215,7 @@ def extract_from_precompiled(precompiled_location: str, package_data: List[str],
             precompiled_location = download_precompiled(tempdir, version)
         extract_from_precompiled(precompiled_location, package_data, tempdir)
 
-sanity_check()
+# sanity_check()
 
 # https://setuptools.pypa.io/en/latest/references/keywords.html
 setup(

tensorrt_llm/_torch/auto_deploy/config/default.yaml

@@ -56,6 +56,7 @@ transforms:
     stage: sharding
     simple_shard_only: false
     use_sharding_from_factory: false
+    support_partial_config: true
     sharding_dims: ['tp', 'ep', 'bmm']
   # TODO: (hg) need to ensure run_shape_prop after sharding.
   sharding_transform_executor:

tensorrt_llm/_torch/auto_deploy/llm_args.py

@@ -165,6 +165,13 @@ class AutoDeployConfig(DynamicYamlMixInForSettings, BaseSettings):
         "AutoDeployConfig.",
     )
 
+    support_partial_config: bool = Field(
+        default=False,
+        description="If True, factory sharding will be applied to the subset of transformations  "
+        "that are currently supported. If False, sharding from factory will be performed only if"
+        "all provided transformations are supported.",
+    )
+
     sharding_dims: List[str] = Field(
         default=["tp", "ep", "dp"],
         description="The sharding methods to apply by the heuristic sharding stage.",

tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py

@@ -126,6 +126,7 @@ class ShardingTransformConfig(TransformConfig):
 
     simple_shard_only: bool = Field(default=False)
     use_sharding_from_factory: bool = Field(default=False)
+    support_partial_config: bool = Field(default=False)
     # Which sharding families to run: any subset of {"tp", "ep", "bmm"}
     sharding_dims: List[str] = Field(default_factory=lambda: ["tp", "ep", "bmm"])
 
@@ -184,6 +185,7 @@ def _apply(
             else ShardingConfigSource.UNKNOWN
         )
         shared_config.sharding_config.simple_shard_only = self.config.simple_shard_only
+        shared_config.sharding_config.support_partial_config = self.config.support_partial_config
         shared_config.sharding_config.use_sharding_from_factory = (
             self.config.use_sharding_from_factory
         )

tensorrt_llm/_torch/auto_deploy/utils/sharding_utils.py

@@ -488,6 +488,7 @@ class ShardingConfig(BaseModel):
     predefined_config: Optional[Dict[str, Any]] = None
     simple_shard_only: bool = Field(default=False)
     use_sharding_from_factory: bool = False
+    support_partial_config: bool = False
     sharding_dims: List[str] = Field(default_factory=list)
     tp_transforms: List[TPShardingInfo] = Field(default_factory=list)
     bmm_transforms: List[BMMShardingInfo] = Field(default_factory=list)
@@ -532,7 +533,7 @@ def validate_config(self) -> bool:
         tp_plan = self.predefined_config["tp_plan"]
 
         values = set(tp_plan.values())
-        allowed_values = {
+        supported_modes = {
             "colwise",  # row split and no collective
             "rowwise",  # column split and all-reduce
             "gather",  # simple shard (row + all_gather)
@@ -544,7 +545,7 @@ def validate_config(self) -> bool:
             # "local_packed_rowwise",
             # "local",
         }
-        if not values.issubset(allowed_values):
+        if not self.support_partial_config and not values.issubset(supported_modes):
             ad_logger.warning("Sharding config contains invalid values. Skipping.")
             # invalidate the config
             self.predefined_config = {}