Interested in contributing to [MODEL-1] and [MODEL-2]

[Mayura Andrew]

Hi Jan,

I hope you’re doing well.

I am reaching out because I would love to start contributing to the project. I was looking through the issue tracker and am very interested in tackling [MODEL-1] HuggingFace Transformer Import and eventually [MODEL-2] Text Generation Example.

As a software engineering student who is passionate about open-source and LLMs, my main objective right now is to go beyond just using API wrappers and truly understand how these models are built and loaded under the hood.

To be completely transparent, I do not have prior knowledge of the lower-level mechanics of ML (like direct safetensors imports without ONNX) just yet. However, I am highly motivated and extremely willing to put in the time to learn. Because I am starting fresh in this specific area, I am really looking for your guidance. If I hit a roadblock, would you be open to pointing me toward the right documentation or answering a few questions to help me get unstuck?

Looking forward to hearing from you and hopefully contributing to the project!

Best regards,

Mayura Alahakoon

https://mayura-andrew.github.io/

[Jan Pfeifer]

(Copying answer send in private in the list, for the record)

hi Mayura, thanks for reaching out, and apologies for the delay in responding (I'm always online in the Slack channel for a more prompt response btw).

Here is a status update on the topics of your interest:

HuggingFace Transformer Import

Coincidentally, I've been working on the HuggingFace Transformer Import the last couple of weeks because I needed to run the KaLM-Gemma3 12B model importing.

The code to read from HuggingFace is entirely within the go-huggingface project. I created the KaLM-Gemma3 example here. It uses the `models/transformer` library in go-huggingface to load the configurations and weights, then calls the actual implementation in GoMLX's `pkg/ml/model/transformer`.

Fundamentally, it lays a framework for implementation. Currently, it implements only one of the transformer "architectures", and supports the set of configurations I found when implementing my model. It now needs support for more models—both to increase support and to check whether current abstractions suffice or require refactoring (likely not, as this was a first attempt).

Text Generation

There is a first attempt at it in GoMLX `pkg/ml/decode`, but it has a couple of serious issues (and maybe others):

- Bucketing: GoMLX still only supports fixed shapes and requires "JIT-recompilation" for different shapes, making bucketing essential for performance. I attempted to address this in go-huggingface's `tokenizers/bucket`. 

- KVCache: It is currently stored as variables in the context. But that doesn't work if more than one inference runs simultaneously, so it's the wrong API. I've been postponing this because I also want to implement a paged KVCache approach, that make more efficient use of memory, but that is more complicated.

Dynamic Shapes Support

The upcoming priority (for this quarter) for GoMLX is to add support for dynamic shapes. That includes: (1) add support for dynamic shapes in the Go backend (`simplego`); (2) Refactor GoMLX /backends package into its own repo (so another project can use it independent of the rest of GoMLX); (3) Add ONNXRuntime as a backend (currently we are able to convert ONNX to GoMLX), now we want to be able to execute a GoMLX computation using ONNX. And maybe also add llama.cpp as a backend as well. Just a heads-up this is coming.

---

Sorry, this was a long email. I appreciate the enthusiasm, working with ML/AI is super fun, full of surprises and one learns a lot. My recommendation is to have a model or project that uses GoMLX and that you want to see through. First, just work on GoMLX to fill missing functionality you need. In other words, start with small changes and get used to developing with or in GoMLX before attempting anything large.

Find me on Slack if you have something in mind and want to chat at some point.

cheers!