Upload UPSTREAM_ISSUE.md with huggingface_hub

UPSTREAM_ISSUE.md

@@ -1,102 +1,134 @@
-# `needle finetune` produces NaN loss from step 2 — lower LR fails sooner
+# `needle finetune` NaNs on the first optimiser step — root cause
 
 **Package:** `cactus-needle` 2.0.0 (PyPI)
-**Reporter:** THOX.ai — training orchestration
-**Severity:** blocks all LoRA fine-tuning on this host
+**Reported by:** THOX.ai — training orchestration
+**Status:** root-caused and measured; THOX has worked around it with an
+independent trainer.
 
 ## Summary
 
-`needle finetune` computes a finite loss on step 1 and `nan` from step 2–4, on a
-dataset that the package's own loader and renderer accept without complaint.
-**Reducing the learning rate 10× makes it fail *earlier*, not later**, which
-argues against ordinary LR instability and points at the loss/merge path.
+`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:
 
-## Reproduction
+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`.
 
-```bash
-pip install cactus-needle==2.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
 
-# public base weights — needle finetune otherwise tries
-# Cactus-Compute/needle-prod/checkpoints/needle2.pkl, which 404s (private)
-huggingface-cli download Cactus-Compute/needle2 weights/needle2.pkl --local-dir ./nb
-
-needle finetune corpus.jsonl \
-  --checkpoint ./nb/weights/needle2.pkl \
-  --epochs 3 --lora-rank 16 --lora-alpha 32 --max-len 512 --lr 1e-4 \
-  --out ./out/lora.pkl
+```
+m̂ / (√v̂ + ε)  =  0 / (√0 + 0)  =  0/0  =  NaN
 ```
 
-### Observed
+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.
 
-```
-training on 216 examples, seq_len 512
-LoRA rank 16 on 10 weight groups (compiling...)
-epoch 1/3  step 1/42  loss 2.9569
-epoch 1/3  step 2/42  loss nan
-epoch 1/3  step 3/42  loss nan
-epoch 1/3  step 4/42  loss nan
-```
+## Isolated reproduction
 
-At `--lr 1e-5 --max-len 1024`:
+```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())
+```
 
 ```
-epoch 1/3  step 1/42  loss 2.5293
-epoch 1/3  step 2/42  loss nan
+0.0
+float16 False
+float32 True
 ```
 
-### Expected
+## Why a learning-rate sweep cannot find this
 
-Loss decreases; no NaN.
+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.
 
-## What we ruled out first
+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.
 
-| hypothesis | how it was checked | result |
-|---|---|---|
-| Learning rate too high | ran 1e-4 and 1e-5 | **rejected** — 1e-5 NaNs *sooner* (step 2 vs 4) |
-| Sequence truncation zeroing the loss mask | tokenised all rows through `render_example` + the package tokenizer | **rejected** — max total **371 tokens** (prompt 321–329, target 17–43) vs `--max-len` 512. No row's prompt reaches the limit, so no target is truncated away and no mask is all-zero |
-| Malformed rows being skipped | counted rows surviving `load_jsonl`'s `if "query" not in example: continue` | **rejected** — 216/216 accepted |
-| Wrong row schema | rendered rows through `render_example` and inspected prompt/target | **rejected** — renders correctly |
+## Suggested fixes
 
-Remaining candidate: numerical instability in `merge_lora` /
-`softmax_cross_entropy_with_integer_labels` under this jax/jaxlib build, or in
-the `1/rank`-scaled LoRA `A` init interacting with the checkpoint dtype.
+Any one of these resolves it; the first is the standard mixed-precision
+discipline and is what THOX adopted:
 
-## Dataset shape
+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.)
 
-216 rows, `{system, tools, query, reasoning, answers}`. 120 rows carry a single
-tool call; **96 carry `answers: []`** (no applicable tool — an intentional refusal
-class). Rendered target for a refusal row is:
+## 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:
 
 ```
-<think>No provided tool satisfies this request.</think>
-<tool_call>[]</tool_call><|im_end|>
+RepositoryNotFoundError: 404
+https://huggingface.co/Cactus-Compute/needle-prod/resolve/main/checkpoints/needle2.pkl
 ```
 
-If an empty `answers` list is not a supported target, that would be a plausible
-cause and is worth confirming — but note step 1 produces a finite loss over a
-batch that already contains such rows.
+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.
 
-## Environment
+## Third issue — `needle_load()` + `needle_complete()` faults on Windows
+
+Loading *any* external `.cact` and then generating crashes the native library:
 
 ```
-cactus-needle 2.0.0
-jax 0.10.2 · jaxlib 0.10.2 · optax 0.2.8 · flax 0.12.8
-Windows-10-10.0.26300-SP0 · jax backend: cpu
-Python 3.11
+OSError: exception: access violation reading 0x0000000003110EA0
+  needle/__init__.py:76 in complete -> _lib().needle_complete(...)
 ```
 
-## Secondary issue — gated default checkpoint
+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.
 
-`DEFAULT_BASE = "checkpoints/needle2.pkl"` resolves against
-`HF_REPO = "Cactus-Compute/needle-prod"` (`needle/model/tokenizer.py:33`), which
-returns **404** for an unauthenticated/ordinary token:
+## Environment
 
 ```
-RepositoryNotFoundError: 404
-https://huggingface.co/Cactus-Compute/needle-prod/resolve/main/checkpoints/needle2.pkl
+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
 ```
-
-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. Pointing `HF_REPO` at
-the public repo would fix `needle finetune` and `needle run` for everyone else.