code cleanup

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

Author: greg-kwasniewski1

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

@@ -1325,17 +1325,6 @@ def _shard_parameter_node(
 
     rank, world_size = config.rank, config.world_size
     allreduce_strategy = config.allreduce_strategy.name
-    # num_users = num_users_of_weight_node(node)
-    # if num_users > 1 or num_users == 0:
-    #     ad_logger.warning(
-    #         f"Weight node {node} has {num_users} users. This is not supported for sharding. Skipping."
-    #     )
-    #     return
-    # # get weight and bias key
-    # weight_key, bias_key = extract_param_names_from_node(node)
-
-    # modname = weight_key.rpartition(".")[0]
-    # submod = gm.get_submodule(modname)
 
     # # Shard weight using the unified function (also updates the parameter)
     # original_weight = gm.get_parameter(weight_key)

tensorrt_llm/_torch/auto_deploy/utils/node_utils.py

@@ -269,7 +269,7 @@ def find_get_attr_node(weight_node: Node) -> Node:
 
 def num_users_of_weight_node(node: Node) -> int:
     """Returns the number of users of the weight node of the given parametrized node."""
-    weight_node = extract_weight_nodes(node)[0]
+    weight_node = extract_weight_nodes(node).weights[0].node
     return len(weight_node.users) if weight_node is not None else 0
 
 

tensorrt_llm/_torch/auto_deploy/utils/quantization_utils.py

@@ -117,7 +117,6 @@ def should_skip_quantization(
     else:
         if not (is_linear_op(node_or_name) or is_bmm_op(node_or_name)):
             return True
-        # param_names, _ = extract_param_names_from_node(node_or_name)
         weight_name = extract_weight_name(node_or_name)
         modname = weight_name.rpartition(".")[0]