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

also, note that logical batch size > physical batch (gradient
accumulation) seems unsupported for optimization (limited to physical
, unlike in ppx - also limited to ctx-size).
training quality/convergence could be improved
by implementing (at cost of some memory, but you can make that up
by using a much smaller physical batch for a net memory savings).
presumably it's physical batch that should be limited to ctx-size?
see llama_context::opt_epoch

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)

100 lines:
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.02611±0.00077 acc=99.82±0.02% t=00:00:45
val:   data=0000008/0000008 loss=4.11112±0.22526 acc=46.36±0.78%

increasing batch+ctx sizes to 1536 (double what fits in memory for
adamw) gets apparently better validation but that could be an artifact
of continuing training from previous weights, i.e. what's train vs val
probably depends on batch size. also amusing - faster due to larger
batch even though larger context would be slower?:
train: data=0000045/0000045 loss=0.02010±0.00138 acc=99.85±0.01% t=00:00:40
val:   data=0000003/0000003 loss=1.96829±1.09488 acc=72.44±0.66%

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

@@ -1094,7 +1094,6 @@ static void common_params_print_completion(common_params_context & ctx_arg) {
                                           "llama-embedding",
                                           "llama-eval-callback",
                                           "llama-export-lora",
-                                          "llama-finetune",
                                           "llama-gen-docs",
                                           "llama-gguf",
                                           "llama-gguf-hash",
@@ -1125,9 +1124,10 @@ static void common_params_print_completion(common_params_context & ctx_arg) {
                                           "llama-speculative-simple",
                                           "llama-tokenize",
                                           "llama-tts",
-                                          "llama-vdot" };
+                                          "llama-vdot",
+                                          "llama-finetune" };
 
-    for (const auto& exe : executables) {
+    for (const auto & exe : executables) {
         printf("complete -F _llama_completions %s\n", exe.c_str());
     }
 }
@@ -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());
+                           params.optimize.optimizer = ggml_opt_get_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 ggml_opt_get_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;