cui-llama.rn 1.1.6 → 1.2.0

package/cpp/ggml.c

@@ -1,7 +1,9 @@
 #define _CRT_SECURE_NO_DEPRECATE // Disables ridiculous "unsafe" warnings on Windows
 #define _USE_MATH_DEFINES // For M_PI on MSVC
 
+#include "ggml-backend.h"
 #include "ggml-impl.h"
+#include "ggml-cpu-impl.h"
 #include "ggml-quants.h"
 #include "ggml.h"
 #include "ggml-aarch64.h"
@@ -61,6 +63,25 @@ int lm_ggml_sve_cnt_b = 0;
 #pragma warning(disable: 4702)
 #endif
 
+// Note: once we move threading into a separate C++ file
+// will use std::hardware_destructive_interference_size instead of hardcoding it here
+// and we'll use C++ attribute syntax.
+#define LM_GGML_CACHE_LINE  64
+
+#if defined(__clang__) || defined(__GNUC__)
+#define LM_GGML_CACHE_ALIGN __attribute__((aligned(LM_GGML_CACHE_LINE)))
+#endif
+
+#if defined(__has_feature)
+#if __has_feature(thread_sanitizer)
+#define LM_GGML_TSAN_ENABLED 1
+#endif
+#else  // __has_feature
+#if defined(__SANITIZE_THREAD__)
+#define LM_GGML_TSAN_ENABLED 1
+#endif
+#endif // __has_feature
+
 #if defined(_WIN32)
 
 #define WIN32_LEAN_AND_MEAN
@@ -70,6 +91,8 @@ int lm_ggml_sve_cnt_b = 0;
 #include <windows.h>
 
 #if !defined(__clang__)
+#define LM_GGML_CACHE_ALIGN __declspec(align(LM_GGML_CACHE_LINE))
+
 typedef volatile LONG atomic_int;
 typedef atomic_int atomic_bool;
 typedef atomic_int atomic_flag;
