examples/finetune -opt SGD (stochastic gradient descent) memory opt

support finetune arg -opt SGD (or sgd). llama 3.2-1b-F32 result:
observed 11gb gpu ram (45 sec/epoch) when using SGD instead of
19gb (55 sec/epoch) using adamw.

(getting the right learning rate for SGD is trickier than for adamw -
too high and you overshoot+oscillate, too low and you waste compute
slowly approaching convergence)

quickly reach 99%+ train accuracy on a tiny wikipedia train
(~58% token accuracy on held-out eval - reasonable)
note: objective loss not directly comparable between adamw, sgd -
check perplexity or accuracy or consider relative improvements
for convergence

train: ... loss=0.00231±0.00032 acc=99.99±0.01% t=00:00:05
val:   ... loss=3.91926±nan acc=58.40±2.18%

on more training data (500 lines), additional catastrophic forgetting before train reaches
99.9% accuracy:

train: data=0000140/0000140 loss=0.02752±0.00094 acc=99.78±0.02% t=00:00:45
val:   data=0000008/0000008 loss=4.16029±0.23384 acc=46.61±0.78%

new finetune args -wd 1e-9 to enable weight decay in sgd or adamw,
and max -epochs N (default 2 as before)

cache (1 - wd*alpha) in 'adamw' opt struct

cache computed per-epoch optimizer opts
(formerly were computed twice per)

add unit tested GGML_OPT_OPTIMIZER_SGD to ggml - avoids allocating
m, v tensors. make ggml_opt_init aware of the optimization method

since opt. memory is pre-allocated, the ggml_opt_get_optimizer_params
would probably be able to change between SGD and AdamW with each epoch
but would need to use adamw for the first (unconfirmed - no arg
to set such a policy yet)

Author: graehl

CMakeLists.txt

@@ -12,6 +12,8 @@ if (NOT XCODE AND NOT MSVC AND NOT CMAKE_BUILD_TYPE)
     set_property(CACHE CMAKE_BUILD_TYPE PROPERTY STRINGS "Debug" "Release" "MinSizeRel" "RelWithDebInfo")
 endif()
 
+message("CMAKE_BUILD_TYPE=${CMAKE_BUILD_TYPE}")
+
 # Add path to modules
 list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake/")
 

common/arg.cpp

@@ -1237,8 +1237,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
     }
     sampler_type_names.pop_back();
 
