[Dynamo] Enable dynamo.export for huggingface models w/ ModelOutput

[BowenBao]

🚀 The feature, motivation and pitch

Initially reported at #96386, dynamo.export raises for any huggingface model that returns a subclass of ModelOutput.

Here is an experimental patch I have done to unblock dynamo.export based ONNX export. The idea is to extend pytree to be able to flatten/unflatten ModelOutput and its subclasses.

pytorch/torch/onnx/_internal/fx/exporter.py

Lines 21 to 68 in 95de585

	@contextlib.contextmanager
	def _patch_pytree_huggingface_modeloutput():
	"""Patch 'torch.utils._pytree' to support family of 'ModelOutput' from HuggingFace 'transformers'.
	The source and details of the issue is described at https://github.com/pytorch/pytorch/issues/96386.
	This patch enables 'torch.utils._pytree' to flatten and unflatten all 'ModelOutput'
	subclasses defined in HuggingFace 'transformers'. Hence resolving the mismatch between
	`dynamo` eager traced model outputs and `dynamo.export` fx graph computed outputs.
	FIXME(bowbao): Remove this patch after above issue is resolved for `dynamo.export`.
	"""
	try:
	from transformers import modeling_outputs # type: ignore[import]
	except ImportError as e:
	# Do nothing if 'transformers' is not installed.
	try:
	yield
	finally:
	pass
	return
	def model_output_flatten(
	output: modeling_outputs.ModelOutput,
	) -> Tuple[List[Any], _pytree.Context]:
	return list(output.values()), (type(output), list(output.keys()))
	def model_output_unflatten(
	values: List[Any], context: _pytree.Context
	) -> modeling_outputs.ModelOutput:
	output_type, keys = context
	return output_type(**dict(zip(keys, values)))
	# All 'ModelOutput' subclasses are defined under module 'modeling_outputs'.
	named_model_output_classes = inspect.getmembers(
	modeling_outputs,
	lambda x: inspect.isclass(x) and issubclass(x, modeling_outputs.ModelOutput),
	)
	for _, class_type in named_model_output_classes:
	_pytree._register_pytree_node(
	class_type, model_output_flatten, model_output_unflatten
	)
	try:
	yield
	finally:
	for _, class_type in named_model_output_classes:
	_pytree.SUPPORTED_NODES.pop(class_type)

I wonder if it is a good idea to bring this native into pytree? The context of this patch needs to cover not only dynamo.export but any other subsequent fx pass that may interact with the formatted output, otherwise it won't recognize the output structure.

Feedbacks and suggestions are welcomed. Feel free to let me know if there are other more suitable solutions.

@jansel @voznesenskym @wconstab

Alternatives

No response

Additional context

No response

cc @zou3519 @ezyang @msaroufim @wconstab @bdhirsh @anijain2305 @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @aakhundov @kadeng @ngimel @Xia-Weiwen @soumith @desertfire

[jansel]

If I recall you can convert a ModelOutput to a tuple(), so a more concise workaround might be something like:

torch.export(lambda *args: orig_model(*args).to_tuple(), ...)

I feel most export targets will have trouble constructing a python class anyway, so getting a basic type in the export might be more useful.

[wconstab]

@BowenBao how do you feel about just using the workaround proposed by @jansel? Is it crucial to have support for exporting a ModelOutput class? Specifically, do you want the 'ModelOutput' class to show up in the runtime where you load the exported module? or you just want a way to export this model?

[BowenBao]

@wconstab thanks for following up. I'm unblocked with workaround combining the idea from the post and suggestion from @jansel.

For your question, I feel it is a question on the design/goal of dynamo.export, which I may have a few misunderstandings from the start. Is there a hidden requirement/design that the graph module produced by dynamo.export should have identical inputs & outputs interface with the original pytorch function? Because that seems like what the code is doing, and becomes the constraint that lead to the failures in the original post.

For me this issue was opened based on a simple context: dynamo.export cannot succeed on various models from Huggingface, the root cause turns out to be ModelOutput. These are basic models so I thought they should work directly.

[wconstab]

Is there a hidden requirement/design that the graph module produced by dynamo.export should have identical inputs & outputs interface with the original pytorch function?

I think it's a good question for follow up. Tagging some others @voznesenskym @SherlockNoMad @suo

[ezyang]

It is a difficult question to answer, because it depends on the intended use case for the exported model.

For example, let's suppose you are exporting a model just to have a "compiled" callable in Python that you can reload later and then execute. Then it would make sense to have identical inputs and outputs, because you have some sort fancy pickler for functions.

But let's say instead your exporting the model to interoperate with some C++ runtime predictor service. Then obviously the Python structs don't matter, you just need to know enough about the output Tensors to reconstruct them in whatever schema you need.

[tugsbayasgalan]

Is this still an issue?