[ONNX] Support constant, buffer, parameters in ExportedProgram serialization to ONNX tensor

[titaiwangms]

torch.onnx.dynamo_export currently supports ExportedProgram with the help of input adapter to prepend the lifted constant, buffer, and parameters:

pytorch/torch/onnx/_internal/io_adapter.py

Line 583 in 796dff7

class PrependParamsBuffersConstantAotAutogradInputStep(InputAdaptStep):

So the inputs need to be adapted before the onnx model is executed:

pytorch/torch/onnx/_internal/exporter.py

Lines 743 to 750 in 796dff7

	onnx_input = self.adapt_torch_inputs_to_onnx(
	*args, model_with_state_dict=model_with_state_dict, **kwargs
	)
	options = options or ONNXRuntimeOptions()
	providers = (
	options.execution_providers or onnxruntime.get_available_providers()
	)
	ort_session = onnxruntime.InferenceSession(onnx_model, providers=providers)

However, this temporary workaround should be fixed with actual serialization of the weights and buffers. A potential fix can be serializing the weights and buffers when we save the model to ONNX:

pytorch/torch/onnx/_internal/exporter.py

Lines 1009 to 1032 in 796dff7

	def save(
	self,
	destination: Union[str, io.BufferedIOBase],
	*,
	include_initializers: bool = True,
	model_state: Optional[Union[Dict[str, Any], str]] = None,
	serializer: Optional[ONNXProgramSerializer] = None,
	) -> None:
	"""Saves the in-memory ONNX model to ``destination`` using specified ``serializer``.
	Args:
	destination: The destination to save the ONNX model. It can be either a string or a file-like object.
	When used with ``model_state``, it must be a string with a full path to the destination.
	If `destination` is a string, besides saving the ONNX model into a file, model weights are also stored
	in separate files in the same directory as the ONNX model. E.g. for `destination="/path/model.onnx"`,
	the initializers are saved in "/path/" folder along with "onnx.model".
	include_initializers: Whether to include initializers in the ONNX graph as external data.
	Cannot be combined with `model_state_dict`.
	model_state: The state_dict of the PyTorch model containing all weights on it.
	It can be either a string with the path to a checkpoint or a dictionary with the actual model state.
	The supported file formats are the same as those supported by `torch.load` and `safetensors.safe_open`.
	Required when :func:`enable_fake_mode` is used but real initializers are needed on the ONNX graph.
	serializer: The serializer to use. If not specified, the model will be serialized as Protobuf.
	"""

Related #126071
Related #126479

[titaiwangms]

cc @justinchuby @shubhambhokare1

[titaiwangms]

Potentially, user might want us to unlift parameters/weights when the model is small.