Fix/save torchao model loading logic

[rolandtannous]

Problem

1- AutoModel: We previously used AutoModel from pretrained to load models in our implementation of the save_pretrained_torchao method. However AutoModel loads the base model without a specific task-head.

2-Merged model directory: save_pretrained_torchao merges the model (if necessary) before converting and does not delete the merged model directory after the process is done which consumes unnecessary disk space and potentially confuses users a

Solution

1-AutoModel:

• We test if the model is a vlm or a text model
• We set auto_model to either AutoModelForCausalLM or AutoModelForImageTextToText based on the test result
• We also set auto_processor to either AutoProcessor or AutoTokenizer based on the above test
  This loads the model properly

is_vlm = False
    if hasattr(self, "config") and hasattr(self.config, "architectures"):
        is_vlm = any(
            x.endswith(("ForConditionalGeneration", "ForVisionText2Text"))
            for x in self.config.architectures
        )
        is_vlm = is_vlm or hasattr(self.config, "vision_config")
    auto_model = AutoModelForImageTextToText if is_vlm else AutoModelForCausalLM
    auto_processor = AutoProcessor if is_vlm else AutoTokenizer

    tokenizer = auto_processor.from_pretrained(arguments["save_directory"])
...
 model = auto_model.from_pretrained(
        arguments["save_directory"],
        device_map = "auto",
        quantization_config = quantization_config,
        **kwargs,
    )

2 - Merged model directory: We remove the save_directory directory tree after the conversion process is complete. The torchao converted model is saved to the torchao_save_directory instead.

if os.path.exists(save_directory):
        try:
            import shutil
            shutil.rmtree(save_directory)
        except:
            pass
            

solves

#3599

[gemini-code-assist]

Summary of Changes

Hello @rolandtannous, 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 refines the save_pretrained_torchao functionality by enhancing model loading to correctly identify and load Vision-Language Models (VLMs) or text models with their appropriate task-specific heads. Additionally, it introduces a cleanup mechanism to automatically remove temporary merged model directories after conversion, optimizing disk space usage. A minor update to the model mapper is also included.

Highlights

• Improved torchao model loading logic: The save_pretrained_torchao method now intelligently selects between AutoModelForCausalLM and AutoModelForImageTextToText (and their respective processors) based on whether the model is a Vision-Language Model (VLM) or a text model. This ensures models are loaded with the correct task-head, addressing an issue where AutoModel loaded only the base model.
• Clean up of merged model directories: After the torchao conversion process, the temporary merged model directory is now automatically removed. This prevents unnecessary disk space consumption and potential user confusion by ensuring only the final torchao converted model directory remains.
• Added BF16 model to mapper: The unsloth/gpt-oss-20b-unsloth-bnb-4bit entry in the model mapper has been updated to include unsloth/gpt-oss-20b-BF16, expanding the recognized model aliases.

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 correctly addresses two issues with the save_pretrained_torchao method. First, it now uses the appropriate AutoModel class (AutoModelForCausalLM or AutoModelForImageTextToText) based on whether the model is a Vision Language Model, ensuring that the correct model head is loaded. Second, it adds a cleanup step to remove the temporary merged model directory, which helps conserve disk space. The logic for detecting VLMs and the overall changes look good. I've left one comment regarding error handling in the cleanup step to make it more robust.