From 285ff850d6f16289037fa1e3ed502946188192a3 Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Tue, 6 Oct 2026 14:20:24 -0500 Subject: [PATCH] metal: OUT_PROD multiplies in float staging; a training gradient overflowed half kernel_mul_mm_f32_f32 stages both inputs through half in threadgroup memory. An activation fits; a training GRADIENT does not. On Qwen3.5 0.8B's LoRA backward (M5, 2026-10-06) one OUT_PROD saw |grad| up to 131273, 95 values past half's 65504. They staged as Inf, and 32,256 outputs came out NaN/Inf. Its inputs, read back right after the op, were finite and correct; the finite outputs carried half's rounding too (3053.17 against 3052.11 computed in double), so every Metal OUT_PROD, the whole LoRA backward, lost ~1e-3 relative precision. - kernel_mul_mm places sb at 64*32*sizeof(S0) (4096 for every half instantiation, as before) - a float-staged instantiation, kernel_mul_mm_f32_f32_fp32 (simdgroup kernel only, not built with the tensor API), with its pipeline getter: 12288 bytes of threadgroup memory - OUT_PROD's internal product uses it whenever the simdgroup kernel runs (the opt-in tensor API keeps its own path) This covers the whole matmul backward: ggml's MUL_MAT backward computes BOTH gradients through OUT_PROD (src0's as out_prod(src1, grad), src1's as out_prod(src0, grad^T)), so no plain MUL_MAT takes a gradient as src1. Test: test_out_prod gains b_range, and cases with |b| up to 1.5e5 (F32 and q8_0 weights, both layouts) run against CPU. Known-positive: with half staging they FAIL (NaN against CPU ~3e5, 33/37); with float staging 37/37 pass. Measured on the 0.8B, window 1024, seed 7: recompute OFF eval 1.9292085 -> 1.9291925 (the precision); layer-3-only recompute NaN -> finite. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo --- ggml/src/ggml-metal/ggml-metal-device.cpp | 35 +++++++++++++++++++++++ ggml/src/ggml-metal/ggml-metal-device.h | 1 + ggml/src/ggml-metal/ggml-metal-ops.cpp | 11 +++++-- ggml/src/ggml-metal/kernels/mul_mm.metal | 15 +++++++++- tests/test-backend-ops.cpp | 28 ++++++++++++++++-- 5 files changed, 83 insertions(+), 7 deletions(-) diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp index c5ff7265d704..22a37453bf39 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.cpp +++ b/ggml/src/ggml-metal/ggml-metal-device.cpp @@ -812,6 +812,41 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_ext(ggml_ return res; } +// F32 x F32 with FLOAT staging (kernel_mul_mm_f32_f32_fp32), for training's internal products: +// a gradient overflows half's range where an activation does not. Simdgroup kernel only: the +// caller uses it when the tensor API is off. +ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_fp32(ggml_metal_library_t lib, const ggml_tensor * op) { + GGML_ASSERT(op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32); + GGML_ASSERT(!ggml_metal_device_get_props(ggml_metal_library_get_device(lib))->has_tensor); + char name[256]; + const bool bc_inp = op->src[0]->ne[0] % 32 != 0; + const bool bc_out = op->ne[0] % 64 != 0 || op->ne[1] % 32 != 0; + GGML_ASSERT(op->src[1]->ne[2] <= INT16_MAX && op->src[1]->ne[3] <= INT16_MAX); + const int16_t ne12 = (int16_t) op->src[1]->ne[2]; + const int16_t ne13 = (int16_t) op->src[1]->ne[3]; + const int16_t r2 = (int16_t) (ne12 / op->src[0]->ne[2]); + const int16_t r3 = (int16_t) (ne13 / op->src[0]->ne[3]); + const char * base = "kernel_mul_mm_f32_f32_fp32"; + snprintf(name, 256, "%s_bci=%d_bco=%d_ne12=%d_ne13=%d_r2=%d_r3=%d", base, bc_inp, bc_out, ne12, ne13, r2, r3); + ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); + if (!res.pipeline) { + ggml_metal_cv_t cv = ggml_metal_cv_init(); + ggml_metal_cv_set_bool(cv, bc_inp, FC_MUL_MM + 0); + ggml_metal_cv_set_bool(cv, bc_out, FC_MUL_MM + 1); + ggml_metal_cv_set_int16(cv, ne12, FC_MUL_MM + 2); + ggml_metal_cv_set_int16(cv, ne13, FC_MUL_MM + 3); + ggml_metal_cv_set_int16(cv, r2, FC_MUL_MM + 4); + ggml_metal_cv_set_int16(cv, r3, FC_MUL_MM + 5); + res = ggml_metal_library_compile_pipeline(lib, base, name, cv); + ggml_metal_cv_free(cv); + } + res.nr0 = 64; + res.nr1 = 32; + res.smem = 64*32*sizeof(float) + 32*32*sizeof(float); // sa + sb; the bc_out tile (8192) fits + res.nsg = N_MM_SIMD_GROUP_X * N_MM_SIMD_GROUP_Y; + return res; +} + ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm(ggml_metal_library_t lib, const ggml_tensor * op) { char base[256]; char name[256]; diff --git a/ggml/src/ggml-metal/ggml-metal-device.h b/ggml/src/ggml-metal/ggml-metal-device.h index d4ed162acbbf..18ae227a5ca9 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.h +++ b/ggml/src/ggml-metal/ggml-metal-device.h @@ -140,6 +140,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_gated_del struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_solve_tri (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_ext (ggml_metal_library_t lib, const struct ggml_tensor * op, int nsg, int nxpsg, int r1ptg); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm (ggml_metal_library_t lib, const struct ggml_tensor * op); +struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_fp32 (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_map0 (ggml_metal_library_t lib, int ne02, int ne20); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id (ggml_metal_library_t lib, const struct ggml_tensor * op); diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index bfe38daac9ab..e8cd0301c9f1 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -2470,7 +2470,7 @@ int ggml_metal_op_pool_2d(ggml_metal_op_t ctx, int idx) { // The F32/quantized simdgroup matmul encode, over an explicit op (its src[0], src[1] and // itself may be synthetic views: OUT_PROD composes it over scratch buffers). -static void ggml_metal_encode_mul_mm(ggml_metal_op_t ctx, const ggml_tensor * op) { +static void ggml_metal_encode_mul_mm(ggml_metal_op_t ctx, const ggml_tensor * op, bool fp32_staging = false) { ggml_metal_library_t lib = ctx->lib; ggml_metal_encoder_t enc = ctx->enc; @@ -2484,7 +2484,8 @@ static void ggml_metal_encode_mul_mm(ggml_metal_op_t ctx, const ggml_tensor * op const int16_t r2 = ne12/ne02; const int16_t r3 = ne13/ne03; - auto pipeline = ggml_metal_library_get_pipeline_mul_mm(lib, op); + auto pipeline = fp32_staging ? ggml_metal_library_get_pipeline_mul_mm_fp32(lib, op) + : ggml_metal_library_get_pipeline_mul_mm(lib, op); ggml_metal_kargs_mul_mm args = { /*.ne00 =*/ ne00, @@ -4185,7 +4186,11 @@ int ggml_metal_op_out_prod(ggml_metal_op_t ctx, int idx) { p.op = GGML_OP_MUL_MAT; p.src[0] = &x; p.src[1] = &y; - ggml_metal_encode_mul_mm(ctx, &p); + // FLOAT staging when the simdgroup kernel runs: X and Y carry gradients, which pass + // half's range (65504) where activations do not; half staging turned them into Inf + // and cost every product ~1e-3 relative precision (2026-10-06, Qwen3.5 0.8B backward) + const bool has_tensor = ggml_metal_device_get_props(ggml_metal_library_get_device(ctx->lib))->has_tensor; + ggml_metal_encode_mul_mm(ctx, &p, !has_tensor); ggml_metal_op_concurrency_reset(ctx); if (!first) { diff --git a/ggml/src/ggml-metal/kernels/mul_mm.metal b/ggml/src/ggml-metal/kernels/mul_mm.metal index f5e2866b8f80..af5f080c92ee 100644 --- a/ggml/src/ggml-metal/kernels/mul_mm.metal +++ b/ggml/src/ggml-metal/kernels/mul_mm.metal @@ -157,8 +157,11 @@ kernel void kernel_mul_mm( ushort tiitg[[thread_index_in_threadgroup]], ushort sgitg[[simdgroup_index_in_threadgroup]]) { + // sa holds NR0 x NK = 64 x 32 staged values of S0; sb follows it. Sized by S0 so a + // float-staged instantiation (training: gradients exceed half's range) lays out right; + // for every half instantiation this is the 4096 it always was. threadgroup S0 * sa = (threadgroup S0 *)(shmem); - threadgroup S1 * sb = (threadgroup S1 *)(shmem + 4096); + threadgroup S1 * sb = (threadgroup S1 *)(shmem + 64*32*sizeof(S0)); constexpr int NR0 = 64; constexpr int NR1 = 32; @@ -770,6 +773,16 @@ template [[host_name("kernel_mul_mm_iq4_nl_f32")]] kernel mul_mm_t kernel_mul_m template [[host_name("kernel_mul_mm_iq4_xs_f32")]] kernel mul_mm_t kernel_mul_mm; template [[host_name("kernel_mul_mm_tq2_0_f32")]] kernel mul_mm_t kernel_mul_mm; +// FLOAT staging, for training: a gradient passes half's range (65504) where an activation +// does not (Qwen3.5 0.8B LoRA backward, 2026-10-06: |grad| 131273, 95 values over, the +// staged Inf turned 32k outputs of one OUT_PROD non-finite), and half staging also cost every +// OUT_PROD ~1e-3 relative precision. Threadgroup memory: 64*32*4 + 32*32*4 = 12288 bytes. +// The simdgroup kernel only (the tensor-API kernel has its own staging); the host asks for +// this variant only when the tensor API is off. +#ifndef GGML_METAL_HAS_TENSOR +typedef decltype(kernel_mul_mm) mul_mm_fp32_t; +template [[host_name("kernel_mul_mm_f32_f32_fp32")]] kernel mul_mm_fp32_t kernel_mul_mm; +#endif template [[host_name("kernel_mul_mm_f32_f16")]] kernel mul_mm_t kernel_mul_mm; template [[host_name("kernel_mul_mm_f16_f16")]] kernel mul_mm_t kernel_mul_mm; template [[host_name("kernel_mul_mm_q1_0_f16")]] kernel mul_mm_t kernel_mul_mm; diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 894a18fe3b20..2d8658a9e2b4 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -5031,9 +5031,23 @@ struct test_out_prod : public test_case { // address against out's buffer passed them all (Metal's F32 transposed path did, and read // stale scratch in every LoRA backward on Apple silicon until it was based on a) const bool a_own_buffer; + // b's values span [-b_range, b_range]: a training gradient passes half's range (65504) + // where an activation does not, and a kernel that staged b through half turned those + // into Inf and the products into NaN (Metal, 2026-10-06) + const float b_range; std::string vars() override { - return VARS_TO_STR9(type_a, type_b, m, n, k, bs, nr, trans_b, a_own_buffer); + return VARS_TO_STR10(type_a, type_b, m, n, k, bs, nr, trans_b, a_own_buffer, b_range); + } + + void initialize_tensors(ggml_context * ctx) override { + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { + if (t->view_src) { + continue; // a view's values are its source's + } + const bool is_b = std::string(ggml_get_name(t)).rfind("b", 0) == 0; + init_tensor_uniform(t, is_b ? -b_range : -1.0f, is_b ? b_range : 1.0f); + } } double max_nmse_err() override { @@ -5046,8 +5060,8 @@ struct test_out_prod : public test_case { int64_t m = 32, int64_t n = 32, int64_t k = 32, std::array bs = {10, 10}, std::array nr = {2, 2}, - bool trans_b = false, bool a_own_buffer = false) - : type_a(type_a), type_b(type_b), m(m), n(n), k(k), bs(bs), nr(nr), trans_b(trans_b), a_own_buffer(a_own_buffer) {} + bool trans_b = false, bool a_own_buffer = false, float b_range = 1.0f) + : type_a(type_a), type_b(type_b), m(m), n(n), k(k), bs(bs), nr(nr), trans_b(trans_b), a_own_buffer(a_own_buffer), b_range(b_range) {} ggml_tensor * build_graph(ggml_context * ctx) override { return build_graph(ctx, nullptr); @@ -5060,6 +5074,7 @@ struct test_out_prod : public test_case { ggml_tensor * b; if (trans_b) { b = ggml_new_tensor_4d(ctx, type_b, k, n, bs[0]*nr[0], bs[1]*nr[1]); + ggml_set_name(b, "b_base"); b = ggml_transpose(ctx, b); } else { b = ggml_new_tensor_4d(ctx, type_b, n, k, bs[0]*nr[0], bs[1]*nr[1]); @@ -9700,6 +9715,13 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_out_prod(GGML_TYPE_Q4_K, GGML_TYPE_F32, 256, 8, 600000, {1, 1}, {1, 1}, trans_b)); } + // gradient-sized b (past half's 65504), both layouts, F32 and q8_0 weights: the values a + // LoRA backward actually hands OUT_PROD (|grad| 131273 measured on the 0.8B) + for (ggml_type type_a : {GGML_TYPE_F32, GGML_TYPE_Q8_0}) { + for (bool trans_b : {false, true}) { + test_cases.emplace_back(new test_out_prod(type_a, GGML_TYPE_F32, 256, 64, 64, {1, 1}, {1, 1}, trans_b, false, 1.5e5f)); + } + } // a in its own buffer (a LoRA weight in the adapter's buffer), F32 and quantized, both // gradient layouts, at LoRA's shape (rank 8 against a 1024-wide activation) for (ggml_type type_a : {GGML_TYPE_F32, GGML_TYPE_Q8_0}) {