diff --git a/engine/ggml/src/ggml.c b/engine/ggml/src/ggml.c index b43016c8..55d50ad8 100644 --- a/engine/ggml/src/ggml.c +++ b/engine/ggml/src/ggml.c @@ -6597,6 +6597,52 @@ static void ggml_compute_backward( ggml_add_or_set(ctx, cgraph, isrc0, tmp); } if (src1_needs_grads) { + // HOT-Step patch: mm-backward — see engine/patches/mm-backward.patch + // + // The default arm below emits the ACTIVATION gradient as OUT_PROD. + // ggml-cuda implements OUT_PROD F32-only (cublasSgemm), so the frozen + // weight src0 must be F32, which drags the FORWARD mul_mat onto TF32 + // tensor cores too. The mul_mat formulation upstream left commented + // out here is dtype-agnostic, so a BF16 src0 rides real BF16 tensor + // cores with no dequant and no F32 window. + // + // The two are provably shape-identical: + // out_prod(src0[n,m,q1,r1], transpose(grad)[p,m,qq,rr]) -> [n,p,qq,rr] + // mul_mat(cont(transpose(src0))[m,n,q1,r1], grad[m,p,qq,rr]) -> [n,p,qq,rr] + // and their broadcast preconditions are the same pair of + // (b->ne[2] % a->ne[2] == 0, b->ne[3] % a->ne[3] == 0) checks. + // + // Measured on an RTX 5090 (engine/src/train/spike-gemmbench.h): 1.67-1.82x + // per layer per step, parity vs the TF32 out_prod arm cosine 0.999996 / + // max_rel ~4e-3. Per-use `cont` is free at these shapes, so no + // pre-transposed weight cache is needed. + // + // Env-gated so an unset environment is byte-identical to upstream. + // Set GGML_BACKWARD_MM=1 (ace-train's `--bwd mm`) to take it. Read once: + // this runs once per mul_mat per graph build, and getenv is not free. + // + // CONTIGUITY GUARD (measured, not theoretical): `grad` becomes + // mul_mat's src1, and the CUDA mul_mat kernels require src1 to be + // row-contiguous — ggml-cuda/mmf.cu:28 asserts `nb10 == ts_src1` + // and ABORTS otherwise. out_prod has no such requirement, which is + // why upstream can hand it an arbitrary transposed view. A + // non-contiguous grad therefore keeps the out_prod arm rather than + // crashing (ggml_cont on grad is not the answer: grad is an + // ACTIVATION-sized tensor, so copying it would cost more than the + // GEMM saves). Found by ace-train train-dit --self-test under + // GGML_BACKWARD_MM=1: the LoKR rung SC3 aborted here. + static int hs_bwd_mm = -1; + if (hs_bwd_mm < 0) { + const char * hs_e = getenv("GGML_BACKWARD_MM"); + hs_bwd_mm = (hs_e && hs_e[0] && strcmp(hs_e, "0") != 0) ? 1 : 0; + } + if (hs_bwd_mm && ggml_is_contiguous(grad)) { + ggml_add_or_set(ctx, cgraph, isrc1, + ggml_mul_mat(ctx, // [n,p,qq,rr] + ggml_cont(ctx, // [m,n,q1,r1] + ggml_transpose(ctx, src0)), // [m,n,q1,r1] + grad)); // [m,p,qq,rr] + } else { ggml_add_or_set(ctx, cgraph, isrc1, // ggml_mul_mat(ctx, // [n,p,qq,rr] // ggml_cont(ctx, // [m,n,q1,r1] @@ -6610,6 +6656,7 @@ static void ggml_compute_backward( src0, // [n,m,q1,r1] ggml_transpose(ctx, // [p,m,qq,rr] grad))); // [m,p,qq,rr] + } } } break; case GGML_OP_SCALE: {