small refactor

Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com>

Author: galagam

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

@@ -31,9 +31,7 @@
 
 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._graph import del_attr_by_name, get_attr_by_name, set_attr_by_name
 from ...utils.logger import ad_logger
 from ...utils.pattern_matcher import ADPatternMatcherPass
 from ..interface import BaseTransform, SharedConfig, TransformInfo, TransformRegistry
@@ -64,21 +62,21 @@ def _ensure_a_fused_param(gm: GraphModule, param_name: str) -> Optional[str]:
 
     new_param_name = param_name.replace("A_log", "A_fused")
     try:
-        _get_attr_by_name(gm, new_param_name)
+        get_attr_by_name(gm, new_param_name)
         return new_param_name
     except AttributeError:
         pass
 
     try:
-        a_log = _get_attr_by_name(gm, param_name)
+        a_log = get_attr_by_name(gm, param_name)
     except AttributeError:
         ad_logger.warning(f"Could not find attribute {param_name} in gm.")
         return None
 
     with torch.no_grad():
         a_fused = -torch.exp(a_log.float())
 
-    _set_attr_by_name(
+    set_attr_by_name(
         gm,
         new_param_name,
         nn.Parameter(a_fused, requires_grad=False),
@@ -102,7 +100,7 @@ def _maybe_remove(name: str) -> None:
         if not name.endswith("A_log") or name in used_a_log_targets:
             return
         try:
-            _del_attr_by_name(gm, name)
+            del_attr_by_name(gm, name)
             removed = True
         except AttributeError:
             ad_logger.warning(f"Failed to delete unused parameter {name} from GraphModule.")

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

@@ -33,7 +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._graph import del_attr_by_name
 from ...utils.logger import ad_logger
 from ...utils.node_utils import (
     LayerSubgraph,
@@ -1469,7 +1469,7 @@ def get_partition(lst, world_size, rank):
         w_up_list_to_remove + w_down_list_to_remove + w_gate_list_to_remove + scales_to_remove
     ):
         try:
-            _del_attr_by_name(gm, expert.target)
+            del_attr_by_name(gm, expert.target)
         except AttributeError:
             ad_logger.warning(
                 f"Failed to delete unused parameter {expert.target} from GraphModule."

tensorrt_llm/_torch/auto_deploy/utils/_graph.py

@@ -401,3 +401,23 @@ def get_lm_head_weights(model: nn.Module) -> torch.Tensor:
     gm, output_node = get_output_node(model)
     lm_head_node = get_lm_head_node(gm, output_node)
     return get_weight_tensor(gm, lm_head_node)
+
+
+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])