B=1 FUSED linear cross-entropy response-only mean log-prob → [1]. Same result
as branchLogpMeanB1 but the LM head is a single CustomVjp with an ANALYTIC
backward — the Liger FusedLinearCrossEntropy structure ported to MLX:
FORWARD token-chunks the M response rows; per chunk materializes only
[Cc, V] logits (= h_c @ Wᵀ via quantizedMatmul, + Gemma softcap),
reduces to logp = target_logit − logsumexp, and DROPS the logits. Returns
Σ logp (the 1/M mean is applied outside, a cheap linear scale).
BACKWARD recomputes [Cc, V] per chunk and forms the gradient in CLOSED
FORM — ∂(Σlogp)/∂logit = onehot_target − softmax (× sech² through the
softcap), scaled by the incoming cotangent — instead of autograd-ing through
logsumexp+gather (what the Checkpoint token-chunked path does). dh_c = grad_logits_c @ dequant(W) uses quantizedMatmul(…, transpose=false), which
IS mlx's own x-vjp of the transpose=true head matmul (the path LoRA
training already differentiates through), so it contracts the vocab axis
correctly without dequantizing the full [V, hidden] head (1.3 GB at 262k).
Peak large-vocab memory is bounded to [chunkSize, V] in BOTH directions, the
full [M, V] logits never materialize, and the head backward needs no autograd
graph. The CustomVjp recomputes in the backward (it must outlive the
enclosing value_and_grad), so it — and h — are pushed into sink, disposed
by the caller after the step's grads are eval'd. The head is NOT a LoRA target,
so only dh flows back (seeding the layer-stack LoRA backward); a head-LoRA
∇A/∇B fold-in is a follow-on. Exact within the established bf16 class.
B=1 FUSED linear cross-entropy response-only mean log-prob → [1]. Same result as
branchLogpMeanB1but the LM head is a singleCustomVjpwith an ANALYTIC backward — the LigerFusedLinearCrossEntropystructure ported to MLX:[Cc, V]logits (=h_c @ WᵀviaquantizedMatmul, + Gemma softcap), reduces tologp = target_logit − logsumexp, and DROPS the logits. Returns Σ logp (the1/Mmean is applied outside, a cheap linear scale).[Cc, V]per chunk and forms the gradient in CLOSED FORM —∂(Σlogp)/∂logit = onehot_target − softmax(×sech²through the softcap), scaled by the incoming cotangent — instead of autograd-ing through logsumexp+gather (what theCheckpointtoken-chunked path does).dh_c = grad_logits_c @ dequant(W)usesquantizedMatmul(…, transpose=false), which IS mlx's own x-vjp of thetranspose=truehead matmul (the path LoRA training already differentiates through), so it contracts the vocab axis correctly without dequantizing the full[V, hidden]head (1.3 GB at 262k). Peak large-vocab memory is bounded to[chunkSize, V]in BOTH directions, the full[M, V]logits never materialize, and the head backward needs no autograd graph. TheCustomVjprecomputes in the backward (it must outlive the enclosing value_and_grad), so it — andh— are pushed intosink, disposed by the caller after the step's grads are eval'd. The head is NOT a LoRA target, so onlydhflows back (seeding the layer-stack LoRA backward); a head-LoRA∇A/∇Bfold-in is a follow-on. Exact within the established bf16 class.