Access to Orbax Checkpoints

[xingjian-zhang]

Simply defines checkpoint paths for many pre-trained models (Gemma 2/3, Qwen, etc.) in simply/config_lib.py, but the actual orbax checkpoint files are not provided anywhere.

Is making these checkpoints publicly available on your roadmap? This would be very useful for external contributors who want to use Simply for research and development.

[kartik-ux]

Thanks for sharing the pretrained model configs in config_lib.py super helpful! 🙌
Quick question: are there any plans to publish the matching Orbax checkpoints as well?
Having them available would make it much easier for the community to experiment, reproduce results, and contribute without needing to write our own conversion pipelines.

[crazydonkey200]

Thanks for the suggestions! I definitely want to make the converted checkpoints and datasets available to make it easier to play with Simply but I am still figuring out the best way to host them.

Currently I am experimenting with using huggingface for hosting and just added a script in the setup folder and instructions for setup in README.md. Please let me know if there are any issues or better alternatives. If it works well, I will add more model checkpoints and datasets this way.

[xingjian-zhang]

Thanks for the update! The solution looks and works great. I confirm that I can set up the model checkpoints by following the instructions in README.md and run experiments on GPUs!

[xingjian-zhang]

Hi @crazydonkey200, given that the current solution works well, could you please upload the Qwen3 checkpoints when you have a chance? It would be great to play with some SOTA models in Simply. Thank you!

[crazydonkey200]

Hi Xingjian, Sorry for the wait. I just added more converted Qwen3 checkpoints and the script for converting huggingface checkpoints from safetensors to orbax format. The uploaded checkpoints are generated using this script. We will add more checkpoints and are currently adding support for MoE models. Please take a look and let me know if there are any issues.

[xingjian-zhang]

Thanks a lot, Chen! I have tested the Qwen3-0.6B model on my side, and it works well. The only issue is that the tokenizer was not correctly set up or registered. So I sent a PR to fix this: #4

[crazydonkey200]

Great! Just merged the PR.