From 0f3861385f9f9d980c4c13d7e1b2dd7e9502560d Mon Sep 17 00:00:00 2001 From: Neo Zhang Date: Mon, 17 Aug 2026 12:31:29 +0800 Subject: [PATCH] support OP OPT_STEP_ADAMW, OPT_STEP_SGD (llama/25268) * fix conflict * fix conflict of ops.md * fix conflict of ops.md * update the ops.md --------- Co-authored-by: Neo Zhang Jianyu --- ggml/src/ggml-sycl/ggml-sycl.cpp | 9 +++ ggml/src/ggml-sycl/opt-step.cpp | 131 +++++++++++++++++++++++++++++++ ggml/src/ggml-sycl/opt-step.hpp | 6 ++ 3 files changed, 146 insertions(+) create mode 100644 ggml/src/ggml-sycl/opt-step.cpp create mode 100644 ggml/src/ggml-sycl/opt-step.hpp diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index d1456a867..5416d4f0f 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -77,6 +77,7 @@ #include "ggml-sycl/fill.hpp" #include "ggml-sycl/cumsum.hpp" #include "ggml-sycl/diag.hpp" +#include "ggml-sycl/opt-step.hpp" #include "ggml-sycl/solve_tri.hpp" #include "ggml-sycl/gated_delta_net.hpp" #include "ggml-sycl/pool.hpp" @@ -5355,6 +5356,12 @@ static bool ggml_sycl_compute_forward(ggml_backend_sycl_context & ctx, struct gg case GGML_OP_GATED_DELTA_NET: ggml_sycl_gated_delta_net(ctx, dst); break; + case GGML_OP_OPT_STEP_ADAMW: + ggml_sycl_opt_step_adamw(ctx, dst); + break; + case GGML_OP_OPT_STEP_SGD: + ggml_sycl_opt_step_sgd(ctx, dst); + break; case GGML_OP_SSM_CONV: ggml_sycl_ssm_conv(ctx, dst); break; @@ -6263,6 +6270,8 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons case GGML_OP_RWKV_WKV7: case GGML_OP_GATED_LINEAR_ATTN: case GGML_OP_GATED_DELTA_NET: + case GGML_OP_OPT_STEP_ADAMW: + case GGML_OP_OPT_STEP_SGD: return true; case GGML_OP_SSM_CONV: return op->type == GGML_TYPE_F32 && diff --git a/ggml/src/ggml-sycl/opt-step.cpp b/ggml/src/ggml-sycl/opt-step.cpp new file mode 100644 index 000000000..6d919a71e --- /dev/null +++ b/ggml/src/ggml-sycl/opt-step.cpp @@ -0,0 +1,131 @@ +#include "opt-step.hpp" + +#define SYCL_OPT_STEP_BLOCK_SIZE 256 + +template +static void opt_step_adamw_f32_kernel( + T * __restrict__ x, + const T * __restrict__ g, + T * __restrict__ g_m, + T * __restrict__ g_v, + const T * __restrict__ pars, + const int64_t k, + const sycl::nd_item<1> & item) { + + const int64_t i = (int64_t) item.get_global_id(0); + if (i >= k) { + return; + } + + const float alpha = pars[0]; + const float beta1 = pars[1]; + const float beta2 = pars[2]; + const float eps = pars[3]; + const float wd = pars[4]; + const float beta1h = pars[5]; + const float beta2h = pars[6]; + + const float gi = g[i]; + const float gmi = g_m[i] * beta1 + gi * (1.0f - beta1); + const float gvi = g_v[i] * beta2 + gi * gi * (1.0f - beta2); + + g_m[i] = gmi; + g_v[i] = gvi; + + const float mh = gmi * beta1h; + const float vh = sycl::sqrt(gvi * beta2h) + eps; + + x[i] = x[i] * (1.0f - alpha * wd) - alpha * mh / vh; +} + +template +static void opt_step_sgd_f32_kernel( + T * __restrict__ x, + const T * __restrict__ g, + const T * __restrict__ pars, + const int64_t k, + const sycl::nd_item<1> & item) { + + const int64_t i = (int64_t) item.get_global_id(0); + if (i >= k) { + return; + } + + x[i] = x[i] * (1.0f - pars[0] * pars[1]) - pars[0] * g[i]; +} + +void ggml_sycl_opt_step_adamw(ggml_backend_sycl_context & ctx, ggml_tensor * dst) { + scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/5); + + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src0_grad = dst->src[1]; + const ggml_tensor * src0_grad_m = dst->src[2]; + const ggml_tensor * src0_grad_v = dst->src[3]; + const ggml_tensor * adamw_params = dst->src[4]; + + GGML_ASSERT(src0->type == GGML_TYPE_F32); + GGML_ASSERT(src0_grad->type == GGML_TYPE_F32); + GGML_ASSERT(src0_grad_m->type == GGML_TYPE_F32); + GGML_ASSERT(src0_grad_v->type == GGML_TYPE_F32); + GGML_ASSERT(adamw_params->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(src0)); + GGML_ASSERT(ggml_is_contiguous(src0_grad)); + GGML_ASSERT(ggml_is_contiguous(src0_grad_m)); + GGML_ASSERT(ggml_is_contiguous(src0_grad_v)); + GGML_ASSERT(ggml_is_contiguous(adamw_params)); + GGML_ASSERT(ggml_are_same_shape(src0, src0_grad)); + GGML_ASSERT(ggml_are_same_shape(src0, src0_grad_m)); + GGML_ASSERT(ggml_are_same_shape(src0, src0_grad_v)); + GGML_ASSERT(ggml_nelements(adamw_params) == 7); + + dpct::queue_ptr stream = ctx.stream(); + SYCL_CHECK(ggml_sycl_set_device(ctx.device)); + + float * src0_d = (float *) src0->data; + const float * src0_grad_d = (const float *) src0_grad->data; + float * src0_grad_m_d = (float *) src0_grad_m->data; + float * src0_grad_v_d = (float *) src0_grad_v->data; + const float * adamw_params_d = (const float *) adamw_params->data; + + const int64_t ne = ggml_nelements(src0); + const int64_t num_blocks = (ne + SYCL_OPT_STEP_BLOCK_SIZE - 1) / SYCL_OPT_STEP_BLOCK_SIZE; + + stream->parallel_for( + sycl::nd_range<1>(num_blocks * SYCL_OPT_STEP_BLOCK_SIZE, SYCL_OPT_STEP_BLOCK_SIZE), + [=](sycl::nd_item<1> item) { + opt_step_adamw_f32_kernel(src0_d, src0_grad_d, src0_grad_m_d, src0_grad_v_d, adamw_params_d, ne, item); + }); +} + +void ggml_sycl_opt_step_sgd(ggml_backend_sycl_context & ctx, ggml_tensor * dst) { + scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/3); + + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src0_grad = dst->src[1]; + const ggml_tensor * sgd_params = dst->src[2]; + + GGML_ASSERT(src0->type == GGML_TYPE_F32); + GGML_ASSERT(src0_grad->type == GGML_TYPE_F32); + GGML_ASSERT(sgd_params->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(src0)); + GGML_ASSERT(ggml_is_contiguous(src0_grad)); + GGML_ASSERT(ggml_is_contiguous(sgd_params)); + GGML_ASSERT(ggml_are_same_shape(src0, src0_grad)); + GGML_ASSERT(ggml_nelements(sgd_params) == 2); + + dpct::queue_ptr stream = ctx.stream(); + SYCL_CHECK(ggml_sycl_set_device(ctx.device)); + + float * src0_d = (float *) src0->data; + const float * src0_grad_d = (const float *) src0_grad->data; + const float * sgd_params_d = (const float *) sgd_params->data; + + const int64_t ne = ggml_nelements(src0); + const int64_t num_blocks = (ne + SYCL_OPT_STEP_BLOCK_SIZE - 1) / SYCL_OPT_STEP_BLOCK_SIZE; + + stream->parallel_for( + sycl::nd_range<1>(num_blocks * SYCL_OPT_STEP_BLOCK_SIZE, SYCL_OPT_STEP_BLOCK_SIZE), + [=](sycl::nd_item<1> item) { + opt_step_sgd_f32_kernel(src0_d, src0_grad_d, sgd_params_d, ne, item); + }); +} diff --git a/ggml/src/ggml-sycl/opt-step.hpp b/ggml/src/ggml-sycl/opt-step.hpp new file mode 100644 index 000000000..dcd633b22 --- /dev/null +++ b/ggml/src/ggml-sycl/opt-step.hpp @@ -0,0 +1,6 @@ +#pragma once + +#include "common.hpp" + +void ggml_sycl_opt_step_adamw(ggml_backend_sycl_context & ctx, ggml_tensor * dst); +void ggml_sycl_opt_step_sgd(ggml_backend_sycl_context & ctx, ggml_tensor * dst);