[Frontend][PyTorch] torch2.0 ExportedProgram PoC

[chunit-quic]

• Add converter for torch2.0 ExportedProgram

[chunit-quic]

Hi community,

This is a PoC to support ExportedProgram, which is a new format in pytorch 2.0. May we know do you have any plan to support ExportedProgram? Or maybe this PR look fine to be a start point? :D

ExportedProgram is a new format in  pytorch2.0. After exporting to ExportedProgram, a model represented by Module or functional will be flatten and represented by core aten op. One of its benefits is that, we can focus on writing converting function for those core aten ops. It can decrease the effort to support different Module and functional.

We are wondering do you have any plan to support this format as well? if there is no such plan, perhaps this PR can be a simple reference to support ExportedProgram conversion. Or if it looks good to you, we can submit a formal one based on your comments, and add this to converter. It will be really nice to have your thoughts. Thanks!

ref:
https://pytorch.org/executorch/stable/tutorials/export-to-executorch-tutorial.html
https://pytorch.org/docs/stable/export.html#exporting-a-pytorch-model
https://pytorch.org/docs/stable/torch.compiler_ir.html

Best,
Joey

[chunit-quic]

Hi @Hzfengsy and @jikechao,

Pardon that disturbing you. I saw you review Relax related PR. May you kindly give us some suggestions about this PR please? Thank you very much!

[yongwww]

Thank you for the proposal!

With the introduction of the SLM, we now are able to utilize the TVM nn.module for supporting models created with Torch, which has been working well afaik. Not sure if the SLM works for the cases you are working on. Here is the llama model implemented via SLM (the SLM llama model).

Looks the Torch FX is the underlying representation of the exported graph, I am wondering if it is possible to use/update the existing FX translator (relax/frontend/torch/fx_translator.py) to support the ExportedProgram.

[chunit-quic]

Hi @yongwww

Thanks for your prompt reply! Here is our thoghts for your reference. :D

Not sure if the SLM works

Pardon that we didn't investigated SLM before. We can let some teamates to read on it later.

I am wondering if it is possible to use/update the existing FX translator

Please allow me to describe more detial about this PR, and what problems we may encounter if we want to integrate it into FX translator.

• Purpose
  This PR aims to support pytorch2.0 new format ExportedProgram, which will be executed by executorch.
  Like tflite relay frontend, which convert tflite model to relay, where tflite model is supposed to executed by tflite-interpreter.
• Problems to intergrate with FX translator

1. Params is stored in different place in fx.GraphModule and ExportedProgram.
   fx.GraphModule: model.named_parameters()
   ExportedProgram: exported_program.graph_signature.inputs_to_parameters or exported_program.graph_signature.inputs_to_buffers
2. Some function has different API, and it will cause ambiguous. e.g., convolution
   fx.GraphModule: conv2d, whose inputs are (input, weight, bias, stride, padding, dilation, group)
   ExportedProgram: aten.convolution, whose inputs are (input, weight, bias, stride, padding, dilation, transposed, output_padding, group)
3. The source to retrieve node attribute is different
   fx.GraphModule: node.kwargs
   ExportedProgram: node.kwargs is empty, we need to get it from node.args

Thank you for reading this long reply. It will be really nice to have your advice!

Best,
Joey

[yongwww]

@chunit-quic thanks for the reply! I'm new to the Torch ExportedProgram, I noticed that the export API is still under development (pytorch/executorch#290). If ExportedProgram is the preferred way to obtain the GraphModule compared to fx.symbolic_trace, we could consider utilizing the export API to get the ExportedProgram and updating the fx translator to work with the GraphModule in the ExportedProgram.

For insights into SLM, please refer to this documentation and explore the code in relax/frontend/nn. If, after further evaluation, SLM aligns with your requirements, we could consider shifting our focus directly to it.

cc: @MasterJH5574 @junrushao @LeshengJin @cyx-6

[chunit-quic]

Hi @yongwww,

Thank you for providing the SLM materials. :)

If ExportedProgram is the preferred way to obtain the GraphModule...

Yes, we think it has more benefit to use ExportedProgram, since every functional/nn.module can be exported to ExportedProgram and their functional op/Module will be decomposed to core aten ops which only has ~200 ops. It means that we only need to write translator for these ~200 ops and then it can support all pytorch model conversion.

However, it is hard to reach the sky in a single bound. If we update existing fx translator, we need to make a whole change, and it may break existing support models. Hopefully, this new translator should have the same capacity of existing translator. If so, we can focus on ExportedProgram rather than fx translator.

For insights into SLM, please refer to this...

Thanks for your introduction, this is awesome feature to debug modularized relay module. However, our issues are more related to unsupported op conversion in the fx translator, and this feature still requires us to write translators for the new functional/nn.Module.

That is the reason we want to leverage the ExportedProgram because it will decompose all functional/nn.Module into limitation set of ops.

[tqchen]

Thanks @chunit-quic for the contribution! I think there are several goals:

• G0: Support translation of PT ops, regardless if they appear in normal FX trace for ExportProgram
• G1: Support ExportProgram usedcase
• G2: Support integration of torch dynamo via FX graph

I think we all agree that G1 is super important. I would like us to think about how we can unify, so we don't have two versions of the torch importer.

Here is one possible way to do so:

• We build a common base class that contains the translation of ops (that can be used in both fx and ExportedProgram)
• The base class allows abstract method for program specific behaviors
  • self.get_parameters(mod) will obtain the parameters
    • For FXGraph we get from model.named_parameters()
    • For ExportedGraph we change according
  • self.get_node_attrs(node) will obtain the attributes
• Support both op variants if needed
  • Because op keys are different, we can support aten.conv2d and conv2d, either as a subclass patch, or support both in base, it does not hurt since both can be useful at different levels, especially for key ops like conv2d

[tqchen]

@chunit-quic let me know if it can help address the problems you see, love to working together and get this feature in!

[chunit-quic]

Hi @tqchen,

Thank you for joining the discussion!
Yes we agree those features are useful at different levels.
We will refactor PR based on your suggestion, so as to integrate fx_translator for both GraphModule and ExportedProgram.

Just one more thing. Would you like us to submit a RFC, new PR or keep updaing this one? Thank you.

[tqchen]

I think update this PR would be sufficient

[tqchen]

@chunit-quic just want to check in and see where we are on this

[chenweng-quic]

Hi @tqchen,

Thank you so much for your attention and effort on this pull request. We truly appreciate your suggestion. Unfortunately, after careful consideration, we've decided to withdraw the pull request at this time. We hope to collaborate again in the future and appreciate your understanding.