From c663363266357a129c766e7fe65ebf614893aea2 Mon Sep 17 00:00:00 2001 From: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com> Date: Tue, 9 Sep 2025 04:21:36 -0700 Subject: [PATCH 1/2] Squashed commit for the partial sharding support Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com> --- .../_torch/auto_deploy/config/default.yaml | 3 +- tensorrt_llm/_torch/auto_deploy/llm_args.py | 30 ++++---- .../auto_deploy/transform/library/sharding.py | 69 +++++++++++++++++-- .../auto_deploy/utils/sharding_utils.py | 5 +- 4 files changed, 81 insertions(+), 26 deletions(-) diff --git a/tensorrt_llm/_torch/auto_deploy/config/default.yaml b/tensorrt_llm/_torch/auto_deploy/config/default.yaml index 451f18b471a4..25f1f44e33b1 100644 --- a/tensorrt_llm/_torch/auto_deploy/config/default.yaml +++ b/tensorrt_llm/_torch/auto_deploy/config/default.yaml @@ -56,7 +56,8 @@ transforms: stage: sharding simple_shard_only: false use_sharding_from_factory: false - sharding_dims: ['tp', 'ep', 'dp'] + support_partial_config: false + sharding_dims: ['tp', 'ep', 'bmm'] # TODO: (hg) need to ensure run_shape_prop after sharding. sharding_transform_executor: stage: sharding diff --git a/tensorrt_llm/_torch/auto_deploy/llm_args.py b/tensorrt_llm/_torch/auto_deploy/llm_args.py index 617f5b58fc37..ea7131ab39a8 100644 --- a/tensorrt_llm/_torch/auto_deploy/llm_args.py +++ b/tensorrt_llm/_torch/auto_deploy/llm_args.py @@ -57,7 +57,7 @@ class AutoDeployConfig(DynamicYamlMixInForSettings, BaseSettings): description="The path to the model checkpoint or the model name from the Hugging Face Hub." ) - model_factory: Literal["AutoModelForCausalLM", "AutoModelForImageTextToText"] = Field( + model_factory: str = Field( default="AutoModelForCausalLM", description="The model factory to use for loading the model.", ) @@ -153,23 +153,6 @@ class AutoDeployConfig(DynamicYamlMixInForSettings, BaseSettings): description="The fraction of available memory to allocate for cache.", ) - simple_shard_only: bool = Field( - default=False, - description="If True, force simple sharding (all_gather) in tensor parallelism. " - "If False, auto-detect and use column+row (all_reduce) sharding when possible.", - ) - - use_sharding_from_factory: bool = Field( - default=False, - description="If True, use sharding from the model factory. If False, use sharding from the " - "AutoDeployConfig.", - ) - - sharding_dims: List[str] = Field( - default=["tp", "ep", "dp"], - description="The sharding methods to apply by the heuristic sharding stage.", - ) - compile_backend: Literal["torch-simple", "torch-compile", "torch-cudagraph", "torch-opt"] = ( Field( default="torch-compile", @@ -211,6 +194,17 @@ def update_attn_page_size(self): self.attn_page_size = self.max_seq_len return self + @field_validator("model_factory", mode="after") + @classmethod + def model_factory_exists(cls, value: str) -> str: + if not ModelFactoryRegistry.has(value): + raise ValueError( + f"'{value}' does not exist in the model factory registry. Available values: " + f"{ModelFactoryRegistry.entries()}." + ) + + return value + ### UTILITY METHODS ############################################################################ def create_factory(self) -> ModelFactory: """Create a model factory from the arguments.""" diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py b/tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py index c37b627240ff..3d467411813c 100644 --- a/tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py +++ b/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,9 @@ 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.sharding_dims = self.config.sharding_dims + shared_config.sharding_config.use_sharding_from_factory = ( self.config.use_sharding_from_factory ) @@ -199,8 +203,6 @@ def _apply( factory_info = detect_sharding_from_factory_config(gm, sharding_config) return gm, factory_info - shared_config.sharding_config.sharding_dims = self.config.sharding_dims - ad_logger.info( f"Running autodeploy sharding heuristics: {shared_config.sharding_config.sharding_dims}" ) @@ -337,8 +339,39 @@ def detect_sharding_from_factory_config( # TODO: Sequence parallelism is not supported yet. ad_logger.warning("Sequence parallelism is not supported yet. Skipping.") elif "local" in config: - # TODO: local refers to hybrid EP+TP parallelism. Not supported yet. - ad_logger.warning("Local EP+TP sharding is not supported yet. Skipping.") + # Check if this applies to shared experts in EP parallelism. + # If yes, apply the TP col-row shard. + if "shared" in module_name: + col_row_action = config.replace("local_", "") + if col_row_action == "colwise": + sharding_config.tp_transforms.append( + TPShardingInfo( + target_node=lin_node.name, + split_dim=SplitDimension.COLUMN, + rank=rank, + world_size=world_size, + dist_op=None, + min_local_shape=min_local_shape, + ) + ) + elif col_row_action == "rowwise": + sharding_config.tp_transforms.append( + TPShardingInfo( + target_node=lin_node.name, + split_dim=SplitDimension.ROW, + rank=rank, + world_size=world_size, + dist_op="all_reduce", + min_local_shape=min_local_shape, + ) + ) + num_row_col_shards += 1 + else: + ad_logger.warning("Invalid sharding config. Skipping.") + else: + # TODO: local refers to hybrid EP+TP parallelism. Not supported yet. + ad_logger.warning("Local EP+TP sharding is not supported yet. Skipping.") + elif "gather" in config: # Simple shard (row + all_gather) sharding_config.tp_transforms.append( @@ -361,9 +394,35 @@ def detect_sharding_from_factory_config( f"Applied {num_shards} TP shards (simple: {num_simple_shards}, " f"row-col pattern: {num_row_col_shards})" ) + + num_matches = len(sharding_config.tp_transforms) + + if sharding_config.support_partial_config: + ad_logger.info( + f"Partial factory config applied only for TP. " + f"Applying heuristics for {sharding_config.sharding_dims}." + ) + + # run EP sharding across ranks + if "ep" in sharding_config.sharding_dims: + ep_info = detect_ep_shard(gm, sharding_config) + else: + ep_info = TransformInfo( + skipped=True, num_matches=0, is_clean=True, has_valid_shapes=True + ) + + # run BMM sharding across ranks + if "bmm" in sharding_config.sharding_dims: + dp_bmm_info = detect_dp_bmm_shard(gm, sharding_config) + else: + dp_bmm_info = TransformInfo( + skipped=True, num_matches=0, is_clean=True, has_valid_shapes=True + ) + num_matches += ep_info.num_matches + dp_bmm_info.num_matches + return TransformInfo( skipped=False, - num_matches=len(sharding_config.tp_transforms), + num_matches=num_matches, is_clean=False, has_valid_shapes=False, ) diff --git a/tensorrt_llm/_torch/auto_deploy/utils/sharding_utils.py b/tensorrt_llm/_torch/auto_deploy/utils/sharding_utils.py index 7f9833a2d18e..f77802f823ed 100644 --- a/tensorrt_llm/_torch/auto_deploy/utils/sharding_utils.py +++ b/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 = {} From 46d629bbcd00e9ac65869b9671b52b7bb05340e4 Mon Sep 17 00:00:00 2001 From: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com> Date: Wed, 10 Sep 2025 05:55:13 -0700 Subject: [PATCH 2/2] brought back sharding args Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com> --- tensorrt_llm/_torch/auto_deploy/llm_args.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/tensorrt_llm/_torch/auto_deploy/llm_args.py b/tensorrt_llm/_torch/auto_deploy/llm_args.py index ea7131ab39a8..a46b4181ca8a 100644 --- a/tensorrt_llm/_torch/auto_deploy/llm_args.py +++ b/tensorrt_llm/_torch/auto_deploy/llm_args.py @@ -153,6 +153,23 @@ class AutoDeployConfig(DynamicYamlMixInForSettings, BaseSettings): description="The fraction of available memory to allocate for cache.", ) + simple_shard_only: bool = Field( + default=False, + description="If True, force simple sharding (all_gather) in tensor parallelism. " + "If False, auto-detect and use column+row (all_reduce) sharding when possible.", + ) + + use_sharding_from_factory: bool = Field( + default=False, + description="If True, use sharding from the model factory. If False, use sharding from the " + "AutoDeployConfig.", + ) + + sharding_dims: List[str] = Field( + default=["tp", "ep", "dp"], + description="The sharding methods to apply by the heuristic sharding stage.", + ) + compile_backend: Literal["torch-simple", "torch-compile", "torch-cudagraph", "torch-opt"] = ( Field( default="torch-compile",