[#10056][fix] AutoDeploy: Handle deletion of nested params in sharding

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

@@ -31,31 +31,13 @@
 
 from ...models.factory import ModelFactory
 from ...shim.interface import CachedSequenceInterface
+from ...utils.attr import del_attr_by_name as _del_attr_by_name
+from ...utils.attr import get_attr_by_name as _get_attr_by_name
+from ...utils.attr import set_attr_by_name as _set_attr_by_name
 from ...utils.logger import ad_logger
 from ...utils.pattern_matcher import ADPatternMatcherPass
 from ..interface import BaseTransform, SharedConfig, TransformInfo, TransformRegistry
 
-
-def _get_attr_by_name(obj, name):
-    for part in name.split("."):
-        obj = getattr(obj, part)
-    return obj
-
-
-def _set_attr_by_name(obj, name, value):
-    parts = name.split(".")
-    for part in parts[:-1]:
-        obj = getattr(obj, part)
-    setattr(obj, parts[-1], value)
-
-
-def _del_attr_by_name(obj, name):
-    parts = name.split(".")
-    for part in parts[:-1]:
-        obj = getattr(obj, part)
-    delattr(obj, parts[-1])
-
-
 _PATTERN_INPUT_NAME = "a_log_like"
 
 

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

@@ -33,6 +33,7 @@
 from ...custom_ops.trtllm_dist import is_trtllm_op_available
 from ...models.factory import ModelFactory, ShardingConfigSource
 from ...shim.interface import CachedSequenceInterface
+from ...utils.attr import del_attr_by_name as _del_attr_by_name
 from ...utils.logger import ad_logger
 from ...utils.node_utils import (
     LayerSubgraph,
@@ -1467,7 +1468,12 @@ def get_partition(lst, world_size, rank):
     for expert in (
         w_up_list_to_remove + w_down_list_to_remove + w_gate_list_to_remove + scales_to_remove
     ):
-        delattr(gm, expert.target)
+        try:
+            _del_attr_by_name(gm, expert.target)
+        except AttributeError:
+            ad_logger.warning(
+                f"Failed to delete unused parameter {expert.target} from GraphModule."
+            )
 
 
 def _slice_expert_dim(gm: GraphModule, tensor_node: Node, lo: int, hi: int) -> Node:

tensorrt_llm/_torch/auto_deploy/utils/attr.py

@@ -0,0 +1,22 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+
+
+def get_attr_by_name(obj, name):
+    for part in name.split("."):
+        obj = getattr(obj, part)
+    return obj
+
+
+def set_attr_by_name(obj, name, value):
+    parts = name.split(".")
+    for part in parts[:-1]:
+        obj = getattr(obj, part)
+    setattr(obj, parts[-1], value)
+
+
+def del_attr_by_name(obj, name):
+    parts = name.split(".")
+    for part in parts[:-1]:
+        obj = getattr(obj, part)
+    delattr(obj, parts[-1])