whisper.rn 0.5.0-rc.9 → 0.5.1

package/cpp/ggml-metal-impl.h

@@ -1,5 +1,5 @@
-#ifndef WSP_GGML_METAL_IMPL
-#define WSP_GGML_METAL_IMPL
+#ifndef WSP_WSP_WSP_GGML_METAL_IMPL
+#define WSP_WSP_WSP_GGML_METAL_IMPL
 
 // kernel parameters for mat-vec threadgroups
 //
@@ -23,6 +23,9 @@
 #define N_R0_Q8_0 4
 #define N_SG_Q8_0 2
 
+#define N_R0_MXFP4 2
+#define N_SG_MXFP4 2
+
 #define N_R0_Q2_K 4
 #define N_SG_Q2_K 2
 
@@ -98,7 +101,7 @@ typedef struct {
     uint64_t nb2;
     uint64_t nb3;
     int32_t  dim;
-} wsp_ggml_metal_kargs_concat;
+} wsp_wsp_wsp_ggml_metal_kargs_concat;
 
 typedef struct {
     int32_t  ne00;
@@ -126,7 +129,17 @@ typedef struct {
     uint64_t nb2;
     uint64_t nb3;
     uint64_t offs;
-} wsp_ggml_metal_kargs_bin;
+    uint64_t o1[8];
+} wsp_wsp_wsp_ggml_metal_kargs_bin;
+
+typedef struct {
+    int64_t ne0;
+    int64_t ne1;
+    size_t nb01;
+    size_t nb02;
+    size_t nb11;
+    size_t nb21;
+} wsp_wsp_wsp_ggml_metal_kargs_add_id;
 
 typedef struct {
     int32_t  ne00;
@@ -145,7 +158,7 @@ typedef struct {
     uint64_t nb1;
     uint64_t nb2;
     uint64_t nb3;
-} wsp_ggml_metal_kargs_repeat;
+} wsp_wsp_wsp_ggml_metal_kargs_repeat;
 
 typedef struct {
     int64_t  ne00;
@@ -164,7 +177,7 @@ typedef struct {
     uint64_t nb1;
     uint64_t nb2;
     uint64_t nb3;
-} wsp_ggml_metal_kargs_cpy;
+} wsp_wsp_wsp_ggml_metal_kargs_cpy;
 
 typedef struct {
     int64_t  ne10;
@@ -179,7 +192,7 @@ typedef struct {
     uint64_t nb3;
     uint64_t offs;
     bool     inplace;
-} wsp_ggml_metal_kargs_set;
+} wsp_wsp_wsp_ggml_metal_kargs_set;
 
 typedef struct {
     int32_t  ne00;
@@ -211,7 +224,7 @@ typedef struct {
     int32_t  sect_1;
     int32_t  sect_2;
     int32_t  sect_3;
-} wsp_ggml_metal_kargs_rope;
+} wsp_wsp_wsp_ggml_metal_kargs_rope;
 
 typedef struct {
     int32_t  ne01;
@@ -229,16 +242,20 @@ typedef struct {
     uint64_t nb21;
     uint64_t nb22;
     uint64_t nb23;
+    int32_t  ne32;
+    int32_t  ne33;
     uint64_t nb31;
+    uint64_t nb32;
+    uint64_t nb33;
     int32_t  ne1;
     int32_t  ne2;
     float    scale;
     float    max_bias;
     float    m0;
     float    m1;
-    uint16_t n_head_log2;
+    int32_t  n_head_log2;
     float    logit_softcap;
-} wsp_ggml_metal_kargs_flash_attn_ext;
+} wsp_wsp_wsp_ggml_metal_kargs_flash_attn_ext;
 
 typedef struct {
     int32_t  ne00;
@@ -255,7 +272,7 @@ typedef struct {
     int32_t  ne1;
     int16_t  r2;
     int16_t  r3;
-} wsp_ggml_metal_kargs_mul_mm;
+} wsp_wsp_wsp_ggml_metal_kargs_mul_mm;
 
 typedef struct {
     int32_t  ne00;
@@ -276,7 +293,7 @@ typedef struct {
     int32_t  ne1;
     int16_t  r2;
     int16_t  r3;
-} wsp_ggml_metal_kargs_mul_mv;
+} wsp_wsp_wsp_ggml_metal_kargs_mul_mv;
 
 typedef struct {
     int32_t  ne00;
@@ -300,7 +317,7 @@ typedef struct {
     int16_t  nsg;
     int16_t  nxpsg;
     int16_t  r1ptg;
-} wsp_ggml_metal_kargs_mul_mv_ext;
+} wsp_wsp_wsp_ggml_metal_kargs_mul_mv_ext;
 
 typedef struct {
     int32_t  ne10;
@@ -311,7 +328,7 @@ typedef struct {
     uint64_t nbh11;
     int32_t  ne20;  // n_expert_used
     uint64_t nb21;
-} wsp_ggml_metal_kargs_mul_mm_id_map0;
+} wsp_wsp_wsp_ggml_metal_kargs_mul_mm_id_map0;
 
 typedef struct {
     int32_t  ne20; // n_expert_used
@@ -322,7 +339,7 @@ typedef struct {
     int32_t  ne0;
     uint64_t nb1;
     uint64_t nb2;
-} wsp_ggml_metal_kargs_mul_mm_id_map1;
+} wsp_wsp_wsp_ggml_metal_kargs_mul_mm_id_map1;
 
 typedef struct {
     int32_t  ne00;
@@ -339,7 +356,7 @@ typedef struct {
     int32_t  neh1;
     int16_t  r2;
     int16_t  r3;
-} wsp_ggml_metal_kargs_mul_mm_id;
+} wsp_wsp_wsp_ggml_metal_kargs_mul_mm_id;
 
 typedef struct {
     int32_t  nei0;
@@ -361,28 +378,36 @@ typedef struct {
     int32_t  ne0;
     int32_t  ne1;
     uint64_t nb1;
-} wsp_ggml_metal_kargs_mul_mv_id;
+} wsp_wsp_wsp_ggml_metal_kargs_mul_mv_id;
 
 typedef struct {
     int32_t  ne00;
     int32_t  ne00_4;
     uint64_t nb01;
     float    eps;
-} wsp_ggml_metal_kargs_norm;
+} wsp_wsp_wsp_ggml_metal_kargs_norm;
 
 typedef struct {
     int32_t  ne00;
     int32_t  ne00_4;
-    uint64_t nb01;
+    uint64_t nb1;
+    uint64_t nb2;
+    uint64_t nb3;
     float    eps;
-} wsp_ggml_metal_kargs_rms_norm;
+    int32_t  nef1[3];
+    int32_t  nef2[3];
+    int32_t  nef3[3];
+    uint64_t nbf1[3];
+    uint64_t nbf2[3];
+    uint64_t nbf3[3];
+} wsp_wsp_wsp_ggml_metal_kargs_rms_norm;
 
 typedef struct {
     int32_t  ne00;
     int32_t  ne00_4;
     uint64_t nb01;
     float    eps;
-} wsp_ggml_metal_kargs_l2_norm;
+} wsp_wsp_wsp_ggml_metal_kargs_l2_norm;
 
 typedef struct {
     int64_t  ne00;
@@ -393,7 +418,7 @@ typedef struct {
     uint64_t nb02;
     int32_t  n_groups;
     float    eps;
-} wsp_ggml_metal_kargs_group_norm;
+} wsp_wsp_wsp_ggml_metal_kargs_group_norm;
 
 typedef struct {
     int32_t  IC;
@@ -402,7 +427,7 @@ typedef struct {
     int32_t  s0;
     uint64_t nb0;
     uint64_t nb1;
-} wsp_ggml_metal_kargs_conv_transpose_1d;
+} wsp_wsp_wsp_ggml_metal_kargs_conv_transpose_1d;
 
 typedef struct {
     uint64_t  ofs0;
@@ -420,7 +445,7 @@ typedef struct {
     int32_t  KH;
     int32_t  KW;
     int32_t  KHW; // KH * KW, pre-computed on CPU to save GPU resources
-} wsp_ggml_metal_kargs_im2col;
+} wsp_wsp_wsp_ggml_metal_kargs_im2col;
 
 typedef struct{
     int32_t  ne00;
@@ -431,7 +456,9 @@ typedef struct{
     uint64_t nb1;
     int32_t  i00;
     int32_t  i10;
-} wsp_ggml_metal_kargs_glu;
+    float    alpha;
+    float    limit;
+} wsp_wsp_wsp_ggml_metal_kargs_glu;
 
 typedef struct {
     int64_t  ne00;
@@ -458,24 +485,36 @@ typedef struct {
     uint64_t nb1;
     uint64_t nb2;
     uint64_t nb3;
-} wsp_ggml_metal_kargs_sum_rows;
+} wsp_wsp_wsp_ggml_metal_kargs_sum_rows;
 
 typedef struct {
-    int64_t  ne00;
-    int64_t  ne01;
-    int64_t  ne02;
+    int32_t  ne00;
+    int32_t  ne01;
+    int32_t  ne02;
+    uint64_t nb01;
+    uint64_t nb02;
+    uint64_t nb03;
+    int32_t  ne11;
+    int32_t  ne12;
+    int32_t  ne13;
+    uint64_t nb11;
+    uint64_t nb12;
+    uint64_t nb13;
+    uint64_t nb1;
+    uint64_t nb2;
+    uint64_t nb3;
     float    scale;
     float    max_bias;
     float    m0;
     float    m1;
-    uint32_t n_head_log2;
-} wsp_ggml_metal_kargs_soft_max;
+    int32_t  n_head_log2;
+} wsp_wsp_wsp_ggml_metal_kargs_soft_max;
 
 typedef struct {
     int64_t  ne00;
     int64_t  ne01;
     int      n_past;
-} wsp_ggml_metal_kargs_diag_mask_inf;
+} wsp_wsp_wsp_ggml_metal_kargs_diag_mask_inf;
 
 typedef struct {
     int64_t  ne00;
@@ -494,32 +533,32 @@ typedef struct {
     uint64_t nb0;
     uint64_t nb1;
     uint64_t nb2;
-} wsp_ggml_metal_kargs_ssm_conv;
+} wsp_wsp_wsp_ggml_metal_kargs_ssm_conv;
 
 typedef struct {
     int64_t  d_state;
     int64_t  d_inner;
+    int64_t  n_head;
+    int64_t  n_group;
     int64_t  n_seq_tokens;
     int64_t  n_seqs;
-    uint64_t nb00;
+    int64_t  s_off;
     uint64_t nb01;
     uint64_t nb02;
-    uint64_t nb10;
+    uint64_t nb03;
     uint64_t nb11;
     uint64_t nb12;
     uint64_t nb13;
-    uint64_t nb20;
     uint64_t nb21;
     uint64_t nb22;
-    uint64_t nb30;
     uint64_t nb31;
-    uint64_t nb40;
     uint64_t nb41;
     uint64_t nb42;
-    uint64_t nb50;
+    uint64_t nb43;
     uint64_t nb51;
     uint64_t nb52;
-} wsp_ggml_metal_kargs_ssm_scan;
+    uint64_t nb53;
+} wsp_wsp_wsp_ggml_metal_kargs_ssm_scan;
 
 typedef struct {
     int64_t  ne00;
@@ -530,7 +569,7 @@ typedef struct {
     uint64_t nb11;
     uint64_t nb1;
     uint64_t nb2;
-} wsp_ggml_metal_kargs_get_rows;
+} wsp_wsp_wsp_ggml_metal_kargs_get_rows;
 
 typedef struct {
     int32_t  nk0;
@@ -546,7 +585,7 @@ typedef struct {
     uint64_t nb1;
     uint64_t nb2;
     uint64_t nb3;
-} wsp_ggml_metal_kargs_set_rows;
+} wsp_wsp_wsp_ggml_metal_kargs_set_rows;
 
 typedef struct {
     int64_t  ne00;
@@ -569,7 +608,7 @@ typedef struct {
     float    sf1;
     float    sf2;
     float    sf3;
-} wsp_ggml_metal_kargs_upscale;
+} wsp_wsp_wsp_ggml_metal_kargs_upscale;
 
 typedef struct {
     int64_t  ne00;
@@ -588,7 +627,7 @@ typedef struct {
     uint64_t nb1;
     uint64_t nb2;
     uint64_t nb3;
-} wsp_ggml_metal_kargs_pad;
+} wsp_wsp_wsp_ggml_metal_kargs_pad;
 
 typedef struct {
     int64_t  ne00;
@@ -609,28 +648,28 @@ typedef struct {
     uint64_t nb3;
     int32_t  p0;
     int32_t  p1;
-} wsp_ggml_metal_kargs_pad_reflect_1d;
+} wsp_wsp_wsp_ggml_metal_kargs_pad_reflect_1d;
 
 typedef struct {
     uint64_t nb1;
     int      dim;
     int      max_period;
-} wsp_ggml_metal_kargs_timestep_embedding;
+} wsp_wsp_wsp_ggml_metal_kargs_timestep_embedding;
 
 typedef struct {
     float    slope;
-} wsp_ggml_metal_kargs_leaky_relu;
+} wsp_wsp_wsp_ggml_metal_kargs_leaky_relu;
 
 typedef struct {
     int64_t  ncols;
     int64_t  ncols_pad;
-} wsp_ggml_metal_kargs_argsort;
+} wsp_wsp_wsp_ggml_metal_kargs_argsort;
 
 typedef struct {
     int64_t  ne0;
     float    start;
     float    step;
-} wsp_ggml_metal_kargs_arange;
+} wsp_wsp_wsp_ggml_metal_kargs_arange;
 
 typedef struct {
     int32_t  k0;
@@ -644,6 +683,6 @@ typedef struct {
     int64_t  OH;
     int64_t  OW;
     int64_t  parallel_elements;
-} wsp_ggml_metal_kargs_pool_2d;
+} wsp_wsp_wsp_ggml_metal_kargs_pool_2d;
 
