[Relax][PyTorch] Add extra_buffers support for exported program

[abeeha123]

Fixes #18357 — KeyError in the Relax PyTorch frontend when loading an HuggingFace BERT model exported with torch. export.
The issue occurs because non-persistent buffers, such as position_ids and token_type_ids, are not included in the state_dict, leading to missing-buffer errors.

This PR adds support for extra non-persistent buffers by introducing an extra_buffers lookup and extending buffer-resolution logic to handle shapes and dtypes correctly.

With this fix, HuggingFace BERT models exported via torch. export no longer fails with missing-buffer KeyErrors.

[gemini-code-assist]

Summary of Changes

Hello @abeeha123, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request addresses a critical KeyError encountered when attempting to load HuggingFace BERT models, exported using torch.export, into the Relax PyTorch frontend. The issue stemmed from the omission of non-persistent buffers, like position_ids and token_type_ids, from the model's state_dict. The solution involves implementing explicit support for these 'extra buffers' and refining the buffer-resolution process to correctly identify and incorporate their shape and data type information, thereby enabling seamless model integration.

Highlights

• Extra Buffers Support: Introduced an extra_buffers mechanism to explicitly handle non-persistent buffers, such as position_ids and token_type_ids, which are not typically included in the state_dict during torch.export.
• Enhanced Buffer Resolution: Extended the buffer-resolution logic within create_input_vars to correctly process shapes and data types for both standard state_dict entries and the newly supported extra_buffers.
• KeyError Fix: Resolved a KeyError that occurred when loading HuggingFace BERT models exported with torch.export into the Relax PyTorch frontend, ensuring successful model loading by accounting for previously missing buffers.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request aims to add support for extra non-persistent buffers in the PyTorch frontend. While the intention is good, the current implementation has a couple of issues. First, the extra_buffers are hardcoded, which limits the feature's reusability. It would be better to pass them as a parameter. Second, the logic to handle these extra buffers is incorrect and will lead to runtime errors due to an undefined variable and incorrect attribute access. I've provided suggestions to address these points.

[gemini-code-assist]

This block has a couple of critical issues:

1. The merged_state variable is used but not defined, which will cause a NameError. You should define it before the loop by merging exported_program.state_dict and extra_buffers. For example: merged_state = {**exported_program.state_dict, **extra_buffers}.
2. The logic to get shape and dtype using hasattr(info, "shape") will not work for the dictionaries in extra_buffers, as dictionaries do not have shape and dtype attributes. This will lead to torch_shape and torch_dtype not being defined, causing another NameError.

Here is a suggested implementation that fixes these issues. Please note that you still need to define merged_state before this block.

Suggested change

	#Buffers
	info=None
	if spec.target in merged_state:
	info=merged_state[spec.target]
	elif spec.target.split(".")[-1] in merged_state:
	info = merged_state[spec.target.split(".")[-1]]
	if info is None:
	raise KeyError(f"Missing target in state_dict or extra buffers: {spec.target}")
	# Handle both original and extra buffer
	if hasattr(info,"shape") and hasattr(info,"dtype"):
	torch_shape=info.shape
	torch_dtype=info.dtype
	info = merged_state.get(spec.target)
	if info is None:
	short_name = spec.target.split(".")[-1]
	info = merged_state.get(short_name)
	if info is None:
	raise KeyError(f"Missing target in state_dict or extra buffers: {spec.target}")
	if isinstance(info, torch.Tensor):
	torch_shape = info.shape
	torch_dtype = info.dtype
	elif isinstance(info, dict):
	torch_shape = info["shape"]
	torch_dtype = info["dtype"]
	else:
	raise TypeError(f"Unsupported type for buffer/parameter info: {type(info)}")

[gemini-code-assist]

This extra_buffers dictionary is hardcoded, which limits the utility of this feature to a specific model configuration (BERT with sequence length 128). To make this more general, extra_buffers should be an optional parameter to from_exported_program and passed down to this function. This would allow users to provide their own non-persistent buffers for different models.

[mshr-h]

Make sense to me. Can you also add testcase for that in https://github.com/apache/tvm/blob/main/tests/python/relax/test_frontend_from_exported_program.py ?