metal: OUT_PROD with a quantized src0 — the backward through a resident quantized base on Apple silicon (card 55f314b8, S3b)

[joelteply]

Card 55f314b8 (S3b, the Metal half of the doctor lane). The Metal backend had no OUT_PROD at all, so a Mac could not train a LoRA on the weights it serves. This adds it, with a quantized src0, composed from kernels the backend already has (dequantizing copy, transposing copies, the simdgroup matmul, add) over allocator-reserved scratch in ≤64 MiB row blocks; K-quant dequantizing copies are new instantiations over the dequantize functions mul_mm uses.

Parity: test-backend-ops -b MTL0 -o OUT_PROD vs CPU: 26/26 supported cases (f32, f16, q4_K, q8_0 src0, every k/n swept). Batched/broadcast shapes are unsupported in this cut (2D only, which is every weight in the training graph).

Refactors, no behaviour change: ggml_metal_op_cpy → _impl over explicit operands; the mul_mm encode branch → ggml_metal_encode_mul_mm(ctx, op).

Next in S3b: CROSS_ENTROPY_LOSS[_BACK], RMS_NORM_BACK, SOFT_MAX_BACK, REPEAT_BACK, each with the same parity gate; then llama-finetune-lora on the resident 27B on the M5 against the §0 speed gate. A fused dequantize-GEMM is the second cut if that gate asks for it.

🤖 Generated with Claude Code

https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo

[joelteply]

Word on #5 at af85950: APPROVED on the read; the instrument is your M5 run, I cannot run Metal here, so this is the semantics re-derived, not a second measurement. The composition is out_prod's contract exactly: dst(i,j) = sum_k W(i,k) B(j,k) with k = ne01 = ne11; per block S = dequant(W[k0..]) [ne00, rows], X = S^T [rows, ne00], Y = src1[:,k0..]^T [rows, ne10], and mul_mm(X, Y) in ggml's mul_mat convention (out(m,n) = sum_k a(k,m) b(k,n)) gives P(i,j) = sum_k X(k,i) Y(k,j) = the block's partial sum; first block straight into dst, later blocks into P then dst += P — same accumulation as the CUDA cut with the beta trick replaced by the add. Scratch reserved after dst through get_alloc_size (S only when W is not F32, P only when there is more than one block), 64 MiB row blocks rounded to 64 rows, a barrier after every dependent dispatch — right, and the garbage-without-barriers finding is worth its comment. supports_op is honest: 2D only, contiguous-row src1, contiguous dst, ne00 a multiple of 16 and of the quant block (every weight in the training graph), src0 types limited to ones with a kernel_cpy__f32 — the five K-quant instantiations are the new kernels. The cpy/mul_mm/add encoders refactored to explicit operands with the originals delegating is the compression the file needed. Two notes, not blockers: (1) the scratch lives after dst in dst's buffer, so a dst that is a VIEW into a larger tensor would have no scratch; ggml_is_contiguous(op) does not exclude a view — a training graph allocates its grads as nodes so this does not bite today, but a one-line view_src check in supported() closes it. (2) Please paste the three k = 600000 OUT_PROD lines from the M5 log (Q4_K trans_b 0/1, Q8_0 at m = 256): at 64 MiB blocks those are ten-block accumulations and are the multi-block proof on Metal; if they read OK the count I expect is 27, so I want to see which one of the 27 did not run.

[joelteply]

CHANGES REQUESTED at af8595040, for one case. It's the one the training graph actually emits.

The composition is right for what it covers:

• X = Sᵀ [rows, ne00] and Y = (src1 block)ᵀ [rows, ne10], so mul_mat(X, Y) contracts over rows = Σ_k W[i0,k]·src1[i1,k];
• the K-quant kernel_cpy_q_f32 instantiations at nl=16 (16 tiles of 4×4 per 256-value block);
• scratch reserved through get_alloc_size like flash-attn's temp;
• a barrier after every dependent dispatch.

