`needle finetune` NaNs on the first optimiser step — root cause

Package: cactus-needle 2.0.0 (PyPI) Reported by: THOX.ai — training orchestration Status: root-caused and measured; THOX has worked around it with an independent trainer.

Summary

needle finetune prints a finite loss on step 1 and nan for every step after, at every learning rate. The cause is a 0/0 inside AdamW, produced by the interaction of three individually reasonable choices:

needle2.pkl stores all 56 tensors as float16.
init_lora (needle/model/finetune.py:236) casts the LoRA factors to weight.dtype, so A and B are float16, and B is zero-initialised.
optax.adamw's default eps=1e-8 is below float16's smallest subnormal (5.96e-8) and rounds to 0.0.

Because B starts at zero, the step-1 gradient with respect to A is exactly zero — dL/dA = scale · grad_out @ Bᵀ. Adam's update for A is therefore

m̂ / (√v̂ + ε)  =  0 / (√0 + 0)  =  0/0  =  NaN

on the first optimiser step. A becomes NaN, the merged weights become NaN, and every subsequent loss is NaN. Step 1 reports a finite loss only because loss is computed before the update is applied.

Isolated reproduction

import numpy as np, jax.numpy as jnp, optax

print(np.float16(1e-8))            # -> 0.0

for dt in (jnp.float16, jnp.float32):
    p = {"A": jnp.zeros((4, 4), dt)}
    g = {"A": jnp.zeros((4, 4), dt)}      # zero because B was zero-init
    opt = optax.adamw(1e-4)
    upd, _ = opt.update(g, opt.init(p), p)
    print(dt.__name__, jnp.isfinite(optax.apply_updates(p, upd)["A"]).all())

0.0
float16 False
float32 True

Why a learning-rate sweep cannot find this

A 0/0 is scale-invariant. Reducing the learning rate by 10× changes nothing. Our earlier observation that --lr 1e-5 appeared to fail "sooner" than 1e-4 was an artifact of the progress cadence (every = total_steps // 50), not a difference in behaviour — both fail on the first optimiser step.

This is worth stating because an LR sweep is the natural first response to a NaN, and here it produces a correct negative result ("not the learning rate") while being structurally incapable of identifying the actual cause.

Suggested fixes

Any one of these resolves it; the first is the standard mixed-precision discipline and is what THOX adopted:

Keep LoRA factors and optimiser state in float32, treating the checkpoint's float16 as a storage format only. Note that needle/model/decode.py:250 already does exactly this for inference (_f32(params)); the training path simply omits it.
Pass an epsilon representable in float16, e.g. optax.adamw(lr, eps=1e-4).
Initialise A to zero and B to random instead, so the step-1 gradient with respect to the zero-initialised factor is non-zero. (Weakest option — it removes this instance without removing the underflow.)

Secondary issue — the default checkpoint is unreachable

DEFAULT_BASE = "checkpoints/needle2.pkl" (finetune.py:18) resolves against HF_REPO = "Cactus-Compute/needle-prod" (tokenizer.py:33), which returns 404 for any account outside Cactus:

RepositoryNotFoundError: 404
https://huggingface.co/Cactus-Compute/needle-prod/resolve/main/checkpoints/needle2.pkl

The equivalent file is public at Cactus-Compute/needle2 under weights/needle2.pkl. Passing --checkpoint with a local path works around it, but the out-of-the-box default is unusable outside Cactus. The same repo is the fallback for get_tokenizer, so a user without the tokenizer cached hits it there too.

Third issue — `needle_load()` + `needle_complete()` faults on Windows

Loading any external .cact and then generating crashes the native library:

OSError: exception: access violation reading 0x0000000003110EA0
  needle/__init__.py:76 in complete -> _lib().needle_complete(...)

This is not specific to a rebuilt artifact. The upstream-shipped needle2.cact and a THOX-exported .cact fail identically, at the same call site, with the same fault:

weights	`needle_load`	`needle_init`	`needle_complete`
stock (no `weights=` argument)	n/a	ok	ok — 966 tok/s prefill, 452 tok/s decode, 35.8 MB peak
upstream-shipped `needle2.cact`	ok (returns 0)	ok	access violation
THOX-exported `.cact`	ok (returns 0)	ok	access violation

So the load and init paths accept the artifact, and the library is functional in its stock configuration; only the externally-loaded weights path faults. Reproduced on Windows 11 (10.0.26300) with the cactus-needle 2.0.0 wheel's auto-fetched native library.

Practical consequence: needle build ... --lora produces an artifact that cannot be exercised on Windows via the Python binding, so end-to-end validation of any fine-tune has to happen elsewhere.

Licensing note

The published wheel is internally inconsistent about its own license: METADATA declares Apache-2.0 while licenses/LICENSE contains the MIT License. Worth reconciling, since the two differ on patent grants.

Environment

cactus-needle 2.0.0
jax 0.10.2 · jaxlib 0.10.2 · optax 0.2.8 · flax 0.12.8
Windows-11-10.0.26300 · jax backend: cpu · Python 3.11

needle finetune NaNs on the first optimiser step — root cause