Run pretrained LLMs with nnx-lm

[jaco-bro]

Hi Flax team,

I’d like to share nnx-lm, a Python package for running LLMs on any hardware backends using NNX.

Motivation:

The currently available Flax examples don’t show how to load pretrained weights from HuggingFace safetensors or evaluate/compare runtime performances in a configurable way across various hardware backends and optimization options.

Features:

• Native safetensors support to load weights from any model repo (no conversion).
• Configurable optimizations through flags to toggle jit and scan and batching.
• XLA-specific optimizations e.g., avoiding dynamic tensor shapes in KV caching.

Goal:

Just as how projects like mlx-lm have gained traction by making pretrained models easy to run and extend, nnx-lm aims to offer a similar experience for Flax users but without introducing yet another model zoo.

I'd love to hear your thoughts. Thank you!