@@ -112,6 +135,9 @@ static atomic_bool atomic_flag_test_and_set(atomic_flag * ptr) {
 static void atomic_flag_clear(atomic_flag * ptr) {
     InterlockedExchange(ptr, 0);
 }
+static void atomic_thread_fence(memory_order mo) {
+    MemoryBarrier();
+}
 #else // clang
 #include <stdatomic.h>
 #endif
@@ -287,7 +313,6 @@ void lm_ggml_abort(const char * file, int line, const char * fmt, ...) {
 #define LM_GGML_DEBUG 0
 #define LM_GGML_GELU_FP16
 #define LM_GGML_GELU_QUICK_FP16
-#define LM_GGML_N_TASKS_MAX (-1)
 
 #define LM_GGML_SOFT_MAX_UNROLL 4
 #define LM_GGML_VEC_DOT_UNROLL  2
@@ -2005,17 +2030,18 @@ struct lm_ggml_threadpool {
 
     // synchronization primitives
     atomic_int n_graph;       // incremented when there is work to be done (i.e each graph)
-    atomic_int n_barrier;
-    atomic_int n_barrier_passed;
+    atomic_int LM_GGML_CACHE_ALIGN n_barrier;
+    atomic_int LM_GGML_CACHE_ALIGN n_barrier_passed;
     atomic_int current_chunk; // currently processing chunk during Mat_Mul, shared between all the threads.
 
     // these are atomic as an annotation for thread-sanitizer
     atomic_bool stop;         // Used for stopping the threadpool altogether
     atomic_bool pause;        // Used for pausing the threadpool or individual threads
+    atomic_bool abort;        // Used for aborting processing of a graph
 
     struct lm_ggml_compute_state * workers;   // per thread state
     int          n_threads_max; // number of threads in the pool
-    int          n_threads_cur; // number of threads used in the current graph
+    atomic_int   n_threads_cur; // number of threads used in the current graph
 
     int32_t      prio;        // Scheduling priority
     uint32_t     poll;        // Polling level (0 - no polling)
@@ -2995,9 +3021,10 @@ static const char * LM_GGML_OP_NAME[LM_GGML_OP_COUNT] = {
 
     "CROSS_ENTROPY_LOSS",
     "CROSS_ENTROPY_LOSS_BACK",
+    "OPT_STEP_ADAMW",
 };
 
-static_assert(LM_GGML_OP_COUNT == 79, "LM_GGML_OP_COUNT != 79");
+static_assert(LM_GGML_OP_COUNT == 80, "LM_GGML_OP_COUNT != 80");
 
 static const char * LM_GGML_OP_SYMBOL[LM_GGML_OP_COUNT] = {
     "none",
@@ -3088,9 +3115,10 @@ static const char * LM_GGML_OP_SYMBOL[LM_GGML_OP_COUNT] = {
 
     "cross_entropy_loss(x,y)",
     "cross_entropy_loss_back(x,y)",
+    "adamw(x)",
 };
 
-static_assert(LM_GGML_OP_COUNT == 79, "LM_GGML_OP_COUNT != 79");
+static_assert(LM_GGML_OP_COUNT == 80, "LM_GGML_OP_COUNT != 80");
 
 static_assert(LM_GGML_OP_POOL_COUNT == 2, "LM_GGML_OP_POOL_COUNT != 2");
 
@@ -3177,41 +3205,43 @@ inline static void lm_ggml_critical_section_start(void) {
     }
 }
 
-#ifdef LM_GGML_USE_OPENMP
-static void lm_ggml_barrier(struct lm_ggml_threadpool * threadpool) {
-    if (threadpool->n_threads_cur == 1) {
+static void lm_ggml_barrier(struct lm_ggml_threadpool * tp) {
+    int n_threads = atomic_load_explicit(&tp->n_threads_cur, memory_order_relaxed);
+    if (n_threads == 1) {
         return;
     }
 
+#ifdef LM_GGML_USE_OPENMP
     #pragma omp barrier
-}
 #else
-static void lm_ggml_barrier(struct lm_ggml_threadpool * threadpool) {
-    if (threadpool->n_threads_cur == 1) {
-        return;
-    }
+    int n_passed = atomic_load_explicit(&tp->n_barrier_passed, memory_order_relaxed);
 
-    atomic_int * n_barrier = &threadpool->n_barrier;
-    atomic_int * n_barrier_passed = &threadpool->n_barrier_passed;
+    // enter barrier (full seq-cst fence)
+    int n_barrier = atomic_fetch_add_explicit(&tp->n_barrier, 1, memory_order_seq_cst);
 
-    int n_threads = threadpool->n_threads_cur;
-    int passed_old = atomic_load_explicit(n_barrier_passed, memory_order_relaxed);
-
-    if (atomic_fetch_add(n_barrier, 1) == n_threads - 1) {
+    if (n_barrier == (n_threads - 1)) {
         // last thread
-        atomic_store(n_barrier, 0);
-        atomic_fetch_add_explicit(n_barrier_passed, 1, memory_order_relaxed);
-    } else {
-        // wait for other threads
-        while (true) {
-            if (atomic_load_explicit(n_barrier_passed, memory_order_relaxed) != passed_old) {
-                return;
-            }
-            lm_ggml_thread_cpu_relax();
-        }
+        atomic_store_explicit(&tp->n_barrier, 0, memory_order_relaxed);
+
+        // exit barrier (fill seq-cst fence)
+        atomic_fetch_add_explicit(&tp->n_barrier_passed, 1, memory_order_seq_cst);
+        return;
     }
-}
+
+    // wait for other threads
+    while (atomic_load_explicit(&tp->n_barrier_passed, memory_order_relaxed) == n_passed) {
+        lm_ggml_thread_cpu_relax();
+    }
+
+    // exit barrier (full seq-cst fence)
+    // TSAN doesn't support standalone fence yet, we use a dummy read-modify-write instead
+    #ifdef LM_GGML_TSAN_ENABLED
+    atomic_fetch_add_explicit(&tp->n_barrier_passed, 0, memory_order_seq_cst);
+    #else
+    atomic_thread_fence(memory_order_seq_cst);
+    #endif
 #endif
+}
 
 // TODO: make this somehow automatically executed
 //       some sort of "sentry" mechanism
@@ -4097,7 +4127,11 @@ static void lm_ggml_set_op_params_f32(struct lm_ggml_tensor * tensor, uint32_t i
 }
 
 struct lm_ggml_tensor * lm_ggml_set_zero(struct lm_ggml_tensor * tensor) {
-    memset(tensor->data, 0, lm_ggml_nbytes(tensor));
+    if (tensor->buffer) {
+        lm_ggml_backend_tensor_memset(tensor, 0, 0, lm_ggml_nbytes(tensor));
+    } else {
+        memset(tensor->data, 0, lm_ggml_nbytes(tensor));
+    }
     return tensor;
 }
 
@@ -8323,11 +8357,46 @@ struct lm_ggml_tensor * lm_ggml_cross_entropy_loss_back(
     return result;
 }
 
-////////////////////////////////////////////////////////////////////////////////
+// opt_step_adamw
 
-void lm_ggml_set_param(
+struct lm_ggml_tensor * lm_ggml_opt_step_adamw(
         struct lm_ggml_context * ctx,
-        struct lm_ggml_tensor * tensor) {
+        struct lm_ggml_tensor  * a,
+        float                 alpha,
+        float                 beta1,
+        float                 beta2,
+        float                 eps,
+        float                 wd) {
+    LM_GGML_ASSERT(a->grad);
+    LM_GGML_ASSERT(alpha >  0.0f);
+    LM_GGML_ASSERT(beta1 >= 0.0f && beta1 <= 1.0f);
+    LM_GGML_ASSERT(beta2 >= 0.0f && beta2 <= 1.0f);
+    LM_GGML_ASSERT(eps   >= 0.0f);
+    LM_GGML_ASSERT(wd    >= 0.0f && wd    <= 1.0f);
+
+    struct lm_ggml_tensor * result = lm_ggml_view_tensor(ctx, a);
+
+    result->op   = LM_GGML_OP_OPT_STEP_ADAMW;
+    result->grad = NULL;
+    result->src[0] = a;
+    result->src[1] = a->grad;
+    result->src[2] = lm_ggml_dup_tensor(ctx, a->grad);
+    result->src[3] = lm_ggml_dup_tensor(ctx, a->grad);
+
+    const int64_t iter = 1;
+    memcpy(&result->op_params[0], &iter, sizeof(int64_t));
+    lm_ggml_set_op_params_f32(result, 2, alpha);
+    lm_ggml_set_op_params_f32(result, 3, beta1);
+    lm_ggml_set_op_params_f32(result, 4, beta2);
+    lm_ggml_set_op_params_f32(result, 5, eps);
+    lm_ggml_set_op_params_f32(result, 6, wd);
+
+    return result;
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+void lm_ggml_set_param(struct lm_ggml_context * ctx, struct lm_ggml_tensor * tensor) {
     tensor->flags |= LM_GGML_TENSOR_FLAG_PARAM;
 
     LM_GGML_ASSERT(tensor->grad == NULL);
@@ -8335,6 +8404,13 @@ void lm_ggml_set_param(
     lm_ggml_format_name(tensor->grad, "%s (grad)", tensor->name);
 }
 
+void lm_ggml_set_loss(struct lm_ggml_tensor * tensor) {
+    LM_GGML_ASSERT(lm_ggml_is_scalar(tensor));
+    LM_GGML_ASSERT(tensor->type == LM_GGML_TYPE_F32);
+    LM_GGML_ASSERT(tensor->grad);
+    tensor->flags |= LM_GGML_TENSOR_FLAG_LOSS;
+}
+
 // lm_ggml_compute_forward_dup
 
 static void lm_ggml_compute_forward_dup_same_cont(
@@ -17409,7 +17485,7 @@ static void lm_ggml_compute_forward_cross_entropy_loss_back_f32(
     const int64_t ir0 = dr*ith;
     const int64_t ir1 = MIN(ir0 + dr, nr);
 
-    float * d   = (float *) opt0->data;
+    const float d_by_nr = ((const float *) opt0->data)[0] / (float) nr;
 
     for (int64_t i1 = ir0; i1 < ir1; i1++) {
         float * ds0 = (float *)((char *) dst->data  + i1*dst->nb[1]);
@@ -17433,7 +17509,7 @@ static void lm_ggml_compute_forward_cross_entropy_loss_back_f32(
 
         // grad(src0) = (softmax(src0) - src1) * grad(cross_entropy_loss(src0, src1)) / nr
         lm_ggml_vec_sub_f32(nc, ds0, ds0, s1);
-        lm_ggml_vec_scale_f32(nc, ds0, d[0] / (float) nr);
+        lm_ggml_vec_scale_f32(nc, ds0, d_by_nr);
 
 #ifndef NDEBUG
         for (int i = 0; i < nc; ++i) {
@@ -17462,6 +17538,94 @@ static void lm_ggml_compute_forward_cross_entropy_loss_back(
     }
 }
 
+static void lm_ggml_compute_forward_opt_step_adamw_f32(
+        const struct lm_ggml_compute_params * params,
+        struct lm_ggml_tensor * dst) {
+
+    const struct lm_ggml_tensor * src0        = dst->src[0];
+    const struct lm_ggml_tensor * src0_grad   = dst->src[1];
+    const struct lm_ggml_tensor * src0_grad_m = dst->src[2];
+    const struct lm_ggml_tensor * src0_grad_v = dst->src[3];
+    LM_GGML_ASSERT(lm_ggml_are_same_shape(src0, src0_grad));
+
+    const int ith = params->ith;
+    const int nth = params->nth;
+
+    const int nr  = lm_ggml_nrows(src0);
+
+    LM_GGML_TENSOR_UNARY_OP_LOCALS
+    LM_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);
+
+    /* const float   gnorm = 1.0f; */
+    int64_t       iter;   memcpy(&iter, &dst->op_params[0], sizeof(int64_t));
+    const float   alpha = lm_ggml_get_op_params_f32(dst, 2);
+    const float   beta1 = lm_ggml_get_op_params_f32(dst, 3);
+    const float   beta2 = lm_ggml_get_op_params_f32(dst, 4);
+    const float   eps   = lm_ggml_get_op_params_f32(dst, 5);
+    const float   wd    = lm_ggml_get_op_params_f32(dst, 6);
+
+    const float beta1h  = alpha/(1.0f - powf(beta1, iter));
+    const float beta2h  =  1.0f/(1.0f - powf(beta2, iter));
+
+    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
+        float       * m = (float       *) ((char       *) src0_grad_m->data + offset);
+        float       * v = (float       *) ((char       *) src0_grad_v->data + offset);
+
+        for (int i00 = 0; i00 < ne00; ++i00) {
+            m[i00] = m[i00]*beta1 +        g[i00]*(1.0f - beta1);
+            v[i00] = v[i00]*beta2 + g[i00]*g[i00]*(1.0f - beta2);
+
+            const float mh =       m[i00]*beta1h;
+            const float vh = sqrtf(v[i00]*beta2h) + eps;
+
+            // 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) - mh/vh;
+        }
+    }
+
+    lm_ggml_barrier(params->threadpool);
+    if (ith != 0) {
+        return;
+    }
+
+    iter++;
+    memcpy(&dst->op_params[0], &iter, sizeof(int64_t));
+}
+
+static void lm_ggml_compute_forward_opt_step_adamw(
+        const struct lm_ggml_compute_params * params,
+        struct lm_ggml_tensor * dst) {
+
+    const struct lm_ggml_tensor * src0 = dst->src[0];
+
+    switch (src0->type) {
+        case LM_GGML_TYPE_F32:
+            {
+                lm_ggml_compute_forward_opt_step_adamw_f32(params, dst);
+            } break;
+        default:
+            {
+                LM_GGML_ABORT("fatal error");
+            }
+    }
+}
 /////////////////////////////////
 
 static void lm_ggml_compute_forward(struct lm_ggml_compute_params * params, struct lm_ggml_tensor * tensor) {
@@ -17807,6 +17971,11 @@ static void lm_ggml_compute_forward(struct lm_ggml_compute_params * params, stru
                 lm_ggml_compute_forward_cross_entropy_loss_back(params, tensor);
             }
             break;
+        case LM_GGML_OP_OPT_STEP_ADAMW:
+            {
+                lm_ggml_compute_forward_opt_step_adamw(params, tensor);
+            }
+            break;
         case LM_GGML_OP_NONE:
             {
                 // nop
@@ -17961,7 +18130,7 @@ void lm_ggml_build_backward_gradient_checkpointing(
         struct lm_ggml_tensor  * * checkpoints,
         int                     n_checkpoints) {
     lm_ggml_graph_cpy(gf, gb_tmp);
-    lm_ggml_build_backward_expand(ctx, gf, gb_tmp, true);
+    lm_ggml_build_backward_expand(ctx, gf, gb_tmp, false, true);
 
     if (n_checkpoints <= 0) {
         lm_ggml_graph_cpy(gb_tmp, gb);
@@ -17999,42 +18168,93 @@ void lm_ggml_build_backward_gradient_checkpointing(
     lm_ggml_hash_map_free(replacements);
 }
 
-// functions to change gradients considering the case that input a might be initial gradient with zero value
-
-static struct lm_ggml_tensor * lm_ggml_add_or_set(struct lm_ggml_context * ctx, struct lm_ggml_tensor * a, struct lm_ggml_tensor * b, struct lm_ggml_hash_set * zero_table) {
+// utility functions to change gradients
+// if a is in acc_table, modify gradients in-place and mark result as gradient accumulator
+// else if a is in zero_table, replace a
+// else, just add/subtract/etc. the gradients
+
+static struct lm_ggml_tensor * lm_ggml_add_or_set(
+        struct lm_ggml_context  * ctx,
+        struct lm_ggml_tensor   * a,
+        struct lm_ggml_tensor   * b,
+        struct lm_ggml_hash_set * zero_table,
+        struct lm_ggml_hash_set * acc_table) {
+    if (lm_ggml_hash_contains(acc_table, a)) {
+        struct lm_ggml_tensor * ret = lm_ggml_add_impl(ctx, a, b, true);
+        const size_t insert_result = lm_ggml_hash_insert(acc_table, ret);
+        LM_GGML_ASSERT(insert_result != LM_GGML_HASHSET_FULL);
+        LM_GGML_ASSERT(insert_result != LM_GGML_HASHSET_ALREADY_EXISTS);
+        return ret;
+    }
     if (lm_ggml_hash_contains(zero_table, a)) {
         return b;
-    } else {
-        return lm_ggml_add_impl(ctx, a, b, false);
     }
+    return lm_ggml_add_impl(ctx, a, b, false);
 }
 
-static struct lm_ggml_tensor * lm_ggml_acc_or_set(struct lm_ggml_context * ctx, struct lm_ggml_tensor * a, struct lm_ggml_tensor * b, size_t nb1, size_t nb2, size_t nb3, size_t offset, struct lm_ggml_hash_set * zero_table) {
+static struct lm_ggml_tensor * lm_ggml_acc_or_set(
+        struct lm_ggml_context  * ctx,
+        struct lm_ggml_tensor   * a,
+        struct lm_ggml_tensor   * b,
+        const  size_t          nb1,
+        const  size_t          nb2,
+        const  size_t          nb3,
+        const  size_t          offset,
+        struct lm_ggml_hash_set * zero_table,
+        struct lm_ggml_hash_set * acc_table) {
+    if (lm_ggml_hash_contains(acc_table, a)) {
+        struct lm_ggml_tensor * ret = lm_ggml_acc_impl(ctx, a, b, nb1, nb2, nb3, offset, true);
+        const size_t insert_result = lm_ggml_hash_insert(acc_table, ret);
+        LM_GGML_ASSERT(insert_result != LM_GGML_HASHSET_FULL);
+        LM_GGML_ASSERT(insert_result != LM_GGML_HASHSET_ALREADY_EXISTS);
+        return ret;
+    }
     if (lm_ggml_hash_contains(zero_table, a)) {
-        struct lm_ggml_tensor * a_zero = lm_ggml_scale(ctx, a, 0.0f);
+        struct lm_ggml_tensor * a_zero = lm_ggml_scale(ctx, a, 0.0f); // FIXME this is going to produce NaN if a contains inf/NaN
         return lm_ggml_acc_impl(ctx, a_zero, b, nb1, nb2, nb3, offset, false);
-    } else {
-        return lm_ggml_acc_impl(ctx, a, b, nb1, nb2, nb3, offset, false);
     }
+    return lm_ggml_acc_impl(ctx, a, b, nb1, nb2, nb3, offset, false);
 }
 
-static struct lm_ggml_tensor * lm_ggml_add1_or_set(struct lm_ggml_context * ctx, struct lm_ggml_tensor * a, struct lm_ggml_tensor * b, struct lm_ggml_hash_set * zero_table) {
+static struct lm_ggml_tensor * lm_ggml_add1_or_set(
+        struct lm_ggml_context  * ctx,
+        struct lm_ggml_tensor   * a,
+        struct lm_ggml_tensor   * b,
+        struct lm_ggml_hash_set * zero_table,
+        struct lm_ggml_hash_set * acc_table) {
+    if (lm_ggml_hash_contains(acc_table, a)) {
+        struct lm_ggml_tensor * ret = lm_ggml_add1_impl(ctx, a, b, true);
+        const size_t insert_result = lm_ggml_hash_insert(acc_table, ret);
+        LM_GGML_ASSERT(insert_result != LM_GGML_HASHSET_FULL);
+        LM_GGML_ASSERT(insert_result != LM_GGML_HASHSET_ALREADY_EXISTS);
+        return ret;
+    }
     if (lm_ggml_hash_contains(zero_table, a)) {
         return lm_ggml_repeat(ctx, b, a);
-    } else {
-        return lm_ggml_add1_impl(ctx, a, b, false);
     }
+    return lm_ggml_add1_impl(ctx, a, b, false);
 }
 
-static struct lm_ggml_tensor * lm_ggml_sub_or_set(struct lm_ggml_context * ctx, struct lm_ggml_tensor * a, struct lm_ggml_tensor * b, struct lm_ggml_hash_set * zero_table) {
+static struct lm_ggml_tensor * lm_ggml_sub_or_set(
+        struct lm_ggml_context  * ctx,
+        struct lm_ggml_tensor   * a,
+        struct lm_ggml_tensor   * b,
+        struct lm_ggml_hash_set * zero_table,
+        struct lm_ggml_hash_set * acc_table) {
+    if (lm_ggml_hash_contains(acc_table, a)) {
+        struct lm_ggml_tensor * ret = lm_ggml_sub_impl(ctx, a, b, true);
+        const size_t insert_result = lm_ggml_hash_insert(acc_table, ret);
+        LM_GGML_ASSERT(insert_result != LM_GGML_HASHSET_FULL);
+        LM_GGML_ASSERT(insert_result != LM_GGML_HASHSET_ALREADY_EXISTS);
+        return ret;
+    }
     if (lm_ggml_hash_contains(zero_table, a)) {
         return lm_ggml_neg(ctx, b);
-    } else {
-        return lm_ggml_sub_impl(ctx, a, b, false);
     }
+    return lm_ggml_sub_impl(ctx, a, b, false);
 }
 
-static void lm_ggml_compute_backward(struct lm_ggml_context * ctx, struct lm_ggml_tensor * tensor, struct lm_ggml_hash_set * zero_table) {
+static void lm_ggml_compute_backward(struct lm_ggml_context * ctx, struct lm_ggml_tensor * tensor, struct lm_ggml_hash_set * zero_table, struct lm_ggml_hash_set * acc_table) {
     struct lm_ggml_tensor * src0 = tensor->src[0];
     struct lm_ggml_tensor * src1 = tensor->src[1];
     struct lm_ggml_tensor * src2 = tensor->src[2];
@@ -18043,38 +18263,38 @@ static void lm_ggml_compute_backward(struct lm_ggml_context * ctx, struct lm_ggm
         case LM_GGML_OP_DUP:
             {
                 if (src0->grad) {
-                    src0->grad = lm_ggml_add_or_set(ctx, src0->grad, tensor->grad, zero_table);
+                    src0->grad = lm_ggml_add_or_set(ctx, src0->grad, tensor->grad, zero_table, acc_table);
                 }
             } break;
         case LM_GGML_OP_ADD:
             {
                 if (src0->grad) {
-                    src0->grad = lm_ggml_add_or_set(ctx, src0->grad, tensor->grad, zero_table);
+                    src0->grad = lm_ggml_add_or_set(ctx, src0->grad, tensor->grad, zero_table, acc_table);
                 }
                 if (src1->grad) {
                     if (lm_ggml_are_same_shape(src0, src1)) {
-                        src1->grad = lm_ggml_add_or_set(ctx, src1->grad,                       tensor->grad,        zero_table);
+                        src1->grad = lm_ggml_add_or_set(ctx, src1->grad,                       tensor->grad,        zero_table, acc_table);
                     } else {
-                        src1->grad = lm_ggml_add_or_set(ctx, src1->grad, lm_ggml_repeat_back(ctx, tensor->grad, src1), zero_table);
+                        src1->grad = lm_ggml_add_or_set(ctx, src1->grad, lm_ggml_repeat_back(ctx, tensor->grad, src1), zero_table, acc_table);
                     }
                 }
             } break;
         case LM_GGML_OP_ADD1:
             {
                 if (src0->grad) {
-                    src0->grad = lm_ggml_add_or_set(ctx, src0->grad, tensor->grad, zero_table);
+                    src0->grad = lm_ggml_add_or_set(ctx, src0->grad, tensor->grad, zero_table, acc_table);
                 }
                 if (src1->grad) {
                     src1->grad = lm_ggml_add_or_set(ctx,
                         src1->grad,
                         lm_ggml_mean(ctx, tensor->grad), // TODO: should probably be sum instead of mean
-                        zero_table);
+                        zero_table, acc_table);
                 }
             } break;
         case LM_GGML_OP_ACC:
             {
                 if (src0->grad) {
-                    src0->grad = lm_ggml_add_or_set(ctx, src0->grad, tensor->grad, zero_table);
+                    src0->grad = lm_ggml_add_or_set(ctx, src0->grad, tensor->grad, zero_table, acc_table);
                 }
                 if (src1->grad) {
                     const size_t nb1     = ((int32_t *) tensor->op_params)[0];
@@ -18096,16 +18316,16 @@ static void lm_ggml_compute_backward(struct lm_ggml_context * ctx, struct lm_ggm
                             lm_ggml_reshape(ctx,
                                 lm_ggml_cont(ctx, tensor_grad_view),
                                 src1->grad),
-                            zero_table);
+                            zero_table, acc_table);
                 }
             } break;
         case LM_GGML_OP_SUB:
             {
                 if (src0->grad) {
-                    src0->grad = lm_ggml_add_or_set(ctx, src0->grad, tensor->grad, zero_table);
+                    src0->grad = lm_ggml_add_or_set(ctx, src0->grad, tensor->grad, zero_table, acc_table);
                 }
                 if (src1->grad) {
-                    src1->grad = lm_ggml_sub_or_set(ctx, src1->grad, tensor->grad, zero_table);
+                    src1->grad = lm_ggml_sub_or_set(ctx, src1->grad, tensor->grad, zero_table, acc_table);
                 }
             } break;
         case LM_GGML_OP_MUL:
@@ -18115,14 +18335,14 @@ static void lm_ggml_compute_backward(struct lm_ggml_context * ctx, struct lm_ggm
                         lm_ggml_add_or_set(ctx,
                                 src0->grad,
                                 lm_ggml_mul(ctx, src1, tensor->grad),
-                                zero_table);
+                                zero_table, acc_table);
                 }
                 if (src1->grad) {
                     src1->grad =
                         lm_ggml_add_or_set(ctx,
                                 src1->grad,
                                 lm_ggml_mul(ctx, src0, tensor->grad),
-                                zero_table);
+                                zero_table, acc_table);
                 }
             } break;
         case LM_GGML_OP_DIV:
@@ -18132,7 +18352,7 @@ static void lm_ggml_compute_backward(struct lm_ggml_context * ctx, struct lm_ggm
                         lm_ggml_add_or_set(ctx,
                                 src0->grad,
                                 lm_ggml_div(ctx, tensor->grad, src1),
-                                zero_table);
+                                zero_table, acc_table);
                 }
                 if (src1->grad) {
                     src1->grad =
@@ -18141,7 +18361,7 @@ static void lm_ggml_compute_backward(struct lm_ggml_context * ctx, struct lm_ggm
                                 lm_ggml_mul(ctx,
                                     tensor->grad,
                                     lm_ggml_div(ctx, tensor, src1)),
-                                zero_table);
+                                zero_table, acc_table);
                 }
             } break;
         case LM_GGML_OP_SQR:
@@ -18153,7 +18373,7 @@ static void lm_ggml_compute_backward(struct lm_ggml_context * ctx, struct lm_ggm
                                 lm_ggml_scale(ctx,
                                     lm_ggml_mul(ctx, src0, tensor->grad),
                                     2.0f),
-                                zero_table);
+                                zero_table, acc_table);
                 }
             } break;
         case LM_GGML_OP_SQRT:
@@ -18167,7 +18387,7 @@ static void lm_ggml_compute_backward(struct lm_ggml_context * ctx, struct lm_ggm
                                         tensor->grad,
                                         tensor),
                                     0.5f),
-                                zero_table);
+                                zero_table, acc_table);
                 }
             } break;
         case LM_GGML_OP_LOG:
@@ -18179,7 +18399,7 @@ static void lm_ggml_compute_backward(struct lm_ggml_context * ctx, struct lm_ggm
                                 lm_ggml_div(ctx,
                                     tensor->grad,
                                     src0),
-                                zero_table);
+                                zero_table, acc_table);
                 }
             } break;
         case LM_GGML_OP_SIN:
@@ -18191,7 +18411,7 @@ static void lm_ggml_compute_backward(struct lm_ggml_context * ctx, struct lm_ggm
                                 lm_ggml_mul(ctx,
                                     tensor->grad,
                                     lm_ggml_cos(ctx, src0)),
-                                zero_table);
+                                zero_table, acc_table);
                 }
             } break;
         case LM_GGML_OP_COS:
@@ -18203,7 +18423,7 @@ static void lm_ggml_compute_backward(struct lm_ggml_context * ctx, struct lm_ggm
                                 lm_ggml_mul(ctx,
                                     tensor->grad,
                                     lm_ggml_sin(ctx, src0)),
-                                zero_table);
+                                zero_table, acc_table);
                 }
             } break;
         case LM_GGML_OP_SUM:
@@ -18213,7 +18433,7 @@ static void lm_ggml_compute_backward(struct lm_ggml_context * ctx, struct lm_ggm
                         lm_ggml_add1_or_set(ctx,
                                 src0->grad,
                                 tensor->grad,
-                                zero_table);
+                                zero_table, acc_table);
                 }
             } break;
         case LM_GGML_OP_SUM_ROWS:
@@ -18225,7 +18445,7 @@ static void lm_ggml_compute_backward(struct lm_ggml_context * ctx, struct lm_ggm
                                 lm_ggml_repeat(ctx,
                                     tensor->grad,
                                     src0->grad),
-                                zero_table);
+                                zero_table, acc_table);
                 }
             } break;
         case LM_GGML_OP_MEAN:
@@ -18240,7 +18460,7 @@ static void lm_ggml_compute_backward(struct lm_ggml_context * ctx, struct lm_ggm
                     src0->grad = lm_ggml_add_or_set(ctx,
                             src0->grad,
                             lm_ggml_repeat_back(ctx, tensor->grad, src0->grad),
-                            zero_table);
+                            zero_table, acc_table);
                 }
             } break;
         case LM_GGML_OP_REPEAT_BACK:
