fix coreml partitioner for llama (#2816)

Summary: Pull Request resolved: #2816

Differential Revision: D55663281

Author: cccclai

backends/apple/coreml/partition/coreml_partitioner.py

@@ -21,6 +21,7 @@
 from torch.export.exported_program import ExportedProgram
 from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner
 from torch.fx.passes.operator_support import OperatorSupportBase
+from torch._export.utils import is_buffer, is_lifted_tensor_constant, is_param
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.WARNING)
@@ -73,7 +74,6 @@ def __init__(
     def partition(self, exported_program: ExportedProgram) -> PartitionResult:
         # Run the CapabilityBasedPartitioner to return the largest possible
         # subgraphs containing the nodes with the tags
-        logger.info("CoreMLPartitioner::partition")
         partition_tags = {}
 
         capability_partitioner = CapabilityBasedPartitioner(
@@ -87,8 +87,20 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult:
                 tag = f"tag{partition.id}"
                 node.meta["delegation_tag"] = tag
                 partition_tags[tag] = self.delegation_spec
-
-        tag_constant_data(exported_program)
+                node_args = node.args
+                for arg in node_args:
+                    # tag the data nodes if they're used by coreml backend
+                    if isinstance(arg, torch.fx.Node):
+                        is_attr = (
+                            arg.op == "placeholder"
+                            and (
+                                is_param(exported_program, arg)
+                                or is_buffer(exported_program, arg)
+                                or is_lifted_tensor_constant(exported_program, arg)
+                            )
+                        ) or (arg.op == "get_attr")
+                        if is_attr:
+                            arg.meta["delegation_tag"] = tag
 
         return PartitionResult(
             tagged_exported_program=exported_program, partition_tags=partition_tags