fix(aprender-train): CUDA RMSNorm honours config.rms_norm_eps (Qwen 1e-6 vs hardcoded Llama 1e-5)

[noahgift]

Summary

Cascade follow-up to PR #1604 (PMAT-CODE-PRETRAIN-CUDA-FORWARD-001 H4D bias fix).

After landing the bias fix, val_loss on populated Qwen2.5-Coder-0.5B moved 18.55 → 17.22 but stayed above ln(vocab) ≈ 11.93 — the model was producing sub-uniform predictions. Bisection of the residual cascade picked layer-0 RMSNorm as the next stage to interrogate. Source inspection found a real bug:

• aprender-train::rms_norm_forward constructed BatchedVectorizedRmsNormKernel::new(hidden_size, batch_size) without .with_epsilon(eps).
• trueno-gpu's BatchedVectorizedRmsNormKernel::new hardcodes epsilon: 1e-5 (Llama default).
• Qwen2 / Qwen2.5 rms_norm_eps = 1e-6 (per HF config + TransformerConfig::qwen2_0_5b()).
• The CPU RMSNorm::new(hidden_size, eps) honours config; CUDA silently substitutes 1e-5.

Five-Whys

1. Why val_loss > ln(vocab) post-bias-fix? CUDA forward still drifts from CPU on populated weights.
2. Why? Some per-layer numerical operation produces different results.
3. Why at RMSNorm? It's the very first stage of layer-0; check it.
4. Why might it differ? Kernel-construction site missing .with_epsilon(eps).
5. Why does this matter? Qwen-magnitude post-embedding activations have mean_sq ~ 4e-4. Eps gap 9e-6 → ~2.25% relative drift in rsqrt denominator, every call, 49 RMSNorms per forward.

Implementation

• NEW: rms_norm_forward_with_eps(.., eps: f32, ..) (eps-aware variant) constructs the kernel via .with_epsilon(eps) and includes eps_bits in the PTX cache key (different eps → different cached module).
• Legacy rms_norm_forward becomes a thin wrapper using 1e-5 (Llama default) for backwards compatibility.
• All 4 production callsites switched: pre-attn + post-attn norms in CudaTransformerBlock::forward, inference-path pre-attn norm in CudaNf4TransformerBlock, and the final RMSNorm before lm_head in CudaTransformerTrainer::eval_batch.

Provable Contract

New: contracts/apr-pretrain-cuda-rmsnorm-eps-parity-v1.yaml (ACTIVE_ALGORITHM_LEVEL).

Three ship-blocking falsifiers:

• FALSIFY-CUDA-RMSNORM-EPS-PARITY-001: pointwise CPU↔CUDA parity within 1e-4 abs at Qwen eps=1e-6.
• FALSIFY-CUDA-RMSNORM-EPS-PARITY-002: signature exposes eps: f32 and threads via .with_epsilon(eps).
• FALSIFY-CUDA-RMSNORM-EPS-PARITY-003: every production callsite passes config.rms_norm_eps.

Test Plan

• Falsifier falsify_cuda_rmsnorm_eps_parity_qwen_1e_minus_6 GREEN locally (lambda-vector RTX 4090) — max abs diff well within 1e-4 at eps=1e-6 on Qwen-magnitude 4×896 synthetic batch
• pv validate contracts/apr-pretrain-cuda-rmsnorm-eps-parity-v1.yaml — 0 errors, 0 warnings
• cargo fmt --all -- --check — clean
• cargo check -p aprender-train --features cuda — 0 errors
• cargo test -p aprender-train --features cuda --lib --release — exit code 0 (modulo known workspace-test trueno SIGSEGV-on-cleanup flake + 2 pre-existing should_panic mismatches in autograd::ops::matmul::tests unrelated to this change)
• LIVE val_loss recheck on populated Qwen 0.5B + 5g.1-v2 corpus (deferred to a follow-up evidence run after this PR + fix(aprender-train): CUDA forward path applies Q/K/V biases (H4D root-cause discharge) #1604 both merge)

Stacking

This PR builds on fix/cuda-forward-parity-qwen-biases (PR #1604). After #1604 merges to main, this branch will rebase clean.

Ship-% Movement

If MERGED + #1604 merged: SHIP-TWO-001 MODEL-2 still 57% pending the residual cascade (uniform → converged), but two of the cascade contributors are discharged (Q/K/V bias dispatch + RMSNorm eps). Next stage: RoPE, attention softmax, or FFN dispatch — picked via element-wise bisection of layer-0 stages.

🤖 Generated with Claude Code