@@ -18250,7 +18470,7 @@ static void lm_ggml_compute_backward(struct lm_ggml_context * ctx, struct lm_ggm
                     src0->grad = lm_ggml_add_or_set(ctx,
                             src0->grad,
                             lm_ggml_repeat(ctx, tensor->grad, src0->grad),
-                            zero_table);
+                            zero_table, acc_table);
                 }
             } break;
         case LM_GGML_OP_CONCAT:
@@ -18275,7 +18495,7 @@ static void lm_ggml_compute_backward(struct lm_ggml_context * ctx, struct lm_ggm
                     src0->grad = lm_ggml_add_or_set(ctx,
                             src0->grad,
                             lm_ggml_rms_norm_back(ctx, src0, tensor->grad, eps),
-                            zero_table);
+                            zero_table, acc_table);
                 }
             } break;
         case LM_GGML_OP_RMS_NORM_BACK:
@@ -18323,7 +18543,7 @@ static void lm_ggml_compute_backward(struct lm_ggml_context * ctx, struct lm_ggm
                         lm_ggml_add_or_set(ctx,
                                 src0->grad, // [n,m,q1,r1]
                                 s1_tg,      // [n,m,q1,r1]
-                                zero_table);
+                                zero_table, acc_table);
                 }
                 if (src1->grad) {
                     src1->grad =
@@ -18341,7 +18561,7 @@ static void lm_ggml_compute_backward(struct lm_ggml_context * ctx, struct lm_ggm
                                     src0,                           // [n,m,q1,r1]
                                     lm_ggml_transpose(ctx,             // [p,m,qq,rr]
                                         tensor->grad)),             // [m,p,qq,rr]
-                                zero_table);
+                                zero_table, acc_table);
                 }
             } break;
         case LM_GGML_OP_MUL_MAT_ID:
