Move all mutable weights to external file.
(
gm: GraphModule,
ep: ExportedProgram,
)
| 65 | |
| 66 | |
| 67 | def external_mutable_weights_pass( |
| 68 | gm: GraphModule, |
| 69 | ep: ExportedProgram, |
| 70 | ) -> PassResult: |
| 71 | """ |
| 72 | Move all mutable weights to external file. |
| 73 | """ |
| 74 | # pass the gm and the ep seperately as the gm is being mutated by a bunch of passes in to_executorch, |
| 75 | # so the gm in the ep is lagging the graph signature is still correct. |
| 76 | # This is really tech debt and all the passes should be refactored to just mutate the ep. |
| 77 | mutated = False |
| 78 | for module in gm.modules(): |
| 79 | if not isinstance(module, torch.fx.GraphModule): |
| 80 | continue |
| 81 | |
| 82 | for node in module.graph.nodes: |
| 83 | if node.op == "placeholder": |
| 84 | spec = node.meta.get("spec") |
| 85 | if ( |
| 86 | isinstance(spec, TensorSpec) |
| 87 | and spec.const |
| 88 | and _is_mutable_weight(node, ep) |
| 89 | ): |
| 90 | node.meta["constant_tag"] = "_default_external_constant" |
| 91 | mutated = True |
| 92 | return PassResult(gm, mutated) |
| 93 | |
| 94 | |
| 95 | # Note: this pass must be run on an unlifted graph, e.g. ep.module(), |
no test coverage detected