Help with Converting NAMO-R1 to MLX

[yourappleintelligence]

I tested the NAMO-R1 model, and the results are amazing! Now, I’m wondering what kind of performance boost and efficiency it could achieve with MLX on Apple Silicon.

The model is in .safetensors format and includes:

• namo.llm (a language model)
• namo.ve (possibly a vision encoder)

I’d love to convert it to MLX, but I’m not sure where to start. I’ve tried extracting the weights and exploring MLX documentation, but I still need guidance on:

• Way for converting safetensors to MLX-compatible format
• Defining the model architecture in mlx.nn
• Mapping and loading the weights correctly

Has anyone here worked on similar model conversions? Any guides, scripts, or insights would be greatly appreciated!

Can’t wait to see what magic NAMO-R1 + MLX can do.

[davidkoski]

Way for converting safetensors to MLX-compatible format

MLX (both python and swift) can load safetensors files -- this is the preferred format. In some cases the MLX code may rename tensors or even apply transforms. You can see an example of renaming here:

• https://github.com/ml-explore/mlx-swift-examples/blob/main/Libraries/MLXVLM/Models/Qwen2VL.swift#L896

Defining the model architecture in mlx.nn

You can look at how the models we have here are built. It is relatively straightforward to port mlx.nn (python) code to MLXNN (swift) code. I have not seen how people port e.g. PyTorch to MLX so I can't comment there, but there are a lot more examples in the python code -- I would probably start there.

Mapping and loading the weights correctly

See the first item -- once you define the modules (layers) loading the weights is pretty straightforward.