Adding Gemma 3n support?

[dwgb93]

⚠️ Please check that this feature request hasn't been suggested before.

• I searched previous Ideas in Discussions didn't find any similar feature requests.
• I searched previous Issues didn't find any similar feature requests.

🔖 Feature description

I saw you guys were mentioned in the Gemma 3n Developer's Guide and wanted to check if you were tracking to add support for the family of MatFormer models recently released by Google.

✔️ Solution

I tried using the gemma3 examples as a starting point, and after upgrading transformers to v4.53.0 and installing timm, I still got opaque errors like

) is not supported. Currently, only the following modules are supported: torch.nn.Linear, torch.nn.Embedding, torch.nn.Conv1d, torch.nn.Conv2d, torch.nn.Conv3d, transformers.pytorch_utils.Conv1D, torch.nn.MultiheadAttention..

I'd love it if, like most other models I've finetuned with axolotl, it just worked.

❓ Alternatives

Unsloth says they have Gemma 3n support, but their page just says "Gemma 3n Notebooks coming soon! Stay tuned."

📝 Additional Context

No response

Acknowledgements

• My issue title is concise, descriptive, and in title casing.
• I have searched the existing issues to make sure this feature has not been requested yet.
• I have provided enough information for the maintainers to understand and evaluate this request.

[NanoCode012]

Thanks. I got a simple config running. However, the loss is very unusually high. Will need to look into this further.

Your issue can be fixed with:

-lora_target_linear: true

and manually target the layers you want

lora_target_modules:
  - q_proj
  - k_proj

[dwgb93]

@NanoCode012 thanks for your help!
I'm still getting stuck though.

I was getting an error about gemma3n not being supported. I upgraded cut-cross-entropy, per your pull request, and tried again using the example config file you made. This gave met the longest error message I've ever seen.

It printed out the entire gemma3n architecture, before ending:

) got multiple values for keyword argument 'return_dict'

Any idea what that means?

[NanoCode012]

Hey @dwgb93 , I would say to wait a bit for us to experiment further with this. There's quite a few things we need to fix first.

) got multiple values for keyword argument 'return_dict'

Do you have more logs for this?

If you want to experiment, the linked PR has a sample config.

[NanoCode012]

I'm not sure how but I triggered this error too with minimal changes.The fix can be found here for now till I figure out if anything else is wrong axolotl-ai-cloud/ml-cross-entropy#5

[dwgb93]

Nice! I'll mess around with what you've provided, but I'll let you guys cook in the meantime.

Thanks for your help!