@@ -18363,7 +18583,7 @@ static void lm_ggml_compute_backward(struct lm_ggml_context * ctx, struct lm_ggm
                         lm_ggml_add_or_set(ctx,
                             src0->grad,
                             lm_ggml_scale_impl(ctx, tensor->grad, s, false),
-                            zero_table);
+                            zero_table, acc_table);
                 }
             } break;
         case LM_GGML_OP_SET:
@@ -18392,7 +18612,7 @@ static void lm_ggml_compute_backward(struct lm_ggml_context * ctx, struct lm_ggm
                             tensor->grad,
                             lm_ggml_neg(ctx, tensor_grad_view),
                             nb1, nb2, nb3, offset, false),
-                        zero_table);
+                        zero_table, acc_table);
                 }
 
                 if (src1->grad) {
@@ -18402,7 +18622,7 @@ static void lm_ggml_compute_backward(struct lm_ggml_context * ctx, struct lm_ggm
                             lm_ggml_reshape(ctx,
                                 lm_ggml_cont(ctx, tensor_grad_view),
                                 src1->grad),
-                            zero_table);
+                            zero_table, acc_table);
                 }
             } break;
         case LM_GGML_OP_CPY:
@@ -18413,7 +18633,7 @@ static void lm_ggml_compute_backward(struct lm_ggml_context * ctx, struct lm_ggm
                 // tensor = src0 * 1 + src1 * 0
                 if (src0->grad) {
                     // dsrc0 = dtensor * 1
-                    src0->grad = lm_ggml_add_or_set(ctx, src0->grad, tensor->grad, zero_table);
+                    src0->grad = lm_ggml_add_or_set(ctx, src0->grad, tensor->grad, zero_table, acc_table);
                 }
                 if (src1->grad) {
                     // dsrc1 = dtensor * 0 -> noop
@@ -18425,7 +18645,7 @@ static void lm_ggml_compute_backward(struct lm_ggml_context * ctx, struct lm_ggm
                 if (src0->grad) {
                     LM_GGML_ASSERT(lm_ggml_is_contiguous(src0->grad));
                     LM_GGML_ASSERT(lm_ggml_is_contiguous(tensor->grad));
-                    src0->grad = lm_ggml_add_or_set(ctx, src0->grad, tensor->grad, zero_table);
+                    src0->grad = lm_ggml_add_or_set(ctx, src0->grad, tensor->grad, zero_table, acc_table);
                 }
             } break;
         case LM_GGML_OP_RESHAPE:
@@ -18439,7 +18659,7 @@ static void lm_ggml_compute_backward(struct lm_ggml_context * ctx, struct lm_ggm
                                     ? tensor->grad
                                     : lm_ggml_cont(ctx, tensor->grad),
                                 src0->grad),
-                        zero_table);
+                        zero_table, acc_table);
                 }
             } break;
         case LM_GGML_OP_VIEW:
@@ -18468,7 +18688,7 @@ static void lm_ggml_compute_backward(struct lm_ggml_context * ctx, struct lm_ggm
                         nb3 = (nb3 / n0) * ng;
                     }
 
-                    src0->grad = lm_ggml_acc_or_set(ctx, src0->grad, tensor->grad, nb1, nb2, nb3, offset, zero_table);
+                    src0->grad = lm_ggml_acc_or_set(ctx, src0->grad, tensor->grad, nb1, nb2, nb3, offset, zero_table, acc_table);
                 }
             } break;
         case LM_GGML_OP_PERMUTE:
@@ -18493,7 +18713,7 @@ static void lm_ggml_compute_backward(struct lm_ggml_context * ctx, struct lm_ggm
                                 axes_backward[1],
                                 axes_backward[2],
                                 axes_backward[3]),
-                            zero_table);
+                            zero_table, acc_table);
                 }
             } break;
         case LM_GGML_OP_TRANSPOSE:
@@ -18503,7 +18723,7 @@ static void lm_ggml_compute_backward(struct lm_ggml_context * ctx, struct lm_ggm
                     src0->grad =
                         lm_ggml_add_or_set(ctx, src0->grad,
                             lm_ggml_transpose(ctx, tensor->grad),
-                        zero_table);
+                        zero_table, acc_table);
                 }
             } break;
         case LM_GGML_OP_GET_ROWS:
@@ -18515,7 +18735,7 @@ static void lm_ggml_compute_backward(struct lm_ggml_context * ctx, struct lm_ggm
                             // last lm_ggml_get_rows_back argument src0->grad is only
                             // necessary to setup correct output shape
                             lm_ggml_get_rows_back(ctx, tensor->grad, src1, src0->grad),
-                        zero_table);
+                        zero_table, acc_table);
                 }
                 if (src1->grad) {
                     // noop
@@ -18539,7 +18759,7 @@ static void lm_ggml_compute_backward(struct lm_ggml_context * ctx, struct lm_ggm
                             /* lm_ggml_diag_mask_inf_impl() shouldn't be here */
                             /* ref:  https://github.com/ggerganov/llama.cpp/pull/4203#discussion_r1412377992 */
                             lm_ggml_diag_mask_zero_impl(ctx, tensor->grad, n_past, false),
-                        zero_table);
+                        zero_table, acc_table);
                 }
             } break;
         case LM_GGML_OP_DIAG_MASK_ZERO:
@@ -18550,7 +18770,7 @@ static void lm_ggml_compute_backward(struct lm_ggml_context * ctx, struct lm_ggm
                     src0->grad =
                         lm_ggml_add_or_set(ctx, src0->grad,
                             lm_ggml_diag_mask_zero_impl(ctx, tensor->grad, n_past, false),
-                        zero_table);
+                        zero_table, acc_table);
                 }
             } break;
         case LM_GGML_OP_SOFT_MAX:
