From bc1c13bb6621a6236d8595cf110ac6ad1af0dd96 Mon Sep 17 00:00:00 2001
From: xaedes <xaedes@gmail.com>
Date: Mon, 1 May 2023 22:22:00 +0200
Subject: [PATCH] train with two examples, creating new tensors each time..

---
 examples/baby-llama/baby-llama.cpp | 152 ++++++++++++++++-------------
 1 file changed, 86 insertions(+), 66 deletions(-)

diff --git a/examples/baby-llama/baby-llama.cpp b/examples/baby-llama/baby-llama.cpp
index c50baf470..0e5b3fbd7 100644
--- a/examples/baby-llama/baby-llama.cpp
+++ b/examples/baby-llama/baby-llama.cpp
@@ -539,7 +539,7 @@ void print_probs(struct ggml_tensor * probs) {
     for (int i=0; i<probs->ne[1]; ++i) {
         for (int k = 0; k < probs->ne[0]; ++k) {
             float p = ggml_get_f32_1d(probs, i*probs->ne[0] + k);
-            printf(" %.1f", p);
+            printf(" %.2f", p);
         }
         printf("\n");
     }
@@ -559,6 +559,21 @@ void print_tokens(struct ggml_tensor * tokens, int n_vocab) {
     }
 }
 
+void get_example_targets(int example_id, struct ggml_tensor * tokens_input, struct ggml_tensor * targets) {
+    int n_tokens = tokens_input->ne[0];
+    int n_vocab = targets->ne[0];
+    ggml_set_zero(targets);
+    for (int i=0; i<n_tokens; ++i) {
+        float x = example_id + i * 3.14159f * 2.0f * 4.0f / n_tokens;
+        float y = sinf(x);//*cosf(x*1.1f+1.0f);
+        float z = (y+1.0f)*0.5f; // scale to [0..1]
+        z = (z < 0.0f) ? 0.0f : (z > 1.0f) ? 1.0f : z; // clamp to [0..1]
+        int token = (int)(z*(float)(n_vocab-1));
+        ggml_set_f32_1d(targets, i*n_vocab + token, +1.0f);
+        ggml_set_i32_1d(tokens_input, i, token);
+    }
+}
+
 int main(int argc, char ** argv) {
     struct ggml_init_params lcparams;
     lcparams.mem_size   = 1024*1024*1024;
@@ -566,19 +581,26 @@ int main(int argc, char ** argv) {
     lcparams.no_alloc   = false;
 
     struct llama_model model;
-    model.hparams.n_vocab = 16;
-    model.hparams.n_ctx   = 64;
-    model.hparams.n_embd  = 64;
+    model.hparams.n_vocab = 8;
+    model.hparams.n_ctx   = 32;
+    model.hparams.n_embd  = 32;
     model.hparams.n_mult  = 2;
     model.hparams.n_head  = 8;
-    model.hparams.n_layer = 16;
+    model.hparams.n_layer = 8;
     model.hparams.n_rot   = 16;
+
+    // model.hparams.n_embd  = 32;
+    // model.hparams.n_mult  = 2;
+    // model.hparams.n_head  = 4;
+    // model.hparams.n_layer = 8;
+    // model.hparams.n_rot   = 8;
+
     model.ctx = ggml_init(lcparams);
     printf("init model\n");
     init_model(&model);
     set_param_model(&model);
 
-    randomize_model(&model, 1337, 0.0f, 2.0f, -1.0f, +1.0f);
+    randomize_model(&model, 1337, 0.0f, 1.0f, -1.0f, +1.0f);
 
     // key + value cache for the self attention
     struct llama_kv_cache kv_self;    
@@ -593,68 +615,66 @@ int main(int argc, char ** argv) {
 
     struct ggml_context * ctx0 = model.ctx; // ggml_init(c0params);
 
-    int n_tokens = 64;
-    struct ggml_tensor * before_opt_best_samples = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
-    struct ggml_tensor * before_opt_probs        = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, model.hparams.n_vocab, n_tokens);
-    struct ggml_tensor * after_opt_best_samples  = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
-    struct ggml_tensor * after_opt_probs         = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, model.hparams.n_vocab, n_tokens);
-    struct ggml_tensor * tokens_input            = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
-    struct ggml_tensor * targets                 = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, model.hparams.n_vocab, n_tokens);
-    for (int i=0; i<n_tokens; ++i) {
-        float x = i * 3.14159f * 2.0f * 4.0f / n_tokens;
-        float y = sinf(x);
-        float z = (y+1.0f)*0.5f;
-        int token = (int)(z*(float)(model.hparams.n_vocab-1));
-        for (int k = 0; k < token; ++k) {
-            ggml_set_f32_1d(targets, i*model.hparams.n_vocab + k, 0.0f);
-        }
-        ggml_set_f32_1d(targets, i*model.hparams.n_vocab + token, +1.0f);
-        for (int k = token+1; k < model.hparams.n_vocab; ++k) {
-            ggml_set_f32_1d(targets, i*model.hparams.n_vocab + k, 0.0f);
-        }
-        ggml_set_i32_1d(tokens_input, i, token);
+    int n_examples = 2;
+    int n_tokens = model.hparams.n_ctx;
+
+    for (int ex=0; ex<n_examples; ++ex) {
+
+        struct ggml_tensor * before_opt_best_samples = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
+        struct ggml_tensor * before_opt_probs        = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, model.hparams.n_vocab, n_tokens);
+        struct ggml_tensor * after_opt_best_samples  = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
+        struct ggml_tensor * after_opt_probs         = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, model.hparams.n_vocab, n_tokens);
+        struct ggml_tensor * tokens_input            = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
+        struct ggml_tensor * targets                 = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, model.hparams.n_vocab, n_tokens);
+
+        int n_past = 0;
+
+        ggml_cgraph gf = {};
+        gf.n_threads = 1;
+
+        get_example_targets(ex, tokens_input, targets);
+        printf("Example %d\n", (ex+1));
+        print_probs(targets);
+        print_tokens(tokens_input, model.hparams.n_vocab);
+
+        struct ggml_tensor * logits = forward(&model, &kv_self, ctx0, &gf, tokens_input, n_tokens, n_past);
+        struct ggml_tensor * e = ggml_sum(ctx0, ggml_sqr(ctx0, ggml_sub(ctx0, targets, logits)));
+
+        ggml_build_forward_expand(&gf, e);
+        ggml_graph_compute(ctx0, &gf);
+
+        float error_before_opt = ggml_get_f32_1d(e, 0);
+        sample_softmax(logits, before_opt_probs, before_opt_best_samples);
+
+        printf("probabilities before optimization:\n");
+        print_probs(before_opt_probs);
+        printf("best samples before optimization:\n");
+        print_tokens(before_opt_best_samples, model.hparams.n_vocab);
+
+        struct ggml_opt_params opt_params_adam = ggml_opt_default_params(GGML_OPT_ADAM);
+        struct ggml_opt_params opt_params_lbfgs = ggml_opt_default_params(GGML_OPT_LBFGS);
+        opt_params_adam.print_forward_graph = false;
+        opt_params_adam.print_backward_graph = false;
+        opt_params_lbfgs.print_forward_graph = false;
+        opt_params_lbfgs.print_backward_graph = false;
+        // ggml_opt(ctx0, opt_params_adam, e);
+        ggml_opt(ctx0, opt_params_lbfgs, e);
+        // 
+        ggml_build_forward_expand(&gf, e);
+        ggml_graph_compute(ctx0, &gf);
+
+        float error_after_opt = ggml_get_f32_1d(e, 0);
+        sample_softmax(logits, after_opt_probs, after_opt_best_samples);
+
+        printf("error_before_opt: %.2f\n", error_before_opt);
+        printf("error_after_opt:  %.2f\n", error_after_opt);
+
+        printf("probabilities after optimization:\n");
+        print_probs(after_opt_probs);
+        printf("best samples after optimization:\n");
+        print_tokens(after_opt_best_samples, model.hparams.n_vocab);
     }
