{"slug": "eqx-zoo-hub-models-in-jax-equinox-verified-against-transformers-and-what-bf16-me", "title": "Eqx-zoo: Hub models in JAX/Equinox, verified against transformers and what verifying bf16 taught me", "summary": "Eqx-zoo, a library by xquantize, loads Hugging Face Hub checkpoints directly as plain Equinox modules in JAX and verifies each model against transformers, supporting Llama, Qwen2, Qwen3 and Qwen3-MoE for generation with KV cache and batching plus BERT, RoBERTa and XLM-RoBERTa embedding models. The author reports that JIT-compiling a whole decoder in XLA made bf16 measurably less accurate than eager execution, up to 2.4x transformers' own bf16 error, while encoders showed no such effect, and that transformers' own bf16 drifted from its float32 after 2 of 20 tokens on Llama 3.2 1B. The library is installable via pip install eqx-zoo and has so far been tested only on CPU, with the author seeking feedback from JAX users on GPU or TPU.", "body_md": "With transformers v5 focusing on PyTorch, I’ve been building eqx-zoo, which loads Hub checkpoints directly (safetensors, by repo name, no conversion) as plain Equinox modules and verifies each model against transformers itself.\n\n**What’s supported:** Llama, Qwen2, Qwen3 and Qwen3-MoE for generation (KV cache, batching) and BERT, RoBERTa and XLM-RoBERTa embedding models, with each checkpoint’s sentence-transformers pooling and normalisation read from its config. The verified checkpoints are in this collection: [Verified in eqx-zoo - a xquantize Collection](https://huggingface.co/collections/xquantize/verified-in-eqx-zoo-6ac2d731fd1ef3e109541bf7)\n\n**How verification works**\n\n- float32: every layer’s output is compared with transformers’ (captured with forward hooks) and greedy generation must reproduce transformers’ output token for token\n- embeddings must match sentence-transformers on padded batches\n- tiny randomly initialised models cover code paths no single checkpoint exercises, with every parameter randomised (default init sets norm weights to 1 and biases to 0, which can hide a dropped scale or bias)\n\n**What bfloat16 taught me**\n\n- Two bf16 implementations round differently, so comparing ours with transformers’ bf16 directly is the wrong test: they’re often further apart than either is from float32. Instead, ours is compared with float32, relative to transformers’ own bf16 error.\n- Exact greedy agreement in bf16 isn’t a sound criterion: transformers’ own bf16 drifted from its float32 after 2 of 20 tokens on Llama 3.2 1B.\n- RMS error is a much more stable statistic than max error, which is dominated by single unlucky roundings.\n- JIT-compiling a whole decoder in XLA made bf16 measurably less accurate than eager execution (up to 2.4x transformers’ own bf16 error), and compiling a single layer alone reproduced it. Encoders, which are post-norm, didn’t show this at all.\n- A subtle bug (LayerNorm statistics in bf16) shifted model-level error less than the difference between x86 and ARM did, so model-level tests couldn’t catch it reliably. At the block level, the same bug was off by thousands of bf16 steps, so it’s now tested there directly.\n\nQuick example:\n\n``` python\nfrom eqx_zoo import CausalLM, generate\n\nmodel = CausalLM.from_pretrained(\"Qwen/Qwen3-0.6B\")\ntokens = generate(model, prompt_ids, max_new_tokens=30)\n```\n\n`pip install eqx-zoo` — [GitHub - xquantize/eqx-zoo: Verified Equinox ports of pretrained models, numerically matched against Hugging Face · GitHub](https://github.com/xquantize/eqx-zoo)\n\nEverything is tested on CPU so far. I’d love feedback from anyone running JAX on GPU or TPU and I’m curious: which Hub models would be most useful to have in JAX next?", "url": "https://wpnews.pro/news/eqx-zoo-hub-models-in-jax-equinox-verified-against-transformers-and-what-bf16-me", "canonical_source": "https://discuss.huggingface.co/t/eqx-zoo-hub-models-in-jax-equinox-verified-against-transformers-and-what-verifying-bf16-taught-me/182905#post_1", "published_at": "2026-10-04 22:54:04+00:00", "updated_at": "2026-10-04 23:11:06.762086+00:00", "lang": "en", "topics": ["machine-learning", "large-language-models", "ai-research", "developer-tools", "ai-infrastructure"], "entities": ["Eqx-zoo", "xquantize", "Equinox", "JAX", "Hugging Face", "transformers", "Llama 3.2 1B", "Qwen3-0.6B"], "also_reported_by": [], "alternates": {"html": "https://wpnews.pro/news/eqx-zoo-hub-models-in-jax-equinox-verified-against-transformers-and-what-bf16-me", "markdown": "https://wpnews.pro/news/eqx-zoo-hub-models-in-jax-equinox-verified-against-transformers-and-what-bf16-me.md", "text": "https://wpnews.pro/news/eqx-zoo-hub-models-in-jax-equinox-verified-against-transformers-and-what-bf16-me.txt", "jsonld": "https://wpnews.pro/news/eqx-zoo-hub-models-in-jax-equinox-verified-against-transformers-and-what-bf16-me.jsonld"}}