Add support for LFM2 hybrid attention-convolution architecture

[noctrex]

LFM2 models use alternating convolution and attention layers with different attribute names than standard transformer models. This change adds fallback paths to handle LFM2's structure:

• Support self_attn.out_proj for attention layers (vs o_proj)
• Support conv.out_proj for convolution layers
• Support feed_forward.w2 for MLP (vs mlp.down_proj)

All attribute accesses are now wrapped in exception suppression to gracefully handle different model architectures. This enables heretic to work with LFM2's hybrid architecture without skipping any layers.

[gemini-code-assist]

Summary of Changes

Hello @noctrex, 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 enhances the heretic library's compatibility with diverse neural network architectures, specifically enabling it to process LFM2 hybrid attention-convolution models. By adapting to LFM2's unique naming conventions for attention, convolution, and MLP layer attributes and implementing robust error handling, the changes ensure that heretic can correctly identify and utilize components from these models without encountering unexpected failures.

Highlights

• LFM2 Attention Layer Support: Introduced support for LFM2's specific attention layer output projection (self_attn.out_proj) as an alternative to the standard o_proj.
• LFM2 Convolution Layer Support: Added support for LFM2's convolution layer output projection (conv.out_proj).
• LFM2 MLP Support: Implemented support for LFM2's MLP down-projection (feed_forward.w2) instead of the typical mlp.down_proj.
• Robust Attribute Access: Wrapped all attribute access attempts in exception suppression blocks to gracefully handle variations in model architectures, preventing crashes when expected attributes are not found.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

• src/heretic/model.py
  • Refactored existing try_add function calls for identifying model components to use contextlib.suppress(Exception).
  • Added new try_add calls to specifically target LFM2's self_attn.out_proj for attention layers.
  • Included a try_add call for LFM2's conv.out_proj to support convolution layers.
  • Implemented a try_add call for LFM2's feed_forward.w2 to support MLP layers.
  • Enclosed the original try_add calls for attn.o_proj and mlp.down_proj within with suppress(Exception) blocks to ensure graceful handling of models that do not use these specific attribute names.

Activity

• No specific activity (comments, reviews, etc.) has been provided in the context for this pull request.

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 adds support for the LFM2 hybrid attention-convolution architecture by adding fallback paths to find the correct modules for attention, convolution, and MLP layers. The changes correctly identify alternative attribute names in LFM2 models. My review includes a suggestion to make the implementation more robust by ensuring the fallback paths are mutually exclusive, which will prevent potential issues with future model architectures.

[p-e-w]

Thanks for the PR! Note that there is already a PR (#43) that implements support for hybrid models. Perhaps you want to check that PR for verification and potentially relevant ideas.

What results are you getting from abliterating LFM2? Have you uploaded the model to Hugging Face?

[noctrex]

Thanks for the PR! Note that there is already a PR (#43) that implements support for hybrid models. Perhaps you want to check that PR for verification and potentially relevant ideas.

What results are you getting from abliterating LFM2? Have you uploaded the model to Hugging Face?

Oh sorry must have missed it. I'll try that one to make some experiment on LFM2, cause mine does not work very well.