[ONNX] ONNX doesn't support exporting non-persistent buffer included models in FakeMode

[titaiwangms]

To avoid out of memory issue during exporting models to ONNX, we need to detach the parameters and persistent buffers with state_dict().

  # Create the toy model with real weight.
  real_model = create_model()

  with tempfile.NamedTemporaryFile(
      prefix=model_name, suffix=".pt"
  ) as tmp_checkpoint_file:
      # Dump state_dict to a file to simulate how HuggingFace model is initialized.
      # The file will be loaded via .load_state_dict(...)
      state_dict = real_model.state_dict()
      torch.save(state_dict, tmp_checkpoint_file.name)

      with torch.onnx.enable_fake_mode() as fake_context:
          fake_args = create_args()
          fake_kwargs = create_kwargs()
          fake_model = create_model()
          if load_checkpoint_during_init:
              fake_model.load_state_dict(torch.load(tmp_checkpoint_file.name))

          # Export the model with fake inputs and parameters
          export_options = torch.onnx.ExportOptions(
              dynamic_shapes=self.dynamic_shapes,
              op_level_debug=self.op_level_debug,
              fake_context=fake_context,
          )

          export_output = torch.onnx.dynamo_export(
              fake_model,
              *fake_args,
              **fake_kwargs,
              export_options=export_options,
          )

However, some models, for example, GPT2, there is non-persistent buffer which can't be detached to state_dict(). Subsequently, ONNX graph complains about the missing buffers, but we don't have it in external data of the model initializer. This kind of case can be be reproduced when we use Config to create_model().

cc @BowenBao @thiagocrepaldi @wschin

[titaiwangms]

Brainstorming on this. Alternative way might be we random generate it, as it's not crucial.

[thiagocrepaldi]

@titaiwangms IIRC this should work with the torch.export.export approach, as the non-persistent buffers would be passed in as input, right?

[titaiwangms]

Based on your PR #115380, it does look like this issue is gone with ExportedProgram. I think we can track it with your issue #115745