66 lines
4.0 KiB
Diff
66 lines
4.0 KiB
Diff
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: {
|