-    params.optimize             = ggml_opt_get_default_optimizer_params(NULL);
-    params.optimize.adamw.alpha = 1e-8;  // default 1e-3 is much too high for LLAMA_EXAMPLE_FINETUNE
+    params.optimize = ggml_opt_get_default_optimizer_params(NULL);
 
     /**
      * filter options by example
@@ -1438,13 +1437,12 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             params.n_predict = value;
         }
     ).set_env("LLAMA_ARG_N_PREDICT"));
-    add_opt(common_arg(
-        {"-b", "--batch-size"}, "N",
-        string_format("logical maximum batch size (default: %d)", params.n_batch),
-        [](common_params & params, int value) {
-            params.n_batch = value;
-        }
-    ).set_env("LLAMA_ARG_BATCH"));
+    add_opt(common_arg({ "-b", "--batch-size" }, "N",
+                       string_format("logical maximum batch size (default: %d) - currently reduced to -ub in optimizer "
+                                     "(TODO: gradient accumulate?)",
+                                     params.n_batch),
+                       [](common_params & params, int value) { params.n_batch = value; })
+                .set_env("LLAMA_ARG_BATCH"));
     add_opt(common_arg(
         {"-ub", "--ubatch-size"}, "N",
         string_format("physical maximum batch size (default: %d)", params.n_ubatch),
@@ -2182,19 +2180,30 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             params.ppl_output_type = value;
         }
     ).set_examples({LLAMA_EXAMPLE_PERPLEXITY}));
-    add_opt(common_arg({ "-lr", "--learning-rate" }, "ALPHA",
-                       string_format("adamw optimizer alpha (default: %.1f)", (double) params.optimize.adamw.alpha),
-                       [](common_params & params, const std::string & value) {
-                           params.optimize.adamw.alpha = std::stof(value);
-                       })
+    add_opt(
+        common_arg(
+            { "-lr", "--learning-rate" }, "ALPHA",
+            string_format(
+                "adamw or sgd optimizer alpha (default: %.2g); note: sgd alpha recommended ~100x (no momentum)",
+                (double) params.optimize.adamw.alpha),
+            [](common_params & params, const std::string & value) { params.optimize.adamw.alpha = std::stof(value); })
+            .set_examples({ LLAMA_EXAMPLE_FINETUNE }));
+    add_opt(common_arg(
+                { "-wd", "--weight-decay" }, "WD",
+                string_format(
+                    "adamw or sgd optimizer weight decay (0 is off; recommend very small e.g. 1e-9) (default: %.2g).",
+                    (double) params.optimize.adamw.wd),
+                [](common_params & params, const std::string & value) { params.optimize.adamw.wd = std::stof(value); })
+                .set_examples({ LLAMA_EXAMPLE_FINETUNE }));
+    add_opt(common_arg({ "-epochs", "--epochs" }, "N",
+                       string_format("optimizer max # of epochs (default: %d)", params.optimize.epochs),
+                       [](common_params & params, int epochs) { params.optimize.epochs = epochs; })
                 .set_examples({ LLAMA_EXAMPLE_FINETUNE }));
     add_opt(common_arg({ "-opt", "--optimizer" }, "sgd|adamw", "adamw or //TODO:sgd",
                        [](common_params & params, const std::string & name) {
                            params.optimize.optimizer = named_ggml_opt_optimizer(name.c_str());
                            if (params.optimize.optimizer == GGML_OPT_OPTIMIZER_COUNT) {
                                throw std::invalid_argument("invalid --optimizer (try adamw)");
-                           } else if (params.optimize.optimizer == GGML_OPT_OPTIMIZER_SGD) {
-                               throw std::invalid_argument("TODO: implement SGD");
                            }
                        })
                 .set_examples({ LLAMA_EXAMPLE_FINETUNE }));

examples/training/finetune.cpp

@@ -38,7 +38,6 @@ int main(int argc, char ** argv) {
     common_init();
     llama_backend_init();
     llama_numa_init(params.numa);
-
     // load the model and apply lora adapter, if any
     common_init_result llama_init = common_init_from_params(params);
     llama_model_ptr   & model = llama_init.model;
@@ -61,7 +60,16 @@ int main(int argc, char ** argv) {
     ggml_opt_dataset_t dataset = common_opt_dataset_init(ctx.get(), tokens, llama_n_ctx(ctx.get())/2);
 
     struct ggml_opt_optimizer_params & optimizer_params = params.optimize;
-    LOG_INF("-optimizer %d -lr: %.1f", optimizer_params.optimizer, (double) optimizer_params.adamw.alpha);
+    if (optimizer_params.optimizer == GGML_OPT_OPTIMIZER_SGD) {
+        double was = (double) optimizer_params.adamw.alpha;
+        double by  = 1e2;
+        double to  = was * by;
+        LOG_INF("sgd multiplying -lr by %.3g (no momentum) from -lr: %.2g to %.2g\n", by, was, to);
+        optimizer_params.adamw.alpha = to;
+    }
+
+    LOG_INF("-optimizer %s -lr %.2g -wd %.2g -epochs %d\n", ggml_opt_optimizer_name(optimizer_params.optimizer),
+            (double) optimizer_params.adamw.alpha, (double) optimizer_params.adamw.wd, optimizer_params.epochs);
 
     struct llama_opt_params lopt_params {
         /*n_ctx_train     =*/ 0,
