Support for Shared Weights between multiple entry points in PTE export

[pkmandke]

Is there (ideally, backend agnostic) support for sharing weights in a multiple entry point PTE exported file? #11738 seems to track this for XNNPack. However, it doesn't seem to be supported based on the last update on the thread (?)

---

On similar lines, alternatively, is it possible to serialize and export just the (FX) graph/s and load the weights separately (say from a flat_tensor)?

A wishful API would be something like this:

ep = export(model, inputs)
ep = _export_forward_backward(ep)

ep = to_edge(ep)
# Optionally ep.to_backend()
ep = ep.to_executorch()

# Now a wishful workflow would look something like:
ep.save_graph_only(path)
ep.save_weights_only(path)

cc @JacobSzwejbka @angelayi

[lucylq]

Hi @pkmandke,

Is there (ideally, backend agnostic) support for sharing weights in a multiple entry point PTE exported file

Yes, weight-sharing across multiple entry points in a PTE is supported for XNNPACK, and I believe is also the default.

The issue you're referencing also tracks sharing memory/state between entrypoints, which doesn't have good support at the moment cc @JacobSzwejbka

is it possible to serialize and export just the (FX) graph/s and load the weights separately (say from a flat_tensor)?

Yes, there is support for saving the program and weights separately. You can see a demo here: https://github.com/pytorch-labs/executorch-examples/tree/main/program-data-separation

This will save your file to flat_tensor (PTD) format. Again this supported for XNNPACK but not other delegates at the moment.

[pkmandke]

Yes, weight-sharing across multiple entry points in a PTE is supported for XNNPACK, and I believe is also the default.

It doesn't seem to work for me. Here is the example I am using along with the latest executorch off main:

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(2048, 128256)
        self.linear2 = nn.Linear(128256, 2048)

    def forward(self, x):
        return self.linear2(F.sigmoid(self.linear(x)))

    def forward_2(self, x):
        return self.linear2(F.relu(self.linear(x)))

net = Net()
x = torch.randn(1, 2048)
ep = export(net, (x,), strict=True)

net.forward = net.forward_2
ep_fwd_2 = export(net, (x,), strict=True)
    
ep = to_edge({"fwd": ep, "fwd_2": ep_fwd_2})
ep = ep.to_backend(XnnpackPartitioner(force_fp32_dynamic_linear=True))
ep = ep.to_executorch(
        config=ExecutorchBackendConfig(
            external_mutable_weights=False
    )
)
with open(output_file, "wb") as fp:
    ep.write_to_file(fp)

With this, the output PTE file is twice the size when exported with the two entry points as compared to just to_edge(ep). Am I missing something?

[pkmandke]

Yes, there is support for saving the program and weights separately. You can see a demo here: https://github.com/pytorch-labs/executorch-examples/tree/main/program-data-separation

Great, thanks @lucylq!

Would this also work for saving and loading part of the (constant) weights from an external PTD? To update only a subset of the weights in the graph. Or perhaps a way to mark them as mutable?

The external_constants flag seems to apply to all the weights..

For instance, a potential workflow could be something like:

ep = export(model, inputs)
ep.to_edge()
ep.to_executorch(list_of_externally_mutable_weights=my_list)

Now, only list_of_externally_mutable_weights will be serialized to a PTD. And the runtime would ideally support loading only this subset via the PTD, while the remaining weights can either already be part of the PTE file or can also be loaded from another PTD.

This flow can be useful for use-cases where LoRA adapters can be switched in LLMs for instance.

[lucylq]

Edit: just saw this task to make weight sharing default #12503

cc @mcr229, was there anything special to enable weight-sharing in XNNPACK?

Thanks for the suggestsions @pkmandke! I'm currently working on the API to enable the LoRA use-cases 😅

[pkmandke]

Thanks @lucylq!

@mcr229 just bumping this up. Any suggestions for enabling weight sharing in XNNPACK?

[lucylq]

Hey @pkmandke, sorry for the late response. I think you can enable it by turning on this flag in the Cmake:

EXECUTORCH_XNNPACK_ENABLE_WEIGHT_CACHE

See:

executorch/backends/xnnpack/CMakeLists.txt

Line 28 in 07b0883

if(EXECUTORCH_XNNPACK_ENABLE_WEIGHT_CACHE)

[lucylq]

Also, if you were still interested in the a lora example/use-case, feel free to take a look at meta-pytorch/executorch-examples#52