# `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:

1. `needle2.pkl` stores **all 56 tensors as float16**.
2. `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.
3. `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

```python
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:

1. **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.
2. Pass an epsilon representable in float16, e.g. `optax.adamw(lr, eps=1e-4)`.
3. 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
```