Add decomposition for dynamo_export + ExportedProgram and remove None from input

[thiagocrepaldi]

Stack from ghstack (oldest at bottom):

• Extend _TestONNXRuntime to reuses all tests for new model format #112289
• Add support for models with mutated buffer on torch.onnx.dynamo_export #112272
• Add torch.onnx.dynamo_export test using ExportedProgram from file #112271
• Rename torch.onnx.ExportOutput* to ONNXProgram* #112263
• -> Add decomposition for dynamo_export + ExportedProgram and remove None from input #112444

This PR introduces the ability to produce GraphModules with Core ATen IR only through decompositions. It also removes None from user inputs as ONNX does not supports them

Tests for these features will be executed when #112289 is merged, but for reference, they are as below:

    def test_log_sigmoid(self):
        # This produces op as `torch.ops.aten.log_sigmoid_forward`, instead of the more
        # conventional `torch.ops.aten.log_sigmoid`.
        class Model(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.m = torch.nn.LogSigmoid()

            def forward(self, x):
                return self.m(x)

        input = torch.randn(2)
        self.run_test_with_fx_to_onnx_exporter_and_onnx_runtime(
            Model(), (input,), model_type=self.model_type
        )

    def test_none_input(self):
        class NoneInputModel(torch.nn.Module):
            def forward(
                self, x: torch.Tensor, y: Optional[torch.Tensor], z: torch.Tensor
            ):
                if y is None:
                    return x + z
                return x + y + z

        self.run_test_with_fx_to_onnx_exporter_and_onnx_runtime(
            NoneInputModel(),
            (torch.randn(1, 2), None, torch.randn(1, 2)),
            model_type=self.model_type,
        )

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/112444

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit e160a05 with merge base c1e2ccd ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[BowenBao]

Put it in pre_export_passes?

[thiagocrepaldi]

This is tricky. Although we have implemented decomposition as a GraphModule pass for the torch._dynamo.export fx_tracer, torch.export does decomp at the ExportedProgram level. Through the torch.export + ExportedProgram flow, we run decompositions before we start doing passes of the GraphModule (using self.pre_export_passes)

In fact, at generate_fx we call return self.pre_export_passes() passing original_model=model and fx_module=model.graph_module. Note that we extract GraphModule from the ExportedProgram.graph_module after the decomposition was ran. Only then we start doing GraphModule passes (and ignore original_model).

In order to move the ExportedProgram decomposition to pre_export_passes we have to

1. change model = model.run_decompositions(options.decomposition_table) call (which assigns a ExportedProgram to another ExportedProgram) to fx_module = original_model.run_decompositions(options.decomposition_table).graph_module(which assigns aGraphModulecoming from aExportedProgram; that is the generate_fx's role, not pre_export_passes`)
2. Always run fx_module: GraphModule = exported_program.run_decompositions as the first pass and
3. Remove fx_module as the input to the pre_export_passes, meaning generate_fx is not generating the GraphModule anymore

Although techinically possible, it mixes responsibilities of generate_fx and pre_export_passes.

[BowenBao]

Okay. We are missing out diagnostics for decomp though.

[thiagocrepaldi]

how should we address this?

[BowenBao]

Need to think more on that. Good opportunity to explore how to track low level info such as these from other modules.

[thiagocrepaldi]

@pytorchbot merge

[pytorchmergebot]

Merge failed

Reason: PR 112444 is out of sync with the corresponding revision e22e881 on branch origin/gh/thiagocrepaldi/9/orig that would be merged into main. This usually happens because there is a non ghstack change in the PR. Please sync them and try again (ex. make the changes on origin/gh/thiagocrepaldi/9/orig and run ghstack).

Details for Dev Infra team

Raised by workflow job

[thiagocrepaldi]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here