RFC: native MLX training backend for `needle finetune` on Apple Silicon

[nandoolle]

Summary

I have a working MLX port of the training-path forward pass plus LoRA finetuning (needle finetune --backend mlx), and would like to know if there is interest in upstreaming it before I invest in rebasing it onto the current architecture.

Motivation

The current macOS GPU path depends on jax-metal, which pins jax==0.4.38, needs ENABLE_PJRT_COMPATIBILITY, and forces flash/remat/scan off. A native MLX backend removes the plugin and the pin, and runs on the current jax for everything else. During my usage of ee221ce, which my fork was based on, the metal extra also produced NaN losses in the first steps but trained cleanly after changing to MLX.

What was built

563 lines, 6 files, fork branch feat/mlx-finetune-backend.

• needle/model/architecture_mlx.py:
  • MLX forward of exactly what finetuning exercises (SimpleAttentionNetwork, quant=False). Flat param dict mirroring the flax tree
  • mx.checkpoint per MHC layer, mirroring nn.remat. dtype=bfloat16 follows the flax semantics (dense/attention in bf16
  • norms, MHC mixing, sinkhorn, logits in f32). Confidence/contrastive heads, MTP and KV-window decode are out of scope.
• needle/model/finetune_mlx.py:
  • same data format, LoRA placement (five stacked attention kernels), AdamW + warmup/cosine, global-norm clip 1.0, masked CE, val holdout. Step compiled with mx.compile. Writes the same adapter .pkl, so needle build --lora works unchanged.
• CLI:
  • --backend {jax,mlx} and --dtype {float32,bfloat16} on finetune (jax remains the default).
• pyproject.toml:
  • train-mlx extra (mlx>=0.29, numpy, sentencepiece).
• tests/test_mlx_parity.py:
  • f32 forward parity vs the JAX model on the published checkpoint, argmax-identical, max-abs logit diff ~3e-3 (skips without jax+mlx+checkpoint).

Known gap

The port targets the needle2 architecture at ee221ce, and since that main has added sliding_window, global_layers, ladder_*, qkv_conv_taps, qk_head_dim, safetensors and depth rungs, none of which the MLX forward implements yet. Rebasing is a rewrite of the forward, hence this RFC first.

Questions

1. Is a second forward implementation acceptable maintenance-wise, given the parity test pins it to the JAX reference?
2. Would you prefer this as a replacement for the metal extra or alongside it?
3. Is the architecture on main stable enough to port against now?

Happy to open a PR against main if there is interest.

[nimakalantar]

Can you provide a link to your branch? Would be excited to test this

[nandoolle]

Sure, there you go.

But be ware that this implementation is a very quickly throw together PoC, and as of now it's far from production ready.

[HenryNdubuaku]

thanks @nandoolle could you PR this? thanks a million!

[HenryNdubuaku]

Yes, interested. I tried the jax-metal route first today and it isn't good enough: it needs jax pinned at 0.4.38, the model only compiles with rematerialisation off, and a training step came out at 1.24 s on the GPU against 0.50 s on the CPU with current jax, with the loss barely moving over ten steps. So the plugin is not a path we can build on.

That leaves a native MLX backend, which is what you're proposing. What it needs to cover: attention with the conv taps, the Hadamard MLP, the engram lookup, the four-lane residual with its Sinkhorn mixing, ladder sampling, and Muon plus Lion. The thing I'd insist on is a parity check against the JAX forward on the same checkpoint, so results stay comparable to everything we've measured.

If you want to take it on, start with the forward pass and the parity test and open a draft; I'll review it. Say roughly how much you want to bite off and I'll stay out of the same files.

[nandoolle]

Hi, @HenryNdubuaku, sorry I've took so long to respond your comment.

This week was been a rush at work but you can be sure I'll be opening the PR as soon as possible.

Thanks for the attention.

[nandoolle]

If you want to take it on, start with the forward pass and the parity test and open a draft; I'll review it. Say roughly how much you want to bite off and I'll stay out of the same files.

@HenryNdubuaku, draft's open at #149