fix(train): NF4 QLoRA CUDA forward computed the WRONG MODEL — NeoX rope pairing + Q/K/V biases + full-warp softmax + causal CPU oracle

[noahgift]

Summary

Cascade defect #4 (after #2249 deadlock, #2250 loss window, #2251 stream race): training ran NaN-free but never learned — CE flat at 13-14 (> ln(V)=11.93, worse than uniform) on trivially predictable targets, then adapter poisoning after ~125 steps. Three-oracle bisection (GPU-fused vs GPU-forward+CPU-CE vs pure-CPU) localized it to the transformer forward and exonerated the fused CE kernel — then per-op layer-0 bisection found four stacked defects:

1. Wrong RoPE pairing (dominant) — batched_rope_neox_forward/_backward instantiated BatchedRopeKernel (GPT-J adjacent pairs, elem0=2*pair_idx) despite the NeoX name. Qwen2 weights need NeoX split-half (i, i+d/2). Fix: new BatchedRopeNeoxKernel/BatchedRopeNeoxBackwardKernel (precise trig per CORRECTNESS-013); wrappers + pre-warm keys rewired. BatchedRopeKernel untouched — realizar uses it for non-NeoX rope types.
2. Dropped Q/K/V biases — CudaNf4TransformerBlock had zero bias support; Qwen2 sets use_bias=true and the model carries blk.N.attn_{q,k,v}.bias (CPU path applies them). Fix: replicated bias buffers + cuda_add_inplace after each projection GEMM, threaded from all 3 NF4 sites + the instruct FP32 site. Backward unchanged (activation-checkpoint re-runs the fixed forward; biases frozen).
3. Partial-warp shfl.sync UB — batched softmax fwd/bwd launched block=(32.min(row_size)) while reducing with membermask 0xFFFFFFFF → undefined for seq<32 → data-dependent NaN rows. Fix: always full 32-lane warp.
4. Non-causal CPU oracle (label leakage) — autograd/ops/attention.rs had no causal mask; the CPU train/eval path attended bidirectionally (toy causal CE 2.13 misreported as 0.17). Fix: attention_causal dispatched for ModelArchitecture::Decoder (encoders stay bidirectional); shared backward exact since masked weights are 0.

Verification (RTX 4090, sm_89)

• Falsifier FALSIFY-CUDA-NF4-TRAIN-LOSS-PARITY-001 (parity_probe.rs): GPU loss + full logits vs an NF4-quantization-matched causal CPU oracle. GREEN: CE 0.6820 vs 0.6557 (|Δ|=0.026 < 0.5), logits relL2 0.047 < 0.10. Mutation-verified RED per individually reverted fix: rope → relL2 0.183; biases → CE 5.17 / relL2 0.97; warp → NaN. Companion probes: per-layer bisect + per-op layer-0.
• Contract contracts/cuda-nf4-train-loss-parity-v1.yaml — pv validate + pv lint contracts/ PASS.
• E2E (apr_code_sft_balanced.jsonl, 160 samples, --max-seq-len 2048): first-step CE 1.58 (pre-fix: 13-14); at lr 2e-5 loss descends 4.31→0.29 by step 50, epoch avg 1.55, 0 NaN — independently reproduced on this cherry-picked content: epoch 1 avg_loss=1.5842, 0 NaN, run still converging.
• aprender-gpu --lib 440/440; aprender-train --lib failures (6) all reproduce on the base commit (pre-existing: release-mode debug_assert class, insta snapshot drift).

Known follow-ups (separate beats, not claimed here)

• Auto-config lr 2e-4 @ rank 256 diverges after ~25 steps (finite) — training-dynamics defect; lr 2e-5 proven stable/decreasing.
• evaluate() val_loss is byte-identical across runs with different lr (pre-existing) — likely no GPU→CPU LoRA sync before eval; needs its own bisection.
• rms_norm_backward hardcodes eps=1e-5 (backward-only precision vs Qwen2 1e-6).
• Single-row SoftmaxKernel broken by construction for all lengths (not on the training path).

With this, apr finetune -m qlora trains the correct model on the 4090 end-to-end — the apr-code tool_call flip is running on it now.

🤖 Generated with Claude Code