diff --git a/tensorrt_llm/_torch/auto_deploy/config/default.yaml b/tensorrt_llm/_torch/auto_deploy/config/default.yaml index 4200be6c6aab..bdc97520fac3 100644 --- a/tensorrt_llm/_torch/auto_deploy/config/default.yaml +++ b/tensorrt_llm/_torch/auto_deploy/config/default.yaml @@ -81,6 +81,7 @@ transforms: sharding_source: ['manual', 'factory', 'heuristic'] support_partial_config: true sharding_dims: ['tp', 'ep', 'bmm'] + shard_all_unprocessed: true allreduce_strategy: 'NCCL' dist_backend: auto requires_shape_prop: true diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py b/tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py index bae85f3a2216..044cf7002079 100644 --- a/tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py @@ -140,6 +140,11 @@ class ShardingTransformConfig(TransformConfig): sharding_dims: List[ShardingDim] = Field( default_factory=lambda: [ShardingDim.TP, ShardingDim.EP, ShardingDim.BMM] ) + shard_all_unprocessed: bool = Field( + default=False, + description="When True, apply simple shard (column split + all_gather) to " + "'leftover' linear nodes that are not part of any layer subgraph.", + ) allreduce_strategy: AllReduceStrategy = Field( default=AllReduceStrategy.AUTO, description="AllReduce strategy for distributed operations. " @@ -2608,7 +2613,8 @@ def detect_column_row_shard( num_attention_shards += 1 # simple shard remaining linear nodes - num_simple_shards += _process_simple_shard(unprocessed_linear_nodes, transform_container) + if config.shard_all_unprocessed: + num_simple_shards += _process_simple_shard(unprocessed_linear_nodes, transform_container) num_column_row_shards += num_ssm_shards num_shards = num_simple_shards + num_column_row_shards ad_logger.info(