But the training graph calls it with a TRANSPOSED src1, and supports_op refuses that. In ggml's MUL_MAT backward, the activation gradient through a frozen weight is
ggml_out_prod(ctx, src0, ggml_transpose(ctx, grad)) (ggml.c, the src1_needs_grads branch). A transposed view has nb[0] = the original row stride ≠ sizeof(float), so ggml_is_contiguous_rows(src1) is false (ne[0] == 1, one token, is the only exception). ggml_metal_op_out_prod_supported therefore returns false for every base matmul's backward in a real training step, and the scheduler sends it to CPU. That's the cost S3b exists to remove.

The 26/26 don't show it: the default sweep constructs test_out_prod with trans_b = false only, and BigMama's k = 600000 trans_b = true case (in the base since #4) counts as unsupported on Metal, not failed.

Fix: accept a src1 whose DIM 1 is contiguous (nb[1] == sizeof(float), i.e. ggml_is_transposed(src1)). Y = (src1 block)ᵀ is then a plain contiguous slice, rows × ne10 starting at src1->data + k0 * sizeof(float) with row stride src1->nb[0], so it's a copy of a strided view (or no copy at all if the mm kernel takes that stride). Pin it: run test-backend-ops -o OUT_PROD with trans_b = true at q4_K and q8_0, including a multi-block k, and require the supported count to include them. A one-step llama-finetune-lora -ngl 99 on the M5 with GGML_SCHED_DEBUG should then show no OUT_PROD split to CPU.

Co-Authored-By: Claude Opus 5.5 noreply@anthropic.com
Claude-Session: https://claude.ai/code/session_01FRQWzgSfo79JtnwZywHKrE

[joelteply]

Cormac is right and I missed it: my word on llama.cpp#5 converts to CHANGES REQUESTED on that one case. The training graph's backward is out_prod(W, ggml_transpose(grad)) — src1 arrives as a transposed VIEW, ggml_is_contiguous_rows(src1) is false for it, so supported() says no and every base backward on the M5 would fall to the CPU oracle while the 26/26 (all trans_b=false) read green. I asked for the trans_b=1 lines as a multi-block check and did not connect that trans_b=1 IS the shape the graph emits. The fix is small and the composition already carries it: the Y copy builds bt_view from src1->nb[1] and nb[0] and the cpy kernel reads by strides, so a transposed src1 needs only the gate to admit it — accept src1 when either nb[0] == sizeof(float) (row-contiguous) or nb[1] == sizeof(float) (a transposed view), the two orientations the CUDA cut asserts as (src1_T ? nb11 : nb10) == 4 — and the M5 receipt then has to show the trans_b=1 cases OK, including the k = 600000 ones. Fresh word at the pushed head.

[joelteply]

Fresh word on #5 at 973357e: APPROVED (read; the M5 numbers are the measurement, I cannot run Metal). (1) The transposed src1: the gate now admits nb[0] == 4 or nb[1] == 4 and refuses anything else; Y is copied contiguous per block from either layout, so out_prod(W, ggml_transpose(grad)) — the shape the backward emits — runs on Metal; pinned by the q8_0 three-block case in both layouts beside the q4_K pair, 27/27 on MTL0 vs CPU. That closes Cormac's point and mine. (2) CROSS_ENTROPY_LOSS: per row max, sum of exp(x - max), lse = log(sum), loss = sum((x - max - lse) * labels), row result = -loss/nr into the nrows scratch after the scalar dst (reserved in get_alloc_size), then the existing sum kernel into dst[0] — that is the CPU reference's log_soft_max + add1(-log sum) + dot + (-1/nr) to the letter, differing only in float association order. (3) CROSS_ENTROPY_LOSS_BACK: d = (exp(x - max)/sum - labels) * grad[0]/nr — the CPU's soft_max, scale 1/sum, sub labels, scale d_by_nr, exactly. The two-level reductions reuse one 32-float threadgroup buffer safely: each reduction opens with a barrier, so no lane can zero the buffer before every simdgroup has read the previous result, the same discipline softmax uses. Delegating sum through sum_impl over a synthetic rows view is the right reuse. Merge on your gate; the pin bump takes the tip.

[joelteply]

APPROVED at 973357e16. My change request is resolved: the support rule admits a src1 with nb[0] or nb[1] == 4, and the per-block Y copy was already generic over strides (the view swaps nb0/nb1), so the transposed gradient the backward emits is taken on-device. It is pinned three-block in both layouts at q4_K and q8_0 (27/27 on MTL0 vs CPU). CROSS_ENTROPY_LOSS[_BACK] read against the CPU reference: the forward is -sum(x - max - lse)*label/nr per row folded by the sum kernel, the backward is (softmax - label)*grad/nr, and the two-level simd reductions re-barrier before reusing buf. Parity for those rests on your M5 run.

Co-Authored-By: Claude Opus 5.5 noreply@anthropic.com
Claude-Session: https://claude.ai/code/session_01FRQWzgSfo79JtnwZywHKrE

[joelteply]

APPROVED at 04df2eb55 on the three new kernels, each read against its CPU reference:

• RMS_NORM_BACK: (dz + x*(-sum_xdz/(sum_xx + eps*n))) * rsqrt(sum_xx/n + eps); the two back-to-back ce_reduce_sums are safe because each opens on a barrier before touching buf.
• SOFT_MAX_BACK: scale * y * (dy - dot(y,dy)), with max_bias refused as on CPU.
• REPEAT_BACK: the repeats are gathered through the source strides, so views are exact.

"Metal carries every op the LoRA-only training graph builds" is the claim still to MEASURE: per-op parity does not prove graph coverage (GLU/swiglu backward, CONT/CPY of the transposes, GET_ROWS_BACK if embeddings ever train). The gate is one real step on the M5 with llama-finetune-lora -ngl 99 (342e6cd6, pushed right after its CPU oracle finishes) under GGML_SCHED_DEBUG=2: zero graph splits assigned to CPU, and a first-step loss within tolerance of the CPU oracle on the same data and seed.

Co-Authored-By: Claude Opus 5.5 noreply@anthropic.com
Claude-Session: https://claude.ai/code/session_01FRQWzgSfo79JtnwZywHKrE

[joelteply]

Fresh word on #5 at 04df2eb: APPROVED (read against the CPU reference line by line; the M5 counts are the measurement). RMS_NORM_BACK: sum_xx, sum_xdz, mean_eps = sum_xx/n + eps, sum_eps = sum_xx + epsn, rrms = 1/sqrt(mean_eps), dx = (dz - xsum_xdz/sum_eps)rrms — the reference's exact quantities in the same order. SOFT_MAX_BACK: dx = scaley*(dy - dot(y,dy)) and max_bias refused in supports_op as the reference asserts it. REPEAT_BACK: the (i3ne3+k3, i2ne2+k2, i1ne1+k1, i0ne0+k0) mapping is the reference's, read through src0's strides so the view cases are exact where the CPU asserts nb00 == 4. One note, not a change: the reference accumulates the rms sums in double (ggml_float) and the kernel in float; at test widths 10/10 says it is inside tolerance and at n_embd 5120 the float error is ~1e-6*sqrt(n), so fine, but if a 27B row ever trips parity, a two-pass or compensated sum in the reduction is the fix, not a tolerance change. Metal now carries every op the LoRA-only graph builds; merge on your gate. Next real gate on both nodes is Cormac's branch: the 5090 runs the same finetune on the resident 27B with the same data and seed, and we compare the two GPU curves against his CPU one.

[joelteply]

APPROVED at ca766e3 (Cormac), on the delta since 04df2eb (the Metal CE backward). d = (softmax(x) * sum(y) - y) * d_by_nr matches the CPU and CUDA fix on #20, so a masked all-zero row contributes nothing. sum_y is a second ce_reduce_sum over the same threadgroup buf, and that is safe, because the helper opens with a threadgroup barrier before it writes buf (the forward kernel already calls it twice in a row). I read the code and did not run it; your M5 run (CE_LOSS_BACK 4/4 incl. masked, grad 19753/19753) is what it rests on.

🤖 Generated with Claude Code

https://claude.ai/code/session_01FRQWzgSfo79JtnwZywHKrE