ai-edge-torch-nightly 0.1.dev202405131930__py3-none-any.whl

ai_edge_torch/generative/utilities/t5_loader.py

@@ -0,0 +1,467 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# Common utility functions for data loading etc.
+from dataclasses import dataclass
+import glob
+import os
+from typing import Callable, Dict
+
+from safetensors import safe_open
+import torch
+
+from ai_edge_torch.generative.layers import model_config
+
+
+def load_safetensors(full_path: str):
+  """Loads safetensors into a single state dictionary.
+
+  Args:
+    full_path (string): the safetensor filename or directory that contains the
+      safetensor files.
+
+  Returns:
+    A state dictionary contating loaded tensors.
+
+  Raises:
+    ValueError: If no tensors are loaded from the provided directory or file.
+  """
+  pattern = (
+      os.path.join(full_path, "*.safetensors")
+      if os.path.isdir(full_path)
+      else full_path
+  )
+  files = []
+  for file in glob.glob(pattern):
+    files.append(file)
+
+  tensors = {}
+  for file in files:
+    with safe_open(file, framework="pt") as fp:
+      for k in fp.keys():
+        assert k not in tensors
+        tensors[k] = fp.get_tensor(k)
+
+  if not tensors:
+    raise ValueError("Failed to load SafeTensors.")
+  return tensors
+
+
+def load_pytorch_statedict(full_path: str):
+  """Loads state dictionary binaries into a single state dictionary.
+
+  Args:
+    full_path (string): the bin filename or directory that contains the bin
+      files.
+
+  Returns:
+    A state dictionary contating loaded tensors.
+
+  Raises:
+    ValueError: If no tensors are loaded from the provided directory or file.
+  """
+  pattern = os.path.join(full_path, "*.bin") if os.path.isdir(full_path) else full_path
+  files = []
+  for file in glob.glob(pattern):
+    files.append(file)
+
+  tensors = {}
+  for file in files:
+    this_file_tensors = torch.load(file, map_location=torch.device("cpu"))
+    for k in this_file_tensors:
+      assert k not in tensors
+    tensors.update(this_file_tensors)
+
+  if not tensors:
+    raise ValueError("Failed to load torch bin files.")
+  return tensors
+
+
+class ModelLoader:
+  """A utility class for loading and converting model checkpoints to ODML
+  transformer layer format.
+  """
+
+  @dataclass
+  class TensorNames:
+    attn_query_proj: str = None
+    attn_key_proj: str = None
+    attn_value_proj: str = None
+    attn_output_proj: str = None
+    relative_attn_bias: str = None
+
+    cross_attn_query_proj: str = None
+    cross_attn_key_proj: str = None
+    cross_attn_value_proj: str = None
+    cross_attn_output_proj: str = None
+
+    ff_up_proj: str = None
+    ff_down_proj: str = None
+    ff_gate_proj: str = None
+
+    pre_attn_norm: str = None
+    pre_cross_attn_norm: str = None
+    pre_ff_norm: str = None
+    embedding: str = None
+    final_norm: str = None
+    lm_head: str = None
+
+  def __init__(self, file_name: str, names: TensorNames) -> None:
+    """ModelLoader constructor. Can be used to load multiple models of the same
+    type.
+
+    Args:
+        file_name (str): Path to the checkpoint. Can be a directory or an
+          exact file.
+        names (TensorNames): An instance of `TensorNames` to determine mappings.
+    """
+    self._file_name = file_name
+    self._names = names
+    self._loader = self._get_loader()
+
+  def load(
+      self, model: torch.nn.Module, strict: bool = True, fuse_attention: bool = True
+  ):
+    """Load the model from the checkpoint
+
+    Args:
+        model (torch.nn.Module): The pytorch model that needs to be loaded.
+        strict (bool, optional): Whether the converted keys are strictly
+          matched. Defaults to True.
+
+    Raises:
+        ValueError: If conversion results in unmapped tensors and strict mode is
+          enabled.
+    """
+    state = self._loader(self._file_name)
+
+    if isinstance(self._names, ModelLoader.TensorNames):
+      converted_state = self._do_load(
+          model, state, self._names, fuse_attention=fuse_attention
+      )
+    elif isinstance(self._names, dict):
+      converted_state = {}
+      for additional_prefix, names in self._names.items():
+        local_converted_state = self._do_load(
+            model,
+            state,
+            self._names[additional_prefix],
+            additional_prefix,
+            fuse_attention=fuse_attention,
+        )
+        converted_state.update(local_converted_state)
+    else:
+      raise ValueError(f"Unkown type for names: {type(self._names)}")
+
+    if strict and state:
+      raise ValueError(
+          f"Failed to map all tensor. Remaining tensor are: {list(state.keys())}"
+      )
+    model.load_state_dict(converted_state, strict=strict)
+
+  def _do_load(self, model, state, names, additional_prefix="", fuse_attention=True):
+    """Load the model from the checkpoint
+
+    Args:
+        model (torch.nn.Module): The pytorch model that needs to be loaded.
+        state (Dict[str, torch.Tensor]): The pytorch state dictionary
+        names (TensorNames]): The TensorNames for the model we are loading.
+
+    Returns:
+        Dict[str, torch.Tensor]: Map of name to tensor for loading.
+    """
+    converted_state = dict()
+    if names.embedding is not None:
+      converted_state["tok_embedding.weight"] = state.pop(f"{names.embedding}.weight")
+    if names.lm_head is not None:
+      converted_state["lm_head.weight"] = state.pop(f"{names.lm_head}.weight")
+      if model.config.lm_head_use_bias:
+        converted_state["lm_head.bias"] = state.pop(f"{names.lm_head}.bias")
+    if names.final_norm is not None:
+      final_norm_name = names.final_norm
+      prefix = additional_prefix
+      converted_state[f"{prefix}final_norm.weight"] = state.pop(
+          f"{final_norm_name}.weight"
+      )
+      if f"{final_norm_name}.bias" in state:
+        converted_state["final_norm.bias"] = state.pop(f"{final_norm_name}.bias")
+
+    if names.relative_attn_bias:
+      rel_attn_name = names.relative_attn_bias
+      prefix = additional_prefix + f"transformer_blocks.0"
+      converted_state[f"{prefix}.atten_func.relative_attention_bias.weight"] = (
+          state.pop(f"{rel_attn_name}.weight")
+      )
+
+    for i in range(model.config.num_layers):
+      self._map_norm(i, model.config, state, converted_state, names, additional_prefix)
+      self._map_feedforward(
+          i, model.config, state, converted_state, names, additional_prefix
+      )
+      self._map_attention(
+          i,
+          model.config,
+          state,
+          converted_state,
+          names,
+          additional_prefix,
+          fuse_attention,
+      )
+      self._map_cross_attention(
+          i,
+          model.config,
+          state,
+          converted_state,
+          names,
+          additional_prefix,
+          fuse_attention,
+      )
+
+    return converted_state
+
+  def _get_loader(self) -> Callable[[str], Dict[str, torch.Tensor]]:
+    """A best effort method for finding appropriate state loader.
+
+    Raises:
+        ValueError: If it fails to find an appropriate loader.
+
+    Returns:
+        Callable[[str], Dict[str, torch.Tensor]]: State loader to be used.
+    """
+    if os.path.isdir(self._file_name):
+      if glob.glob(os.path.join(self._file_name, "*.safetensors")):
+        return load_safetensors
+      if glob.glob(os.path.join(self._file_name, "*.bin")):
+        return load_pytorch_statedict
+
+    if self._file_name.endswith(".safetensors"):
+      return load_safetensors
+
+    if self._file_name.endswith(".bin"):
+      return load_pytorch_statedict
+
+    raise ValueError(f"File format not supported.")
+
+  def _map_feedforward(
+      self,
+      idx: int,
+      config: model_config.ModelConfig,
+      state: Dict[str, torch.Tensor],
+      converted_state: Dict[str, torch.Tensor],
+      names: TensorNames,
+      additional_prefix: str = "",
+  ):
+    prefix = additional_prefix + f"transformer_blocks.{idx}"
+    if names.ff_up_proj is None or names.ff_down_proj is None:
+      return
+    if config.ff_config.type == model_config.FeedForwardType.SEQUENTIAL:
+      ff_up_proj_name = names.ff_up_proj.format(idx)
+      ff_down_proj_name = names.ff_down_proj.format(idx)
+      converted_state[f"{prefix}.ff.w1.weight"] = state.pop(f"{ff_up_proj_name}.weight")
+      converted_state[f"{prefix}.ff.w2.weight"] = state.pop(
+          f"{ff_down_proj_name}.weight"
+      )
+      if config.ff_config.use_bias:
+        converted_state[f"{prefix}.ff.w1.bias"] = state.pop(f"{ff_up_proj_name}.bias")
+        converted_state[f"{prefix}.ff.w2.bias"] = state.pop(f"{ff_down_proj_name}.bias")
+    else:
+      if names.ff_gate_proj is not None:
+        ff_up_proj_name = names.ff_up_proj.format(idx)
+        ff_down_proj_name = names.ff_down_proj.format(idx)
+        ff_gate_proj_name = names.ff_gate_proj.format(idx)
+        converted_state[f"{prefix}.ff.w3.weight"] = state.pop(
+            f"{ff_up_proj_name}.weight"
+        )
+        converted_state[f"{prefix}.ff.w2.weight"] = state.pop(
+            f"{ff_down_proj_name}.weight"
+        )
+        converted_state[f"{prefix}.ff.w1.weight"] = state.pop(
+            f"{ff_gate_proj_name}.weight"
+        )
+        if config.ff_config.use_bias:
+          converted_state[f"{prefix}.ff.w3.bias"] = state.pop(f"{ff_up_proj_name}.bias")
+          converted_state[f"{prefix}.ff.w2.bias"] = state.pop(
+              f"{ff_down_proj_name}.bias"
+          )
+          converted_state[f"{prefix}.ff.w1.bias"] = state.pop(
+              f"{ff_gate_proj_name}.bias"
+          )
+
+  def _map_attention(
+      self,
+      idx: int,
+      config: model_config.ModelConfig,
+      state: Dict[str, torch.Tensor],
+      converted_state: Dict[str, torch.Tensor],
+      names: TensorNames,
+      additional_prefix: str = "",
+      fuse_attention: bool = True,
+  ):
+    if (
+        names.attn_query_proj is None
+        or names.attn_key_proj is None
+        or names.attn_value_proj is None
+    ):
+      return
+    prefix = additional_prefix + f"transformer_blocks.{idx}"
+    q_name = names.attn_query_proj.format(idx)
+    k_name = names.attn_key_proj.format(idx)
+    v_name = names.attn_value_proj.format(idx)
+    # model.encoder.transformer_blocks[0].atten_func.q.weight
+    if fuse_attention:
+      converted_state[f"{prefix}.atten_func.attn.weight"] = self._fuse_qkv(
+          config,
+          state.pop(f"{q_name}.weight"),
+          state.pop(f"{k_name}.weight"),
+          state.pop(f"{v_name}.weight"),
+      )
+      if config.attn_config.qkv_use_bias:
+        converted_state[f"{prefix}.atten_func.attn.bias"] = self._fuse_qkv(
+            config,
+            state.pop(f"{q_name}.bias"),
+            state.pop(f"{k_name}.bias"),
+            state.pop(f"{v_name}.bias"),
+        )
+    else:
+      converted_state[f"{prefix}.atten_func.q.weight"] = state.pop(f"{q_name}.weight")
+      converted_state[f"{prefix}.atten_func.k.weight"] = state.pop(f"{k_name}.weight")
+      converted_state[f"{prefix}.atten_func.v.weight"] = state.pop(f"{v_name}.weight")
+      if config.attn_config.qkv_use_bias:
+        converted_state[f"{prefix}.atten_func.q.bias"] = state.pop(f"{q_name}.bias")
+        converted_state[f"{prefix}.atten_func.k.bias"] = state.pop(f"{k_name}.bias")
+        converted_state[f"{prefix}.atten_func.v.bias"] = state.pop(f"{v_name}.bias")
+
+    o_name = names.attn_output_proj.format(idx)
+    converted_state[f"{prefix}.atten_func.proj.weight"] = state.pop(f"{o_name}.weight")
+    if config.attn_config.output_proj_use_bias:
+      converted_state[f"{prefix}.atten_func.proj.bias"] = state.pop(f"{o_name}.bias")
+
+  def _map_cross_attention(
+      self,
+      idx: int,
+      config: model_config.ModelConfig,
+      state: Dict[str, torch.Tensor],
+      converted_state: Dict[str, torch.Tensor],
+      names: TensorNames,
+      additional_prefix: str = "",
+      fuse_attention: bool = True,
+  ):
+    if (
+        names.cross_attn_query_proj is None
+        or names.cross_attn_key_proj is None
+        or names.cross_attn_value_proj is None
+    ):
+      return
+    prefix = additional_prefix + f"transformer_blocks.{idx}"
+    q_name = names.cross_attn_query_proj.format(idx)
+    k_name = names.cross_attn_key_proj.format(idx)
+    v_name = names.cross_attn_value_proj.format(idx)
+
+    if fuse_attention:
+      converted_state[f"{prefix}.cross_atten_func.attn.weight"] = self._fuse_qkv(
+          config,
+          state.pop(f"{q_name}.weight"),
+          state.pop(f"{k_name}.weight"),
+          state.pop(f"{v_name}.weight"),
+      )
+      if config.attn_config.qkv_use_bias:
+        converted_state[f"{prefix}.cross_atten_func.attn.bias"] = self._fuse_qkv(
+            config,
+            state.pop(f"{q_name}.bias"),
+            state.pop(f"{k_name}.bias"),
+            state.pop(f"{v_name}.bias"),
+        )
+    else:
+      converted_state[f"{prefix}.cross_atten_func.q.weight"] = state.pop(
+          f"{q_name}.weight"
+      )
+      converted_state[f"{prefix}.cross_atten_func.k.weight"] = state.pop(
+          f"{k_name}.weight"
+      )
+      converted_state[f"{prefix}.cross_atten_func.v.weight"] = state.pop(
+          f"{v_name}.weight"
+      )
+      if config.attn_config.qkv_use_bias:
+        converted_state[f"{prefix}.cross_atten_func.q.bias"] = state.pop(
+            f"{q_name}.bias"
+        )
+        converted_state[f"{prefix}.cross_atten_func.k.bias"] = state.pop(
+            f"{k_name}.bias"
+        )
+        converted_state[f"{prefix}.cross_atten_func.v.bias"] = state.pop(
+            f"{v_name}.bias"
+        )
+
+    o_name = names.cross_attn_output_proj.format(idx)
+    converted_state[f"{prefix}.cross_atten_func.proj.weight"] = state.pop(
+        f"{o_name}.weight"
+    )
+    if config.attn_config.output_proj_use_bias:
+      converted_state[f"{prefix}.cross_atten_func.proj.bias"] = state.pop(
+          f"{o_name}.bias"
+      )
+
+  def _map_norm(
+      self,
+      idx: int,
+      config: model_config.ModelConfig,
+      state: Dict[str, torch.Tensor],
+      converted_state: Dict[str, torch.Tensor],
+      names: TensorNames,
+      additional_prefix: str = "",
+  ):
+    prefix = additional_prefix + f"transformer_blocks.{idx}"
+    if names.pre_attn_norm is not None:
+      pre_attn_norm_name = names.pre_attn_norm.format(idx)
+      converted_state[f"{prefix}.atten_func.pre_atten_norm.weight"] = state.pop(
+          f"{pre_attn_norm_name}.weight"
+      )
+      if f"{pre_attn_norm_name}.bias" in state:
+        converted_state[f"{prefix}.atten_func.pre_atten_norm.bias"] = state.pop(
+            f"{pre_attn_norm_name}.bias"
+        )
+
+    if names.pre_cross_attn_norm:
+      pre_cross_attn_norm_name = names.pre_cross_attn_norm.format(idx)
+      converted_state[f"{prefix}.cross_atten_func.pre_atten_norm.weight"] = state.pop(
+          f"{pre_cross_attn_norm_name}.weight"
+      )
+      if f"{pre_cross_attn_norm_name}.bias" in state:
+        converted_state[f"{prefix}.cross_atten_func.pre_atten_norm.bias"] = state.pop(
+            f"{pre_cross_attn_norm_name}.bias"
+        )
+
+    if names.pre_ff_norm is not None:
+      pre_ff_norm_name = names.pre_ff_norm.format(idx)
+      converted_state[f"{prefix}.pre_ff_norm.weight"] = state.pop(
+          f"{pre_ff_norm_name}.weight"
+      )
+      if f"{pre_ff_norm_name}.bias" in state:
+        converted_state[f"{prefix}.pre_ff_norm.bias"] = state.pop(
+            f"{pre_ff_norm_name}.bias"
+        )
+
+  def _fuse_qkv(
+      self,
+      config: model_config.ModelConfig,
+      q: torch.Tensor,
+      k: torch.Tensor,
+      v: torch.Tensor,
+  ) -> torch.Tensor:
+    q_per_kv = config.attn_config.num_heads // config.attn_config.num_query_groups
+    qs = torch.split(q, config.head_dim * q_per_kv)
+    ks = torch.split(k, config.head_dim)
+    vs = torch.split(v, config.head_dim)
+    cycled = [t for group in zip(qs, ks, vs) for t in group]
+    return torch.cat(cycled)