@@ -18560,7 +18780,7 @@ static void lm_ggml_compute_backward(struct lm_ggml_context * ctx, struct lm_ggm
                     src0->grad =
                         lm_ggml_add_or_set(ctx, src0->grad,
                             lm_ggml_soft_max_back(ctx, tensor->grad, tensor),
-                        zero_table);
+                        zero_table, acc_table);
                 }
 
             } break;
@@ -18601,7 +18821,7 @@ static void lm_ggml_compute_backward(struct lm_ggml_context * ctx, struct lm_ggm
                                 attn_factor,
                                 beta_fast,
                                 beta_slow),
-                            zero_table);
+                            zero_table, acc_table);
                 }
             } break;
         case LM_GGML_OP_ROPE_BACK:
@@ -18637,7 +18857,7 @@ static void lm_ggml_compute_backward(struct lm_ggml_context * ctx, struct lm_ggm
                                 beta_fast,
                                 beta_slow,
                                 false),
-                            zero_table);
+                            zero_table, acc_table);
                 }
             } break;
         case LM_GGML_OP_CLAMP:
@@ -18662,7 +18882,7 @@ static void lm_ggml_compute_backward(struct lm_ggml_context * ctx, struct lm_ggm
                     src1->grad = lm_ggml_add_or_set(ctx,
                             src1->grad,
                             lm_ggml_im2col_back(ctx, src0, tensor->grad, src1->ne, s0, s1, p0, p1, d0, d1, is_2D),
-                            zero_table);
+                            zero_table, acc_table);
                 }
             } break;
         case LM_GGML_OP_IM2COL_BACK:
@@ -18691,7 +18911,7 @@ static void lm_ggml_compute_backward(struct lm_ggml_context * ctx, struct lm_ggm
                     src0->grad = lm_ggml_add_or_set(ctx,
                             src0->grad,
                             lm_ggml_pool_2d_back(ctx, tensor->grad, src0, op, k0, k1, s0, s1, p0, p1),
-                            zero_table);
+                            zero_table, acc_table);
                 }
             } break;
         case LM_GGML_OP_POOL_2D_BACK:
@@ -18756,7 +18976,7 @@ static void lm_ggml_compute_backward(struct lm_ggml_context * ctx, struct lm_ggm
                     src0->grad = lm_ggml_add_or_set(ctx,
                             src0->grad,
                             grad_q,
-                            zero_table);
+                            zero_table, acc_table);
                 }
                 if (src1->grad) {
                     struct lm_ggml_tensor * view_k = lm_ggml_view_1d(ctx, flash_grad, elem_k, offs_k);
@@ -18764,7 +18984,7 @@ static void lm_ggml_compute_backward(struct lm_ggml_context * ctx, struct lm_ggm
                     src1->grad = lm_ggml_add_or_set(ctx,
                             src1->grad,
                             grad_k,
-                            zero_table);
+                            zero_table, acc_table);
                 }
                 if (src2->grad) {
                     struct lm_ggml_tensor * view_v = lm_ggml_view_1d(ctx, flash_grad, elem_v, offs_v);
@@ -18772,7 +18992,7 @@ static void lm_ggml_compute_backward(struct lm_ggml_context * ctx, struct lm_ggm
                     src2->grad = lm_ggml_add_or_set(ctx,
                             src2->grad,
                             grad_v,
-                            zero_table);
+                            zero_table, acc_table);
                 }
             } break;
         case LM_GGML_OP_FLASH_ATTN_BACK:
@@ -18798,7 +19018,7 @@ static void lm_ggml_compute_backward(struct lm_ggml_context * ctx, struct lm_ggm
                                             lm_ggml_mul(ctx,
                                                 lm_ggml_sgn(ctx, src0),
                                                 tensor->grad),
-                                            zero_table);
+                                            zero_table, acc_table);
                             }
                         } break;
                     case LM_GGML_UNARY_OP_SGN:
@@ -18810,7 +19030,7 @@ static void lm_ggml_compute_backward(struct lm_ggml_context * ctx, struct lm_ggm
                     case LM_GGML_UNARY_OP_NEG:
                         {
                             if (src0->grad) {
-                                src0->grad = lm_ggml_sub_or_set(ctx, src0->grad, tensor->grad, zero_table);
+                                src0->grad = lm_ggml_sub_or_set(ctx, src0->grad, tensor->grad, zero_table, acc_table);
                             }
                         } break;
                     case LM_GGML_UNARY_OP_STEP:
@@ -18835,7 +19055,7 @@ static void lm_ggml_compute_backward(struct lm_ggml_context * ctx, struct lm_ggm
                                         lm_ggml_mul(ctx,
                                             lm_ggml_step(ctx, src0),
                                             tensor->grad),
-                                        zero_table);
+                                        zero_table, acc_table);
                             }
                         } break;
                     case LM_GGML_UNARY_OP_SIGMOID:
@@ -18857,7 +19077,7 @@ static void lm_ggml_compute_backward(struct lm_ggml_context * ctx, struct lm_ggm
                                 src0->grad = lm_ggml_add_or_set(ctx,
                                         src0->grad,
                                         lm_ggml_silu_back(ctx, src0, tensor->grad),
-                                        zero_table);
+                                        zero_table, acc_table);
                             }
                         } break;
                     case LM_GGML_UNARY_OP_EXP:
@@ -18866,7 +19086,7 @@ static void lm_ggml_compute_backward(struct lm_ggml_context * ctx, struct lm_ggm
                                 src0->grad = lm_ggml_add_or_set(ctx,
                                         src0->grad,
                                         lm_ggml_mul(ctx, tensor, tensor->grad),
-                                        zero_table);
+                                        zero_table, acc_table);
                             }
                         } break;
                     default:
@@ -18896,13 +19116,17 @@ static void lm_ggml_compute_backward(struct lm_ggml_context * ctx, struct lm_ggm
                                     src0,
                                     src1,
                                     tensor->grad),
-                                zero_table);
+                                zero_table, acc_table);
                 }
             } break;
         case LM_GGML_OP_CROSS_ENTROPY_LOSS_BACK:
             {
                 LM_GGML_ABORT("fatal error"); // not supported
             }
