wip

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

Author: greg-kwasniewski1

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

@@ -1344,6 +1344,7 @@ def _shard_parameter_node(
 
     # Shard weight using the unified function (also updates the parameter)
     original_weight = gm.get_parameter(weight_key)
+
     _, weight_new_shape = shard_weight_tensor(
         gm=gm,
         weight_tensor=original_weight,
@@ -1892,6 +1893,8 @@ def _process_ssm_sharding(
         if "out_proj" not in str(n)
     ]
     for weight_node in weight_nodes:
+        # if is_any_ssm_op(list(weight_node.users)[0]):
+        #     continue
         weight_key = weight_node.target
         # Get the weight parameter
         try:

tensorrt_llm/_torch/auto_deploy/utils/node_utils.py

@@ -131,6 +131,8 @@ def get_quantization_params_from_linear_node(linear_op: torch.fx.node.Node):
 
 def extract_weight_node(node: Node) -> int:
     """Extracts the weight node from the given parametrized node"""
+    gm = node.graph.owning_module
+    param_names = {name for name, _ in gm.named_parameters()}
 
     def find_get_attr_node(weight_node: Node) -> Node:
         """Recursively traverse inputs of allowed nodes to find a node with 'get_attr' op."""
@@ -141,7 +143,7 @@ def find_get_attr_node(weight_node: Node) -> Node:
             torch.ops.aten.view.default,
         }
 
-        if weight_node.op == "get_attr":
+        if weight_node.op == "get_attr" and weight_node.target in param_names:
             return weight_node
 
         # If node is not in the list of allowable ops then return None
@@ -325,6 +327,15 @@ def is_any_ssm_op(node: Node) -> bool:
     )
 
 
+def is_any_conv_op(node: Node) -> bool:
+    return is_op(
+        node,
+        ops=[
+            torch.ops.auto_deploy.torch_causal_conv1d,
+        ],
+    )
+
+
 def is_any_attention_op(node: Node) -> bool:
     return is_op(
         node,