From 60624a5f4208f6bee6d19c52176c82d6024560b7 Mon Sep 17 00:00:00 2001 From: lucylq Date: Wed, 6 Aug 2025 11:21:05 -0700 Subject: [PATCH 1/2] Use unlifted export pass to tag delegated constants Use the unlifted pass to tag constants for delegates. Implications: - Tagging must happen on the unlifted ep.module(), before going into to_edge_transform_and_lower/to_edge. Why? - The unlifted graph contains constants in getattr nodes, which is a convenient way to isolate constants. - After going into to_edge_transform_and_lower/to_edge, transforms happen on the graph_module, which is lifted. - The lifted graph requires the ep graph signature to differentiate constants via the `is_param` function. However, in to_edge.transform, we do not have access to the ep. Baking the ep as an argument via partial function doesn't work, as the ep from earlier may be outdated. This means we are comparing an older ep to a newer graph_module, which may not have corresponding graph signatures etc. Differential Revision: [D79736684](https://our.internmc.facebook.com/intern/diff/D79736684/) [ghstack-poisoned] --- docs/source/using-executorch-export.md | 12 ++++++---- exir/passes/external_constants_pass.py | 31 ------------------------- test/models/export_delegated_program.py | 13 ++++------- 3 files changed, 12 insertions(+), 44 deletions(-) diff --git a/docs/source/using-executorch-export.md b/docs/source/using-executorch-export.md index da9cadf3ec2..aacff4e6fd5 100644 --- a/docs/source/using-executorch-export.md +++ b/docs/source/using-executorch-export.md @@ -129,14 +129,16 @@ To generate a `model.pte`, `model.ptd` pair with the weights inside `model.ptd`, ```python from executorch.exir.passes.external_constants_pass import ( - delegate_external_constants_pass, + delegate_external_constants_pass_unlifted, ) -partial_function = partial( - delegate_external_constants_pass, - ep=exported_program, +# Tag the unlifted ep.module(). +tagged_module = exported_program.module() +delegate_external_constants_pass_unlifted( + tagged_module, gen_tag_fn=lambda x: "model", # This is the filename the weights will be saved to. In this case, weights will be saved as "model.ptd" ) - +# Re-export to get the EP. +exported_program = export(tagged_module, inputs, dynamic_shapes=dynamic_shapes) executorch_program = to_edge_transform_and_lower( exported_program, transform_passes = [partial_function], diff --git a/exir/passes/external_constants_pass.py b/exir/passes/external_constants_pass.py index 414e131d6f5..998844c8ac2 100644 --- a/exir/passes/external_constants_pass.py +++ b/exir/passes/external_constants_pass.py @@ -88,37 +88,6 @@ def external_mutable_weights_pass( return PassResult(gm, mutated) -def delegate_external_constants_pass( - gm: GraphModule, - ep: ExportedProgram, - gen_tag_fn: Optional[Callable[[torch.fx.Node], str]] = None, -) -> PassResult: - """ - Tag external constants before to_backend. - - Note: this pass must be run after run_decompositions(), as tags on - constants are removed then. - - Args: - gm: GraphModule to tag. - ep: ExportedProgram, to distinguish if a node is a constant. - gen_tag_fn: node -> str callable indicating the tag for the node. - Returns: - PassResult: The resulting gm, and if it was mutated or not. - """ - mutated = False - for module in gm.modules(): - if not isinstance(module, torch.fx.GraphModule): - continue - for node in module.graph.nodes: - if node.op == "placeholder" and is_param_node(ep, node): - if gen_tag_fn is not None: - node.meta.setdefault("custom", {}) - node.meta["custom"]["delegate_constant_tag"] = gen_tag_fn(node) - mutated = True - return PassResult(gm, mutated) - - # Note: this pass must be run on an unlifted graph, e.g. ep.module(), # and not on a lifted graph, e.g. ep.graph_module. # This is using 'get_attr' to tag constants, which only appears in diff --git a/test/models/export_delegated_program.py b/test/models/export_delegated_program.py index cbfdfaedab3..44f5c89f7cd 100644 --- a/test/models/export_delegated_program.py +++ b/test/models/export_delegated_program.py @@ -28,7 +28,7 @@ ExecutorBackend, ) from executorch.exir.passes.external_constants_pass import ( - delegate_external_constants_pass, + delegate_external_constants_pass_unlifted, ) from executorch.exir.program import ExecutorchProgramManager from torch import nn @@ -172,18 +172,15 @@ def forward(self, *args, **kwargs): from executorch.backends.xnnpack.partition.xnnpack_partitioner import ( XnnpackPartitioner, ) - - transform_passes = [] if external_constants: - partial_function = partial( - delegate_external_constants_pass, - ep=exported_program, + tagged_module = exported_program.module() + delegate_external_constants_pass_unlifted( + gm=tagged_module, gen_tag_fn=lambda x: module_class.__name__, ) - transform_passes.append(partial_function) + exported_program = export(tagged_module, args=inputs, strict=True) executorch_program = to_edge_transform_and_lower( exported_program, - transform_passes=transform_passes, compile_config=edge_config, partitioner=[XnnpackPartitioner()], ).to_executorch(config=et_config) From 34254bd8c1b2abad5b563e67221ae72404e78f60 Mon Sep 17 00:00:00 2001 From: lucylq Date: Wed, 6 Aug 2025 11:26:23 -0700 Subject: [PATCH 2/2] Update on "Use unlifted export pass to tag delegated constants" Use the unlifted pass to tag constants for delegates. Implications: - Tagging must happen on the unlifted ep.module(), before going into to_edge_transform_and_lower/to_edge. Why? - The unlifted graph contains constants in getattr nodes, which is a convenient way to isolate constants. - After going into to_edge_transform_and_lower/to_edge, transforms happen on the graph_module, which is lifted. - The lifted graph requires the ep graph signature to differentiate constants via the `is_param` function. However, in to_edge.transform, we do not have access to the ep. Baking the ep as an argument via partial function doesn't work, as the ep from earlier may be outdated. This means we are comparing an older ep to a newer graph_module, which may not have corresponding graph signatures etc. Differential Revision: [D79736684](https://our.internmc.facebook.com/intern/diff/D79736684/) [ghstack-poisoned]