{"slug": "putting-my-jax-trained-models-on-the-hugging-face-hub", "title": "Putting my JAX-trained models on the Hugging Face Hub", "summary": "Developer gpjt has uploaded PyTorch-compatible versions of all his JAX-trained GPT-2 models to the Hugging Face Hub, including models from his blog series on writing an LLM from scratch and Chinchilla scaling experiments. The uploads enable use with Transformers' AutoModelForCausalLM, which has been PyTorch-only since version 5.", "body_md": "I hadn't uploaded the models that I trained using JAX to the Hugging Face Hub because\nTransformers has been PyTorch-only since [version 5](https://huggingface.co/blog/transformers-v5)\n(though they say they're working to add interoperability with JAX in the future), so it\nwould have been tough to get them working natively with `AutoModelForCausalLM`\n\nand the like.\n\nBut then it dawned on me that I'd already written [a conversion script](https://github.com/gpjt/jax-gpt2-from-scratch/blob/main/convert_model_to_pytorch.py) that could\ntake my JAX safetensors files and convert them into ones compatible with my PyTorch\ncode. It's actually those converted models that I use for my evals!\n\nSo, I've now uploaded PyTorch-compatible versions of all of my JAX-trained models:\n\n\"[Writing an LLM from scratch, part 34b -- from bigrams to GPT-2, one component at a time (in JAX)](/2026/07/llm-from-scratch-34b-building-and-training-gpt-2-small-in-jax)\"\n\n`gpjt/jax-no-mha-bias-no-dropout`\n\n`gpjt/jax-no-mha-bias-with-dropout`\n\n`gpjt/jax-with-mha-bias-no-dropout`\n\n\"[Why do OpenAI's GPT-2 weights beat mine? Part three: testing overtraining](/2026/07/why-do-openai-gpt2-weights-beat-mine-3-overtraining)\"\n\n`gpjt/jax-with-mha-bias-no-dropout-extended`\n\n`gpjt/jax-with-mha-bias-no-dropout-2-epoch`\n\n\"[A quick(ish) Chinchilla check](/2026/08/chinchilla-check)\"\n\n`gpjt/jax-with-mha-bias-larger-chinchilla-1`\n\n`slightly-larger`\n\nmodel.`gpjt/jax-with-mha-bias-larger-chinchilla-2`\n\n`slightly-smaller`\n\nmodel.I've also added links to the posts in question.", "url": "https://wpnews.pro/news/putting-my-jax-trained-models-on-the-hugging-face-hub", "canonical_source": "https://www.gilesthomas.com/2026/09/jax-models-on-hugging-face", "published_at": "2026-09-03 17:25:17.019712+00:00", "updated_at": "2026-09-03 17:25:18.820862+00:00", "lang": "en", "topics": ["artificial-intelligence", "machine-learning", "large-language-models", "ai-products", "developer-tools"], "entities": ["Hugging Face Hub", "Transformers", "JAX", "PyTorch", "GPT-2", "gpjt"], "alternates": {"html": "https://wpnews.pro/news/putting-my-jax-trained-models-on-the-hugging-face-hub", "markdown": "https://wpnews.pro/news/putting-my-jax-trained-models-on-the-hugging-face-hub.md", "text": "https://wpnews.pro/news/putting-my-jax-trained-models-on-the-hugging-face-hub.txt", "jsonld": "https://wpnews.pro/news/putting-my-jax-trained-models-on-the-hugging-face-hub.jsonld"}}