@@ -77,7 +85,7 @@ int main(int argc, char ** argv) {
     ggml_opt_result_t result_train = ggml_opt_result_init();
     ggml_opt_result_t result_eval  = ggml_opt_result_init();
 
-    for (int epoch = 0; epoch < 2; ++epoch) {
+    for (unsigned epoch = 0; epoch < optimizer_params.epochs; ++epoch) {
         llama_opt_epoch(ctx.get(), dataset, result_train, result_eval, idata_split,
             ggml_opt_epoch_callback_progress_bar, ggml_opt_epoch_callback_progress_bar);
         fprintf(stderr, "\n");

ggml/include/ggml-opt.h

@@ -74,28 +74,33 @@ extern "C" {
         GGML_OPT_BUILD_TYPE_OPT     = 30,
     };
 
-    enum ggml_opt_optimizer {
+    enum ggml_opt_optimizer_type {
         GGML_OPT_OPTIMIZER_ADAMW,
         GGML_OPT_OPTIMIZER_SGD,
 
         GGML_OPT_OPTIMIZER_COUNT
     };
 
     // "adamw" or "sgd" (case insensitive)
-    GGML_API const char *            ggml_opt_optimizer_name(enum ggml_opt_optimizer);
-    GGML_API enum ggml_opt_optimizer named_ggml_opt_optimizer(const char *);
+    GGML_API const char *                 ggml_opt_optimizer_name(enum ggml_opt_optimizer_type);
+    GGML_API enum ggml_opt_optimizer_type named_ggml_opt_optimizer(const char *);
 
     // parameters that control which optimizer is used and how said optimizer tries to find the minimal loss
     struct ggml_opt_optimizer_params {
-        // AdamW optimizer parameters
+        // SGD and AdamW optimizer parameters
         struct {
             float alpha; // learning rate
-            float beta1;
-            float beta2;
+            float beta1;  // adamw
+            float beta2;  // adamw
             float eps;   // epsilon for numerical stability
-            float wd;    // weight decay for AdamW, use 0.0f to disable
+            float wd;    // weight decay for SGD or AdamW, use 0.0f to disable
         } adamw;
-        enum ggml_opt_optimizer optimizer;
+
+        // only GGML_OPT_OPTIMIZER_ADAMW allocates m, v per parameter
+        enum ggml_opt_optimizer_type optimizer;
+
+        // affects finetune.cpp only so far:
+        unsigned                epochs;  // max # of epochs sampling over training data
     };
 
     // callback to calculate optimizer parameters prior to a backward pass
@@ -126,6 +131,8 @@ extern "C" {
 
         ggml_opt_get_optimizer_params get_opt_pars; // callback for calculating optimizer parameters
         void * get_opt_pars_ud;                     // userdata for calculating optimizer parameters
+        struct ggml_opt_optimizer_params
+            opt_params;  // holds result of get_opt_pars(get_opt_pars_ud) after ggml_opt_init (could call get_opt_pars repeatedly instead)
     };
 
     // get parameters for an optimization context with defaults set where possible

ggml/include/ggml.h

@@ -450,7 +450,7 @@ extern "C" {
         GGML_OP_REPEAT_BACK,
         GGML_OP_CONCAT,
         GGML_OP_SILU_BACK,
-        GGML_OP_NORM, // normalize
+        GGML_OP_NORM,  // normalize
         GGML_OP_RMS_NORM,
         GGML_OP_RMS_NORM_BACK,
         GGML_OP_GROUP_NORM,
@@ -486,7 +486,7 @@ extern "C" {
         GGML_OP_POOL_1D,
         GGML_OP_POOL_2D,
         GGML_OP_POOL_2D_BACK,
-        GGML_OP_UPSCALE, // nearest interpolate
+        GGML_OP_UPSCALE,  // nearest interpolate
         GGML_OP_PAD,
         GGML_OP_PAD_REFLECT_1D,
         GGML_OP_ARANGE,
@@ -517,6 +517,7 @@ extern "C" {
         GGML_OP_CROSS_ENTROPY_LOSS,
         GGML_OP_CROSS_ENTROPY_LOSS_BACK,
         GGML_OP_OPT_STEP_ADAMW,
+        GGML_OP_OPT_STEP_SGD,
 
         GGML_OP_COUNT,
     };
@@ -2063,6 +2064,11 @@ extern "C" {
             struct ggml_tensor  * v,
             struct ggml_tensor  * adamw_params); // parameters such a the learning rate
 
+    // SGD (with weight decay) step
+    GGML_API struct ggml_tensor * ggml_opt_step_sgd(
+        struct ggml_context * ctx, struct ggml_tensor * a, struct ggml_tensor * grad,
+        struct ggml_tensor * adamw_params);  // parameters: alpha, the learning rate, and wd, weight decay
+
     //
     // automatic differentiation
     //

ggml/src/ggml-cpu/ggml-cpu.c

@@ -2057,6 +2057,11 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm
                 ggml_compute_forward_opt_step_adamw(params, tensor);
             }
             break;
+        case GGML_OP_OPT_STEP_SGD:
+            {
+                ggml_compute_forward_opt_step_sgd(params, tensor);
+            }
+            break;
         case GGML_OP_NONE:
             {
                 // nop
@@ -2341,6 +2346,7 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) {
         case GGML_OP_CROSS_ENTROPY_LOSS:
         case GGML_OP_CROSS_ENTROPY_LOSS_BACK:
         case GGML_OP_OPT_STEP_ADAMW:
+        case GGML_OP_OPT_STEP_SGD:
             {
                 n_tasks = n_threads;
             } break;

ggml/src/ggml-cpu/ops.cpp

@@ -8831,7 +8831,7 @@ static void ggml_compute_forward_opt_step_adamw_f32(
     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);
+    GGML_ASSERT(ggml_nelements(adamw_params) == 8);
 
     const int ith = params->ith;
     const int nth = params->nth;
@@ -8849,14 +8849,14 @@ static void ggml_compute_forward_opt_step_adamw_f32(
     const int ir1 = MIN(ir0 + dr, nr);
 
     const float * adamw_params_ptr = ggml_get_data_f32(adamw_params);
+
     const float alpha  = adamw_params_ptr[0];
     const float beta1  = adamw_params_ptr[1];
     const float beta2  = adamw_params_ptr[2];
     const float eps    = adamw_params_ptr[3];
-    const float wd     = adamw_params_ptr[4];
     const float beta1h = adamw_params_ptr[5];
     const float beta2h = adamw_params_ptr[6];
-
+    const float keep   = adamw_params_ptr[7];
     for (int ir = ir0; ir < ir1; ++ir) {
         const int64_t i03 = ir/(ne02*ne01);
         const int64_t i02 = (ir - i03*ne02*ne01)/ne01;
@@ -8879,7 +8879,7 @@ static void ggml_compute_forward_opt_step_adamw_f32(
             // The weight decay is applied independently of the Adam momenta m and v.
             // This is NOT equivalent to l2 regularization that adds w[i00]*w[i00] to the loss.
             // See: https://arxiv.org/pdf/1711.05101v3.pdf
-            w[i00] = w[i00]*(1.0f - alpha*wd) - alpha*mh/vh;
+            w[i00] = w[i00] * keep - alpha * mh / vh;
         }
     }
 }
@@ -8901,3 +8901,63 @@ void ggml_compute_forward_opt_step_adamw(
             }
     }
 }
+
+static void ggml_compute_forward_opt_step_sgd_f32(const ggml_compute_params * params, ggml_tensor * dst) {
+    const ggml_tensor * src0         = dst->src[0];
+    const ggml_tensor * src0_grad    = dst->src[1];
+    const ggml_tensor * adamw_params = dst->src[2];
+
+    GGML_ASSERT(ggml_are_same_shape(src0, src0_grad));
+    GGML_ASSERT(ggml_nelements(adamw_params) == 8);
+
+    const int ith = params->ith;
+    const int nth = params->nth;
+
+    const int nr = ggml_nrows(src0);
+
+    GGML_TENSOR_UNARY_OP_LOCALS
+    GGML_ASSERT(nb00 == sizeof(float));
+
+    // rows per thread
+    const int dr = (nr + nth - 1) / nth;
+
+    // row range for this thread
+    const int ir0 = dr * ith;
+    const int ir1 = MIN(ir0 + dr, nr);
+
+    // using adamw param subset we care about - alpha, wd - could have a separate struct
+    const float * adamw_params_ptr = ggml_get_data_f32(adamw_params);
+    const float   alpha            = adamw_params_ptr[0];
+    const float   keep             = adamw_params_ptr[7];
+
+    for (int ir = ir0; ir < ir1; ++ir) {
+        const int64_t i03 = ir / (ne02 * ne01);
+        const int64_t i02 = (ir - i03 * ne02 * ne01) / ne01;
+        const int64_t i01 = (ir - i03 * ne02 * ne01 - i02 * ne01);
+
+        const size_t offset = i03 * nb03 + i02 * nb02 + i01 * nb01;
+
+        float *       w = (float *) ((char *) src0->data + offset);                   // weight
+        const float * g = (const float *) ((const char *) src0_grad->data + offset);  // grad
+
+        for (int i00 = 0; i00 < ne00; ++i00) {
+            w[i00] = w[i00] * keep - alpha * g[i00];
+        }
+    }
+}
+
+void ggml_compute_forward_opt_step_sgd(const ggml_compute_params * params, ggml_tensor * dst) {
+    const ggml_tensor * src0 = dst->src[0];
+
+    switch (src0->type) {
+        case GGML_TYPE_F32:
+            {
+                ggml_compute_forward_opt_step_sgd_f32(params, dst);
+            }
+            break;
+        default:
+            {
+                GGML_ABORT("fatal error - sgd is F32 only");
+            }
+    }
+}

ggml/src/ggml-cpu/ops.h

@@ -104,7 +104,7 @@ void ggml_compute_forward_custom(const struct ggml_compute_params * params, stru
 void ggml_compute_forward_cross_entropy_loss(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_cross_entropy_loss_back(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_opt_step_adamw(const struct ggml_compute_params * params, struct ggml_tensor * dst);
-
+void ggml_compute_forward_opt_step_sgd(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 #ifdef __cplusplus
 }
 #endif

ggml/src/ggml-cuda/ggml-cuda.cu

@@ -24,6 +24,7 @@
 #include "ggml-cuda/mmvq.cuh"
 #include "ggml-cuda/norm.cuh"
 #include "ggml-cuda/opt-step-adamw.cuh"
+#include "ggml-cuda/opt-step-sgd.cuh"
 #include "ggml-cuda/out-prod.cuh"
 #include "ggml-cuda/pad.cuh"
 #include "ggml-cuda/pool2d.cuh"
@@ -2352,6 +2353,9 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg
         case GGML_OP_OPT_STEP_ADAMW:
             ggml_cuda_opt_step_adamw(ctx, dst);
             break;
+        case GGML_OP_OPT_STEP_SGD:
+            ggml_cuda_opt_step_sgd(ctx, dst);
+            break;
         default:
             return false;
     }
@@ -3256,6 +3260,7 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
         case GGML_OP_CROSS_ENTROPY_LOSS:
         case GGML_OP_CROSS_ENTROPY_LOSS_BACK:
         case GGML_OP_OPT_STEP_ADAMW:
+        case GGML_OP_OPT_STEP_SGD:
             return true;
         default:
             return false;

ggml/src/ggml-cuda/opt-step-adamw.cu

@@ -17,9 +17,9 @@ static __global__ void opt_step_adamw_f32(
     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 keep   = pars[7];
 
     const float gi = g[i];
     const float gmi = g_m[i]*beta1 +    gi*(1.0f - beta1);
@@ -31,7 +31,11 @@ static __global__ void opt_step_adamw_f32(
     const float mh =       gmi*beta1h;
     const float vh = sqrtf(gvi*beta2h) + eps;
 
-    x[i] = x[i]*(1.0f - alpha*wd) - alpha*mh/vh;
+#if 1
+    x[i] = x[i] * (1.f - alpha * pars[4]) - alpha * mh / vh;
+#else
+    x[i] = x[i] * keep - alpha * mh / vh;
+#endif
 }
 
 static void opt_step_adamw_f32_cuda(
@@ -62,14 +66,13 @@ void ggml_cuda_opt_step_adamw(ggml_backend_cuda_context & ctx, ggml_tensor * dst
     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);
+    GGML_ASSERT(ggml_nelements(adamw_params) == 8);
 
     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;
-
     cudaStream_t stream = ctx.stream();
 
     const int64_t ne = ggml_nelements(src0);