fixed latent MoE sharding

Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>

Author: greg-kwasniewski1

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

@@ -1246,6 +1246,7 @@ def _shard_parameter_node(
 
     # Shard weight using the unified function (also updates the parameter)
     weight_nodes = extract_weight_nodes(node)
+
     for weight_node in weight_nodes.weights:
         _, weight_new_shape = shard_weight_tensor(
             gm=gm,
@@ -2237,7 +2238,7 @@ def detect_column_row_shard(
         attention_nodes = list(filtered_nodes(layer_subgraph, is_any_attention_op))
         min_local_shape = 1
 
-        if config.simple_shard_only:
+        if config.simple_shard_only or layer.layer_type == LayerType.UNKNOWN:
             ad_logger.debug(
                 f"Forcing Simple Shard on nodes: {nodes_linear} with layer type: {layer.layer_type}"
             )

tensorrt_llm/_torch/auto_deploy/utils/node_utils.py

@@ -535,6 +535,9 @@ def get_all_layer_subgraphs(gm: GraphModule) -> List[List[Node]]:
     assert gm.graph.nodes, "Graph is empty"
     layer_subgraphs = []
     linear_nodes = list(filtered_nodes(gm.graph.nodes, is_any_lin_op))
+
+    # find the embedding size of this model. Extract it from the input of the first linear node.
+    embd = get_weight_shape(linear_nodes[0], dim=-1)
     unprocessed_linear_nodes = set(linear_nodes)
     assert len(linear_nodes) > 0, "Could not find any linear nodes in the graph"
 
@@ -546,7 +549,7 @@ def get_all_layer_subgraphs(gm: GraphModule) -> List[List[Node]]:
         # opening is the list of linear nodes
         # layer_subgraph is the list of nodes between the opening and closing linear nodes
         # closing is the last linear node in the layer
-        layer_subgraph = get_layer_after_linear_node(linear_nodes, terminating_indices)
+        layer_subgraph = get_layer_after_linear_node(linear_nodes, terminating_indices, embd=embd)
         if layer_subgraph.opening_nodes is not None and len(layer_subgraph.opening_nodes) > 0:
             unprocessed_linear_nodes -= (
                 set(layer_subgraph.opening_nodes)
@@ -812,6 +815,7 @@ def get_weight_shape(node: Node, dim: Optional[int] = None) -> Optional[Union[in
 def get_layer_after_linear_node(
     linear_nodes: List[Node],
     terminating_indices: List[int],
+    embd: Optional[int] = None,
     match_on_shapes: bool = True,
     enforce_strict_linear_history: bool = True,
 ) -> LayerSubgraph:
@@ -886,8 +890,9 @@ def filter_condition(node: Node, embd: Optional[int] = None, dim: Optional[int]
                 layer_type=LayerType.UNKNOWN,
             )
         if match_on_shapes:
-            # get embedding size of the opening linear node
-            embd = get_weight_shape(linear_nodes[start_lin_index], dim=-1)
+            if embd is None:
+                # get embedding size of the opening linear node
+                embd = get_weight_shape(linear_nodes[start_lin_index], dim=-1)
             # partial init boundary_condition and filter_condition
             boundary_condition = partial(boundary_condition, embd=embd, dim=0)
             filter_condition = partial(filter_condition, embd=embd, dim=0)
@@ -940,7 +945,12 @@ def filter_condition(node: Node, embd: Optional[int] = None, dim: Optional[int]
     ssm_nodes = list(filtered_nodes(interior_nodes, is_any_ssm_op))
     attention_nodes = list(filtered_nodes(interior_nodes, is_any_attention_op))
     intermediate_lin_nodes = list(filtered_nodes(interior_nodes, is_any_lin_op))
-    intermediate_weight_nodes = list(filtered_nodes(interior_nodes, is_weight_node))
+    intermediate_weight_nodes = list(
+        filtered_nodes(
+            interior_nodes, lambda n: is_weight_node(n) and
+            not is_any_lin_op(list(n.users)[0])
+        )
+    )
 
     layer_type = LayerType.MLP
     min_local_shape = 1