+        case LM_GGML_OP_OPT_STEP_ADAMW:
+            {
+                LM_GGML_ABORT("fatal error"); // not supported
+            }
         case LM_GGML_OP_NONE:
             {
                 // nop
@@ -18992,7 +19216,7 @@ void lm_ggml_build_forward_expand(struct lm_ggml_cgraph * cgraph, struct lm_ggml
     lm_ggml_build_forward_impl(cgraph, tensor, true);
 }
 
-void lm_ggml_build_backward_expand(struct lm_ggml_context * ctx, struct lm_ggml_cgraph * gf, struct lm_ggml_cgraph * gb, bool keep) {
+void lm_ggml_build_backward_expand(struct lm_ggml_context * ctx, struct lm_ggml_cgraph * gf, struct lm_ggml_cgraph * gb, bool accumulate, bool keep) {
     LM_GGML_ASSERT(gf->n_nodes > 0);
     LM_GGML_ASSERT(gf->grads);
 
@@ -19008,21 +19232,35 @@ void lm_ggml_build_backward_expand(struct lm_ggml_context * ctx, struct lm_ggml_
         }
     }
 
-    // remember original gradients which start with zero values
+    // keep tables of original gradients for replacement/accumulation logic
     struct lm_ggml_hash_set zero_table = lm_ggml_hash_set_new(gf->size);
+    struct lm_ggml_hash_set acc_table  = lm_ggml_hash_set_new(gf->size);
     for (int i = 0; i < gf->n_nodes; i++) {
-        if (gf->grads[i]) {
-            lm_ggml_hash_insert(&zero_table, gf->grads[i]);
+        struct lm_ggml_tensor * node = gf->nodes[i];
+
+        if (node->grad) {
+            {
+                const size_t insert_result = lm_ggml_hash_insert(&zero_table, node->grad);
+                LM_GGML_ASSERT(insert_result != LM_GGML_HASHSET_FULL);
+                LM_GGML_ASSERT(insert_result != LM_GGML_HASHSET_ALREADY_EXISTS);
+            }
+
+            // only gradients of trainable parameters should be accumulated
+            if (accumulate && (node->flags & LM_GGML_TENSOR_FLAG_PARAM)) {
+                const size_t insert_result = lm_ggml_hash_insert(&acc_table, node->grad);
+                LM_GGML_ASSERT(insert_result != LM_GGML_HASHSET_FULL);
+                LM_GGML_ASSERT(insert_result != LM_GGML_HASHSET_ALREADY_EXISTS);
+            }
         }
     }
 
     for (int i = gf->n_nodes - 1; i >= 0; i--) {
         struct lm_ggml_tensor * node = gf->nodes[i];
 
-        // inplace operations to add gradients are not created by lm_ggml_compute_backward
+        // inplace operations to add gradients are not created by lm_ggml_compute_backward except for gradient accumulation
         // use allocator to automatically make inplace operations
         if (node->grad) {
-            lm_ggml_compute_backward(ctx, node, &zero_table);
+            lm_ggml_compute_backward(ctx, node, &zero_table, &acc_table);
         }
     }
 
@@ -19036,8 +19274,30 @@ void lm_ggml_build_backward_expand(struct lm_ggml_context * ctx, struct lm_ggml_
     }
 
     lm_ggml_hash_set_free(&zero_table);
+    lm_ggml_hash_set_free(&acc_table);
+}
+
+void lm_ggml_build_opt_adamw(
+        struct lm_ggml_context * ctx,
+        struct lm_ggml_cgraph  * gf,
+        struct lm_ggml_cgraph  * gb,
+        float                 alpha,
+        float                 beta1,
+        float                 beta2,
+        float                 eps,
+        float                 wd) {
+    for (int i = 0; i < gf->n_nodes; i++) {
+        struct lm_ggml_tensor * node = gf->nodes[i];
+
+        if (node->flags & LM_GGML_TENSOR_FLAG_PARAM) {
+            LM_GGML_PRINT_DEBUG("%s: found root node %p\n", __func__, (void *) node);
+            struct lm_ggml_tensor * opt_step = lm_ggml_opt_step_adamw(ctx, node, alpha, beta1, beta2, eps, wd);
+            lm_ggml_build_forward_expand(gb, opt_step);
+        }
+    }
 }
 
+
 static void * incr_ptr_aligned(void ** p, size_t size, size_t align) {
     void * ptr = *p;
     ptr = (void *) LM_GGML_PAD((uintptr_t) ptr, align);
@@ -19165,10 +19425,28 @@ void lm_ggml_graph_reset(struct lm_ggml_cgraph * cgraph) {
     LM_GGML_ASSERT(cgraph->grads != NULL);
 
     for (int i = 0; i < cgraph->n_nodes; i++) {
-        struct lm_ggml_tensor * grad = cgraph->grads[i];
+        struct lm_ggml_tensor * node = cgraph->nodes[i];
+
+        // initial gradients of loss should be 1, 0 otherwise
+        if (node->grad) {
+            if (node->flags & LM_GGML_TENSOR_FLAG_LOSS) {
+                LM_GGML_ASSERT(node->grad->buffer);
+                LM_GGML_ASSERT(node->type == LM_GGML_TYPE_F32);
+                LM_GGML_ASSERT(lm_ggml_is_scalar(node));
+
+                const float onef = 1.0f;
+                lm_ggml_backend_tensor_set(node->grad, &onef, 0, lm_ggml_nbytes(node->grad));
+            } else {
+                lm_ggml_set_zero(node->grad);
+            }
+        }
 
-        if (grad) {
-            lm_ggml_set_zero(grad);
+        LM_GGML_ASSERT(node);
+        if (node->op == LM_GGML_OP_OPT_STEP_ADAMW) {
+            // set iteration to 1 and clear momenta
+            lm_ggml_set_op_params_i32(node, 0, 1);
+            lm_ggml_set_zero(node->src[2]);
+            lm_ggml_set_zero(node->src[3]);
         }
     }
 }
@@ -19461,6 +19739,7 @@ static int lm_ggml_get_n_tasks(struct lm_ggml_tensor * node, int n_threads) {
             } break;
         case LM_GGML_OP_CROSS_ENTROPY_LOSS:
         case LM_GGML_OP_CROSS_ENTROPY_LOSS_BACK:
+        case LM_GGML_OP_OPT_STEP_ADAMW:
             {
                 n_tasks = n_threads;
             } break;
@@ -19756,8 +20035,8 @@ void lm_ggml_threadpool_resume(struct lm_ggml_threadpool * threadpool) {
 
 struct lm_ggml_cplan lm_ggml_graph_plan(
           const struct lm_ggml_cgraph * cgraph,
-                           int       n_threads,
-    struct lm_ggml_threadpool * threadpool) {
+                               int   n_threads,
+            struct lm_ggml_threadpool * threadpool) {
 
     if (threadpool == NULL) {
         LM_GGML_PRINT_DEBUG("Threadpool is not specified. Will create a disposable threadpool : n_threads %d\n", n_threads);
@@ -19932,34 +20211,33 @@ struct lm_ggml_cplan lm_ggml_graph_plan(
 
 static thread_ret_t lm_ggml_graph_compute_thread(void * data) {
     struct lm_ggml_compute_state * state = (struct lm_ggml_compute_state *) data;
+    struct lm_ggml_threadpool    * tp    = state->threadpool;
 
-    const struct lm_ggml_cgraph * cgraph = state->threadpool->cgraph;
-    const struct lm_ggml_cplan  * cplan  = state->threadpool->cplan;
+    const struct lm_ggml_cgraph * cgraph = tp->cgraph;
+    const struct lm_ggml_cplan  * cplan  = tp->cplan;
 
     set_numa_thread_affinity(state->ith);
 
     struct lm_ggml_compute_params params = {
         /*.ith       =*/ state->ith,
-        /*.nth       =*/ state->threadpool->n_threads_cur,
+        /*.nth       =*/ atomic_load_explicit(&tp->n_threads_cur, memory_order_relaxed),
         /*.wsize     =*/ cplan->work_size,
         /*.wdata     =*/ cplan->work_data,
-        /*.threadpool=*/ state->threadpool,
+        /*.threadpool=*/ tp,
     };
 
-    for (int node_n = 0; node_n < cgraph->n_nodes; node_n++) {
+    for (int node_n = 0; node_n < cgraph->n_nodes && !tp->abort; node_n++) {
         struct lm_ggml_tensor * node = cgraph->nodes[node_n];
 
         lm_ggml_compute_forward(&params, node);
 
-        if (state->ith == 0 && cplan->abort_callback && cplan->abort_callback(cplan->abort_callback_data)) {
-            state->threadpool->ec = LM_GGML_STATUS_ABORTED;
+        if (state->ith == 0 && cplan->abort_callback &&
+                cplan->abort_callback(cplan->abort_callback_data)) {
+            tp->abort = true;
+            tp->ec    = LM_GGML_STATUS_ABORTED;
         }
 
         lm_ggml_barrier(state->threadpool);
-
-        if (state->threadpool->ec != LM_GGML_STATUS_SUCCESS) {
-            break;
-        }
     }
 
     return 0;
@@ -19967,7 +20245,15 @@ static thread_ret_t lm_ggml_graph_compute_thread(void * data) {
 
 #ifndef LM_GGML_USE_OPENMP
 
-static inline bool lm_ggml_graph_compute_ready(struct lm_ggml_compute_state * state) {
+// check if thread is active
+static inline bool lm_ggml_graph_compute_thread_active(struct lm_ggml_compute_state * state) {
+    struct lm_ggml_threadpool * threadpool = state->threadpool;
+    int n_threads = atomic_load_explicit(&threadpool->n_threads_cur, memory_order_relaxed);
+    return (state->ith < n_threads);
+}
+
+// check if thread is ready to proceed (exit from polling or sleeping)
+static inline bool lm_ggml_graph_compute_thread_ready(struct lm_ggml_compute_state * state) {
     struct lm_ggml_threadpool * threadpool = state->threadpool;
 
     if (state->pending || threadpool->stop || threadpool->pause) { return true; }
@@ -19975,21 +20261,37 @@ static inline bool lm_ggml_graph_compute_ready(struct lm_ggml_compute_state * st
     // check for new graph/work
     int new_graph = atomic_load_explicit(&threadpool->n_graph, memory_order_relaxed);
     if (new_graph != state->last_graph) {
-        state->pending    = (state->ith < threadpool->n_threads_cur);
+        state->pending    = lm_ggml_graph_compute_thread_active(state);
         state->last_graph = new_graph;
     }
 
     return state->pending;
 }
 
+// sync thread state after polling
+static inline void lm_ggml_graph_compute_thread_sync(struct lm_ggml_compute_state * state) {
+    // TSAN doesn't support standalone fence yet, we use a dummy read-modify-write instead
+    #ifdef LM_GGML_TSAN_ENABLED
+    atomic_fetch_add_explicit(&state->threadpool->n_graph, 0, memory_order_seq_cst);
+    #else
+    atomic_thread_fence(memory_order_seq_cst);
+    #endif
+    UNUSED(state);
+}
+
 static inline bool lm_ggml_graph_compute_poll_for_work(struct lm_ggml_compute_state * state) {
     struct lm_ggml_threadpool * threadpool = state->threadpool;
 
+    // Skip polling for unused threads
+    if (!lm_ggml_graph_compute_thread_active(state)) {
+        return state->pending;
+    }
+
     // This seems to make 0 ... 100 a decent range for polling level across modern processors.
     // Perhaps, we can adjust it dynamically based on load and things.
     const uint64_t n_rounds = 1024UL * 128 * threadpool->poll;
 
-    for (uint64_t i=0; !lm_ggml_graph_compute_ready(state) && i<n_rounds; i++) {
+    for (uint64_t i=0; !lm_ggml_graph_compute_thread_ready(state) && i < n_rounds; i++) {
         // No new work. Keep polling.
         lm_ggml_thread_cpu_relax();
     }
@@ -20001,13 +20303,14 @@ static inline bool lm_ggml_graph_compute_check_for_work(struct lm_ggml_compute_s
     struct lm_ggml_threadpool * threadpool = state->threadpool;
 
     if (lm_ggml_graph_compute_poll_for_work(state)) {
+        lm_ggml_graph_compute_thread_sync(state);
         return state->pending;
     }
 
     lm_ggml_mutex_lock_shared(&threadpool->mutex);
-    while (!lm_ggml_graph_compute_ready(state)) {
+    while (!lm_ggml_graph_compute_thread_ready(state)) {
         // No new work. Wait for the signal.
-        LM_GGML_PRINT_DEBUG("thread #%d waiting for work\n", state->ith);
+        LM_GGML_PRINT_DEBUG("thread #%d waiting for work (sleeping)\n", state->ith);
         lm_ggml_cond_wait(&threadpool->cond, &threadpool->mutex);
     }
     lm_ggml_mutex_unlock_shared(&threadpool->mutex);
@@ -20054,13 +20357,20 @@ static thread_ret_t lm_ggml_graph_compute_secondary_thread(void* data) {
 }
 
 // Start processing new graph
-static void lm_ggml_graph_compute_kickoff(struct lm_ggml_threadpool * threadpool)
+static void lm_ggml_graph_compute_kickoff(struct lm_ggml_threadpool * threadpool, int n_threads)
 {
-    // always take the mutex here because the worker threads are doing hybrid poll/wait
+    // Always take the mutex here because the worker threads are doing hybrid poll/wait
 
     lm_ggml_mutex_lock(&threadpool->mutex);
 
-    atomic_fetch_add_explicit(&threadpool->n_graph, 1, memory_order_relaxed);
+    LM_GGML_PRINT_DEBUG("threadpool: n_threads_cur %d n_threads %d\n", threadpool->n_threads_cur, n_threads);
+
+    // Update the number of active threads
+    atomic_store_explicit(&threadpool->n_threads_cur, n_threads, memory_order_relaxed);
+
+    // Indicate the graph is ready to be processed
+    // We need the full seq-cst fence here because of the polling threads (used in thread_sync)
+    atomic_fetch_add_explicit(&threadpool->n_graph, 1, memory_order_seq_cst);
 
     if (threadpool->pause) {
        // Update main thread prio and affinity to match the threadpool settings
@@ -20119,6 +20429,7 @@ static struct lm_ggml_threadpool * lm_ggml_threadpool_new_impl(
         threadpool->current_chunk    = 0;
         threadpool->stop             = false;
         threadpool->pause            = tpp->paused;
+        threadpool->abort            = false;
         threadpool->workers          = NULL;
         threadpool->n_threads_max    = tpp->n_threads;
         threadpool->n_threads_cur    = tpp->n_threads;
@@ -20194,15 +20505,11 @@ enum lm_ggml_status lm_ggml_graph_compute(struct lm_ggml_cgraph * cgraph, struct
         // No worker threads should be accessing the parameters below at this stage
         threadpool->cgraph           = cgraph;
         threadpool->cplan            = cplan;
-        threadpool->n_threads_cur    = n_threads;
         threadpool->current_chunk    = 0;
+        threadpool->abort            = false;
         threadpool->ec               = LM_GGML_STATUS_SUCCESS;
     }
 
-    if (n_threads > threadpool->n_threads_max) {
-        LM_GGML_PRINT("WARNING: cplan is requesting more threads than the threadpool contains. Expect a bad time!\n");
-    }
-
 #ifdef LM_GGML_USE_OPENMP
     if (n_threads > 1) {
         #pragma omp parallel num_threads(n_threads)
@@ -20211,17 +20518,23 @@ enum lm_ggml_status lm_ggml_graph_compute(struct lm_ggml_cgraph * cgraph, struct
             {
                 // update the number of threads from the actual number of threads that we got from OpenMP
                 n_threads = omp_get_num_threads();
-                threadpool->n_threads_cur = n_threads;
+                atomic_store_explicit(&threadpool->n_threads_cur, n_threads, memory_order_relaxed);
             }
 
             lm_ggml_graph_compute_thread(&threadpool->workers[omp_get_thread_num()]);
         }
     } else {
+        atomic_store_explicit(&threadpool->n_threads_cur, 1, memory_order_relaxed);
         lm_ggml_graph_compute_thread(&threadpool->workers[0]);
     }
 #else
+    if (n_threads > threadpool->n_threads_max) {
+        LM_GGML_PRINT("WARNING: cplan requested more threads (%d) than available (%d)\n", n_threads, threadpool->n_threads_max);
+        n_threads = threadpool->n_threads_max;
+    }
+
     // Kick all threads to start the new graph
-    lm_ggml_graph_compute_kickoff(threadpool);
+    lm_ggml_graph_compute_kickoff(threadpool, n_threads);
 
     // This is a work thread too
     lm_ggml_graph_compute_thread(&threadpool->workers[0]);
@@ -21823,7 +22136,7 @@ enum lm_ggml_opt_result lm_ggml_opt_resume(
     lm_ggml_build_forward_expand(gf, f);
 
     struct lm_ggml_cgraph * gb = lm_ggml_graph_dup(ctx, gf);
-    lm_ggml_build_backward_expand(ctx, gf, gb, true);
+    lm_ggml_build_backward_expand(ctx, gf, gb, false, true);
 
     return lm_ggml_opt_resume_g(ctx, opt, f, gf, gb, NULL, NULL);
 }