ai_edge_torch/hlfb/__init__.py

@@ -0,0 +1,16 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from torch_xla.experimental.mark_pattern_utils import StableHLOCompositeBuilder

ai_edge_torch/hlfb/mark_pattern/__init__.py

@@ -0,0 +1,139 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+import copy
+from typing import Any
+import uuid
+
+import torch
+from torch_xla.experimental import xla_marker
+
+from ai_edge_torch.hlfb.mark_pattern.pattern import Pattern
+from ai_edge_torch.hlfb.mark_pattern.pattern import ScalarAttrTracker  # NOQA
+
+
+@torch._dynamo.assume_constant_result
+def _get_uuid() -> str:
+  return uuid.uuid4().hex
+
+
+# TODO: Move to a general fx utils file.
+def _prepose_placeholder_nodes(graph: torch.fx.Graph):
+  nodes = [node for node in graph.nodes if node.op == "placeholder"] + [
+      node for node in graph.nodes if node.op != "placeholder"
+  ]
+
+  for a, b in zip(nodes, nodes[1:]):
+    if a.next is not b:
+      a.append(b)
+  return graph
+
+
+def _insert_marker(
+    graph_module: torch.fx.GraphModule,
+    node: torch.fx.Node,
+    name: str,
+    pos: int,
+    id: str,
+    is_input: bool,
+    attr: dict[str, Any] = None,
+):
+  attr = xla_marker.serialize_composite_attr(attr) if attr else None
+  with graph_module.graph.inserting_after(node):
+    new_node = graph_module.graph.call_function(
+        torch.ops.xla.mark_tensor,
+        args=(node,),
+        kwargs={
+            "name": name,
+            "pos": pos,
+            "id": id,
+            "is_input": is_input,
+            "attr": attr,
+        },
+    )
+
+  new_node.meta = node.meta
+  return new_node
+
+
+def mark_pattern(
+    graph_module: torch.fx.GraphModule,
+    pattern: Pattern,
+) -> torch.fx.GraphModule:
+  """Mark all existences of pattern graph in the GraphModule with fx pattern matching.
+  The marked subgraphs will be lowered in StableHLO composite ops.
+  Args:
+    graph_module (torch.fx.GraphModule): GraphModule to be matched and marked.
+    pattern (ai_edge_torch.hlfb.mark_pattern.Pattern): Pattern to match.
+  Returns:
+    The modified graph_module with additional marker ops in graph.
+  """
+  # Create a copy of graph_module and sanitize it for pattern matching.
+  graph_module_to_match = copy.deepcopy(graph_module)
+  for n, m in zip(graph_module.graph.nodes, graph_module_to_match.graph.nodes):
+    m.meta["ORIGINAL_NODE"] = n
+
+  # Sanitize graph_module to match in the same way as pattern's graph_module.
+  graph_module_to_match = passes.remove_clone_ops(graph_module_to_match)
+
+  match_with_attrs = pattern.match(graph_module_to_match)
+
+  for match, attr in match_with_attrs:
+    match_id = _get_uuid()
+
+    # NOTE: Current graph rewriter (_insert_marker) does not work perfectly
+    # with continuous matches e.g. matching (a + b) on (w + x + y + z). The
+    # rewritten results may be undetermined with false negative - some
+    # matches may not be marked in the lowering, while the marked ones would
+    # always be correct.
+    # TODO(cnchan): completely support mark_pattern on continuous matches.
+    for i, pattern_input_node in enumerate(pattern.input_nodes):
+      input_node = match.nodes_map[pattern_input_node]
+      new_input_node = _insert_marker(
+          graph_module,
+          input_node.meta["ORIGINAL_NODE"],
+          name=pattern.name,
+          pos=i,
+          id=match_id,
+          is_input=True,
+      )
+
+      # Only replace input by the marker node for those nodes used in the pattern.
+      in_pattern_nodes = set(match.nodes_map.values())
+      for user in input_node.users.keys():
+        if user in in_pattern_nodes:
+          user.meta["ORIGINAL_NODE"].replace_input_with(
+              input_node.meta["ORIGINAL_NODE"], new_input_node
+          )
+
+    for i, pattern_output_node in enumerate(pattern.output_nodes):
+      output_node = match.nodes_map[pattern_output_node]
+      new_output_node = _insert_marker(
+          graph_module,
+          output_node.meta["ORIGINAL_NODE"],
+          name=pattern.name,
+          pos=i,
+          id=match_id,
+          is_input=False,
+          attr=attr,  # torch_xla internal: only output marker needs attr.
+      )
+      output_node.meta["ORIGINAL_NODE"].replace_all_uses_with(new_output_node)
+      new_output_node.update_arg(0, output_node.meta["ORIGINAL_NODE"])
+
+  graph_module.graph.eliminate_dead_code()
+  _prepose_placeholder_nodes(graph_module.graph)
+
+  graph_module.graph.lint()
+  graph_module.recompile()
+  return graph_module

ai_edge_torch/hlfb/mark_pattern/passes.py

@@ -0,0 +1,42 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+import torch
+
+
+def remove_clone_ops(gm: torch.fx.GraphModule):
+  # torch export adds additional aten.clone nodes to produce contiguous in memory tensors
+  # depending on tensor sizes for runtime efficiency. However, these unpredictable clone
+  # nodes can break the pattern matching. Thus remove all clones in model and pattern graphs.
+  for node in gm.graph.nodes:
+    if node.op == "call_function" and node.name.startswith("clone"):
+      node.replace_all_uses_with(node.args[0])
+      gm.graph.erase_node(node)
+
+  gm.graph.lint()
+  gm.recompile()
+  return gm
+
+
+def remove_dangling_args(gm: torch.fx.GraphModule):
+  nodes_to_erase = []
+  for node in gm.graph.nodes:
+    if node.op == "placeholder" and len(node.users) == 0:
+      nodes_to_erase.append(node)
+  for node in nodes_to_erase:
+    gm.graph.erase_node(node)
+
+  gm.graph.lint()
+  gm.recompile()
+  return gm