-    print_probs(targets);
-    print_tokens(tokens_input, model.hparams.n_vocab);
 
-    int n_past = 0;
-
-    ggml_cgraph gf = {};
-    gf.n_threads = 1;
-
-    struct ggml_tensor * logits = forward(&model, &kv_self, ctx0, &gf, tokens_input, n_tokens, n_past);
-    
-    struct ggml_tensor * e = ggml_sum(ctx0, ggml_sqr(ctx0, ggml_sub(ctx0, targets, logits)));
-
-    ggml_build_forward_expand(&gf, e);
-    ggml_graph_compute(ctx0, &gf);
-
-    float error_before_opt = ggml_get_f32_1d(e, 0);
-    sample_softmax(logits, before_opt_probs, before_opt_best_samples);
-
-    printf("probabilities before optimization:\n");
-    print_probs(before_opt_probs);
-    printf("best samples before optimization:\n");
-    print_tokens(before_opt_best_samples, model.hparams.n_vocab);
-
-    struct ggml_opt_params opt_params_adam = ggml_opt_default_params(GGML_OPT_ADAM);
-    struct ggml_opt_params opt_params_lbfgs = ggml_opt_default_params(GGML_OPT_LBFGS);
-    ggml_opt(ctx0, opt_params_lbfgs, e);
-    // ggml_opt(ctx0, opt_params_adam, e);
-    // 
-    ggml_build_forward_expand(&gf, e);
-    ggml_graph_compute(ctx0, &gf);
-
-    float error_after_opt = ggml_get_f32_1d(e, 0);
-    sample_softmax(logits, after_opt_probs, after_opt_best_samples);
-
-    printf("error_before_opt: %.2f\n", error_before_opt);
-    printf("error_after_opt:  %.2f\n", error_after_opt);
-
-    printf("probabilities after optimization:\n");
-    print_probs(after_opt_probs);
-    printf("best samples after optimization:\n");
-    print_tokens(after_opt_best_samples, model.hparams.n_vocab);
 
     ggml_free(ctx0);