-#endif // WSP_GGML_METAL_IMPL
+#endif // WSP_WSP_WSP_GGML_METAL_IMPL

package/cpp/ggml-metal.h

@@ -39,18 +39,13 @@ extern "C" {
 // user-code should use only these functions
 //
 
+// TODO: remove in the future
 WSP_GGML_BACKEND_API wsp_ggml_backend_t wsp_ggml_backend_metal_init(void);
 
 WSP_GGML_BACKEND_API bool wsp_ggml_backend_is_metal(wsp_ggml_backend_t backend);
 
-WSP_GGML_DEPRECATED(
-        WSP_GGML_BACKEND_API wsp_ggml_backend_buffer_t wsp_ggml_backend_metal_buffer_from_ptr(void * data, size_t size, size_t max_size),
-        "obsoleted by the new device interface - https://github.com/ggml-org/llama.cpp/pull/9713");
-
 WSP_GGML_BACKEND_API void wsp_ggml_backend_metal_set_abort_callback(wsp_ggml_backend_t backend, wsp_ggml_abort_callback abort_callback, void * user_data);
 
-WSP_GGML_BACKEND_API wsp_ggml_backend_buffer_type_t wsp_ggml_backend_metal_buffer_type(void);
-
 // helper to check if the device supports a specific family
 // ideally, the user code should be doing these checks
 // ref: https://developer.apple.com/metal/Metal-Feature-Set-Tables.pdf

package/cpp/ggml-opt.cpp

@@ -64,9 +64,11 @@ struct wsp_ggml_opt_context {
     int32_t opt_i              = 0;
     bool    loss_per_datapoint = false;
 
-    wsp_ggml_opt_get_optimizer_params get_opt_pars = nullptr;
-    void * get_opt_pars_ud                     = nullptr;
-    struct wsp_ggml_tensor * adamw_params          = nullptr;
+    wsp_ggml_opt_get_optimizer_params get_opt_pars    = nullptr;
+    void *                        get_opt_pars_ud = nullptr;
+    struct wsp_ggml_tensor *          opt_step_params = nullptr; // Stores output of get_opt_pars.
+
+    enum wsp_ggml_opt_optimizer_type optimizer = WSP_GGML_OPT_OPTIMIZER_TYPE_ADAMW;
 };
 
 struct wsp_ggml_opt_result {
@@ -229,9 +231,13 @@ struct wsp_ggml_opt_optimizer_params wsp_ggml_opt_get_default_optimizer_params(v
     result.adamw.eps   = 1e-8f;
     result.adamw.wd    = 0.0f;
 
+    result.sgd.alpha   = 1e-3f;
+    result.sgd.wd      = 0.0f;
+
     return result;
 }
 
+
 struct wsp_ggml_opt_optimizer_params wsp_ggml_opt_get_constant_optimizer_params(void * userdata) {
     return *((struct wsp_ggml_opt_optimizer_params *) userdata);
 }
@@ -249,6 +255,7 @@ struct wsp_ggml_opt_params wsp_ggml_opt_default_params(
         /*opt_period      =*/ 1,
         /*get_opt_pars    =*/ wsp_ggml_opt_get_default_optimizer_params,
         /*get_opt_pars_ud =*/ nullptr,
+        /*optimizer       =*/ WSP_GGML_OPT_OPTIMIZER_TYPE_ADAMW,
     };
 }
 
@@ -316,9 +323,14 @@ static void wsp_ggml_opt_build(wsp_ggml_opt_context_t opt_ctx) {
     WSP_GGML_ASSERT(opt_ctx->ctx_compute && "no compute context set, either use static graphs or set one with wsp_ggml_opt_prepare_alloc");
     WSP_GGML_ASSERT((!opt_ctx->static_graphs || opt_ctx->inputs->data) && "when using static graphs the inputs must be allocated statically");
 
+    const enum wsp_ggml_opt_optimizer_type optimizer = opt_ctx->optimizer;
+
     const bool accumulate = opt_ctx->build_type_alloc >= WSP_GGML_OPT_BUILD_TYPE_GRAD &&
         !(opt_ctx->static_graphs && opt_ctx->build_type_alloc == WSP_GGML_OPT_BUILD_TYPE_OPT && opt_ctx->opt_period == 1);
 
+    const bool need_momenta = opt_ctx->build_type_alloc == WSP_GGML_OPT_BUILD_TYPE_OPT &&
+        opt_ctx->optimizer == WSP_GGML_OPT_OPTIMIZER_TYPE_ADAMW;
+
     wsp_ggml_set_input(opt_ctx->inputs);
     wsp_ggml_set_output(opt_ctx->outputs);
 
@@ -340,8 +352,7 @@ static void wsp_ggml_opt_build(wsp_ggml_opt_context_t opt_ctx) {
         //   - pred (if using static graphs)
         //   - ncorrect (if using static graphs, 2 tensors).
         constexpr size_t n_loss = 1;
-        const size_t tensors_per_param = (accumulate ? 1 : 0) +
-            (opt_ctx->build_type_alloc == WSP_GGML_OPT_BUILD_TYPE_OPT ? 2 : 0);
+        const size_t tensors_per_param = (accumulate ? 1 : 0) + (need_momenta ? 2 : 0);
         const size_t tensors_const = opt_ctx->static_graphs ? 9 : 0;
         const size_t size_meta = (n_loss + tensors_per_param*n_param + tensors_const) * wsp_ggml_tensor_overhead();
         struct wsp_ggml_init_params params = {
@@ -458,7 +469,7 @@ static void wsp_ggml_opt_build(wsp_ggml_opt_context_t opt_ctx) {
             }
         }
 
-        if (opt_ctx->build_type_alloc >= WSP_GGML_OPT_BUILD_TYPE_OPT) {
+        if (need_momenta && opt_ctx->build_type_alloc >= WSP_GGML_OPT_BUILD_TYPE_OPT) {
             opt_ctx->grad_m.resize(n_nodes);
             opt_ctx->grad_v.resize(n_nodes);
             for (int i = 0; i < n_nodes; ++i) {
@@ -492,23 +503,36 @@ static void wsp_ggml_opt_build(wsp_ggml_opt_context_t opt_ctx) {
     // gb_opt == graph backward optimize, forward pass, then backward pass to calculate gradients, then optimizer step.
     opt_ctx->gb_opt = wsp_ggml_graph_dup(opt_ctx->ctx_compute, opt_ctx->gb_grad, /*force_grads =*/ true);
 
-    opt_ctx->adamw_params = wsp_ggml_new_tensor_1d(opt_ctx->ctx_cpu, WSP_GGML_TYPE_F32, 7);
-    wsp_ggml_set_input(opt_ctx->adamw_params);
-    wsp_ggml_set_name(opt_ctx->adamw_params, "adamw_params");
-
+    opt_ctx->opt_step_params = wsp_ggml_new_tensor_1d(opt_ctx->ctx_cpu, WSP_GGML_TYPE_F32, need_momenta ? 7 : 2);
+    wsp_ggml_tensor * adamw_params = opt_ctx->opt_step_params;
+    wsp_ggml_set_input(adamw_params);
+    const char * optimizer_name = wsp_ggml_opt_optimizer_name(opt_ctx->optimizer);
+    wsp_ggml_format_name(adamw_params, "%s_params", optimizer_name);
     for (int i = opt_ctx->gf->n_nodes-1; i >= 0; --i) {
         struct wsp_ggml_tensor * node = opt_ctx->gb_opt->nodes[i];
         struct wsp_ggml_tensor * grad = wsp_ggml_graph_get_grad(opt_ctx->gb_opt, node);
 
         if (grad && (node->flags & WSP_GGML_TENSOR_FLAG_PARAM)) {
-            struct wsp_ggml_tensor * m        = opt_ctx->grad_m[i];
-            struct wsp_ggml_tensor * v        = opt_ctx->grad_v[i];
-            struct wsp_ggml_tensor * opt_step = wsp_ggml_opt_step_adamw(opt_ctx->ctx_compute, node, grad, m, v, opt_ctx->adamw_params);
-
-            wsp_ggml_set_name(m,        (std::string("AdamW m for ")    + std::string(node->name)).c_str());
-            wsp_ggml_set_name(v,        (std::string("AdamW v for ")    + std::string(node->name)).c_str());
-            wsp_ggml_set_name(opt_step, (std::string("AdamW step for ") + std::string(node->name)).c_str());
-
+            struct wsp_ggml_tensor * m = nullptr;
+            struct wsp_ggml_tensor * v = nullptr;
+            if (need_momenta) {
+                m = opt_ctx->grad_m[i];
+                v = opt_ctx->grad_v[i];
+                wsp_ggml_format_name(m, "AdamW m for %s", node->name);
+                wsp_ggml_format_name(v, "AdamW v for %s", node->name);
+            }
+            struct wsp_ggml_tensor * opt_step;
+            switch (optimizer) {
+                case WSP_GGML_OPT_OPTIMIZER_TYPE_ADAMW:
+                    opt_step = wsp_ggml_opt_step_adamw(opt_ctx->ctx_compute, node, grad, m, v, adamw_params);
+                    break;
+                case WSP_GGML_OPT_OPTIMIZER_TYPE_SGD:
+                    opt_step = wsp_ggml_opt_step_sgd(opt_ctx->ctx_compute, node, grad, adamw_params);
+                    break;
+                default:
+                    WSP_GGML_ABORT("fatal error");
+            }
+            wsp_ggml_format_name(opt_step, "%s step for %s", optimizer_name, node->name);
             wsp_ggml_build_forward_expand(opt_ctx->gb_opt, opt_step);
         }
     }
@@ -534,6 +558,7 @@ wsp_ggml_opt_context_t wsp_ggml_opt_init(struct wsp_ggml_opt_params params) {
     result->opt_period       = params.opt_period;
     result->get_opt_pars     = params.get_opt_pars;
     result->get_opt_pars_ud  = params.get_opt_pars_ud;
+    result->optimizer        = params.optimizer;
 
     WSP_GGML_ASSERT(result->opt_period >= 1);
 
@@ -756,29 +781,43 @@ void wsp_ggml_opt_alloc(wsp_ggml_opt_context_t opt_ctx, bool backward) {
 void wsp_ggml_opt_eval(wsp_ggml_opt_context_t opt_ctx, wsp_ggml_opt_result_t result) {
     WSP_GGML_ASSERT(opt_ctx->eval_ready);
     if (opt_ctx->allocated_graph == opt_ctx->gb_opt) {
-        struct wsp_ggml_opt_optimizer_params opt_pars = opt_ctx->get_opt_pars(opt_ctx->get_opt_pars_ud);
-
-        WSP_GGML_ASSERT(opt_pars.adamw.alpha >  0.0f);
-        WSP_GGML_ASSERT(opt_pars.adamw.beta1 >= 0.0f);
-        WSP_GGML_ASSERT(opt_pars.adamw.beta1 <= 1.0f);
-        WSP_GGML_ASSERT(opt_pars.adamw.beta2 >= 0.0f);
-        WSP_GGML_ASSERT(opt_pars.adamw.beta2 <= 1.0f);
-        WSP_GGML_ASSERT(opt_pars.adamw.eps   >= 0.0f);
-        WSP_GGML_ASSERT(opt_pars.adamw.wd    >= 0.0f);
-        WSP_GGML_ASSERT(opt_pars.adamw.wd    <= 1.0f);
-
-        // beta1, beta2 after applying warmup
-        const float beta1h = 1.0f/(1.0f - powf(opt_pars.adamw.beta1, opt_ctx->iter));
-        const float beta2h = 1.0f/(1.0f - powf(opt_pars.adamw.beta2, opt_ctx->iter));
-
-        float * adamw_par_data = wsp_ggml_get_data_f32(opt_ctx->adamw_params);
-        adamw_par_data[0] = opt_pars.adamw.alpha;
-        adamw_par_data[1] = opt_pars.adamw.beta1;
-        adamw_par_data[2] = opt_pars.adamw.beta2;
-        adamw_par_data[3] = opt_pars.adamw.eps;
-        adamw_par_data[4] = opt_pars.adamw.wd;
-        adamw_par_data[5] = beta1h;
-        adamw_par_data[6] = beta2h;
+        const wsp_ggml_opt_optimizer_params & opt_pars = opt_ctx->get_opt_pars(opt_ctx->get_opt_pars_ud);
+
+        switch (opt_ctx->optimizer) {
+            case WSP_GGML_OPT_OPTIMIZER_TYPE_ADAMW: {
+                WSP_GGML_ASSERT(opt_pars.adamw.alpha > 0.0f);
+                WSP_GGML_ASSERT(opt_pars.adamw.beta1 >= 0.0f);
+                WSP_GGML_ASSERT(opt_pars.adamw.beta1 <= 1.0f);
+                WSP_GGML_ASSERT(opt_pars.adamw.beta2 >= 0.0f);
+                WSP_GGML_ASSERT(opt_pars.adamw.beta2 <= 1.0f);
+                WSP_GGML_ASSERT(opt_pars.adamw.eps >= 0.0f);
+                WSP_GGML_ASSERT(opt_pars.adamw.wd >= 0.0f);
+                WSP_GGML_ASSERT(opt_pars.adamw.wd <= 1.0f);
+
+                // beta1, beta2 after applying warmup
+                const float beta1h = 1.0f / (1.0f - powf(opt_pars.adamw.beta1, opt_ctx->iter));
+                const float beta2h = 1.0f / (1.0f - powf(opt_pars.adamw.beta2, opt_ctx->iter));
+
+                float * adamw_par_data = wsp_ggml_get_data_f32(opt_ctx->opt_step_params);
+                adamw_par_data[0] = opt_pars.adamw.alpha;
+                adamw_par_data[1] = opt_pars.adamw.beta1;
+                adamw_par_data[2] = opt_pars.adamw.beta2;
+                adamw_par_data[3] = opt_pars.adamw.eps;
+                adamw_par_data[4] = opt_pars.adamw.wd;
+                adamw_par_data[5] = beta1h;
+                adamw_par_data[6] = beta2h;
+            } break;
+            case WSP_GGML_OPT_OPTIMIZER_TYPE_SGD: {
+                WSP_GGML_ASSERT(opt_pars.sgd.alpha > 0.0f);
+                WSP_GGML_ASSERT(opt_pars.sgd.wd >= 0.0f);
+                WSP_GGML_ASSERT(opt_pars.sgd.wd <= 1.0f);
+                float * sgd = wsp_ggml_get_data_f32(opt_ctx->opt_step_params);
+                sgd[0] = opt_pars.sgd.alpha;
+                sgd[1] = opt_pars.sgd.wd;
+            } break;
+            default:
+                WSP_GGML_ABORT("fatal error");
+        }
     }
 
     wsp_ggml_backend_sched_graph_compute(opt_ctx->backend_sched, opt_ctx->allocated_graph_copy);
@@ -963,6 +1002,7 @@ void wsp_ggml_opt_fit(
         wsp_ggml_tensor                   * outputs,
         wsp_ggml_opt_dataset_t              dataset,
         enum wsp_ggml_opt_loss_type         loss_type,
+        enum wsp_ggml_opt_optimizer_type    optimizer,
         wsp_ggml_opt_get_optimizer_params   get_opt_pars,
         int64_t                         nepoch,
         int64_t                         nbatch_logical,
@@ -993,6 +1033,7 @@ void wsp_ggml_opt_fit(
     params.opt_period      = opt_period;
     params.get_opt_pars    = get_opt_pars;
     params.get_opt_pars_ud = &epoch;
+    params.optimizer       = optimizer;
     wsp_ggml_opt_context_t opt_ctx = wsp_ggml_opt_init(params);
 
     // Shuffling the data is generally useful but there is only a point if not all data is used in a single batch.
@@ -1035,3 +1076,18 @@ void wsp_ggml_opt_fit(
     wsp_ggml_opt_result_free(result_train);
     wsp_ggml_opt_result_free(result_val);
 }
+
+enum wsp_ggml_opt_optimizer_type wsp_ggml_opt_context_optimizer_type(wsp_ggml_opt_context_t c) {
+    return c->optimizer;
+}
+
+WSP_GGML_API const char * wsp_ggml_opt_optimizer_name(enum wsp_ggml_opt_optimizer_type o) {
+    switch (o) {
+        case WSP_GGML_OPT_OPTIMIZER_TYPE_ADAMW:
+            return "adamw";
+        case WSP_GGML_OPT_OPTIMIZER_TYPE_SGD:
+            return "sgd";
+        default:
+            return "undefined";
+    };
+}