lora: support tucker decomposition (from loha)

Author: stduhpf

ggml_extend.hpp

@@ -52,6 +52,32 @@
 #define __STATIC_INLINE__ static inline
 #endif
 
+__STATIC_INLINE__ void print_ggml_tensor(struct ggml_tensor* tensor, bool shape_only, const char* mark);
+
+// n-mode trensor-matrix product
+// example: 2-mode product
+// A: [ne03, k, ne01, ne00]
+// B: k rows, m columns => [k, m]
+// result is [ne03, m, ne01, ne00]
+__STATIC_INLINE__ struct ggml_tensor* ggml_mul_n_mode(struct ggml_context* ctx, struct ggml_tensor* a, struct ggml_tensor* b, int mode = 0) {
+    // reshape A
+    // swap 0th and nth axis
+    a = ggml_cont(ctx, ggml_permute(ctx, a, mode, mode != 1 ? 1 : 0, mode != 2 ? 2 : 0, mode != 3 ? 3 : 0));
+    int ne1 = a->ne[1];
+    int ne2 = a->ne[2];
+    int ne3 = a->ne[3];
+    // make 2D
+    a = ggml_cont(ctx, ggml_reshape_2d(ctx, a, a->ne[0], (ne3 * ne2 * ne1)));
+
+    struct ggml_tensor* result = ggml_cont(ctx, ggml_transpose(ctx, ggml_mul_mat(ctx, a, b)));
+
+    // reshape output (same shape as a after permutation except first dim)
+    result = ggml_reshape_4d(ctx, result, result->ne[0], ne1, ne2, ne3);
+    // swap back 0th and nth axis
+    result = ggml_permute(ctx, result, mode, mode != 1 ? 1 : 0, mode != 2 ? 2 : 0, mode != 3 ? 3 : 0);
+    return result;
+}
+
 __STATIC_INLINE__ void ggml_log_callback_default(ggml_log_level level, const char* text, void* user_data) {
     (void)level;
     (void)user_data;
@@ -319,7 +345,7 @@ __STATIC_INLINE__ void sd_apply_mask(struct ggml_tensor* image_data,
         for (int iy = 0; iy < height; iy++) {
             float m = ggml_tensor_get_f32(mask, ix, iy);
             for (int k = 0; k < channels; k++) {
-                float value = ((float)(m < 254.5/255)) * (ggml_tensor_get_f32(image_data, ix, iy, k) - .5) + .5;
+                float value = ((float)(m < 254.5 / 255)) * (ggml_tensor_get_f32(image_data, ix, iy, k) - .5) + .5;
                 ggml_tensor_set_f32(output, value, ix, iy, k);
             }
         }

lora.hpp

@@ -244,12 +244,15 @@ struct LoraModel : public GGMLRunner {
             std::vector<std::string> keys = to_lora_keys(k_tensor, version);
             if (keys.size() == 0)
                 continue;
+
+            ggml_tensor* lora_mid  = NULL;  // tau for tucker decomposition
             ggml_tensor* lora_up   = NULL;
             ggml_tensor* lora_down = NULL;
             for (auto& key : keys) {
                 std::string alpha_name         = "";
                 std::string scale_name         = "";
                 std::string split_q_scale_name = "";
+                std::string lora_mid_name      = "";
                 std::string lora_down_name     = "";
                 std::string lora_up_name       = "";
 
@@ -584,8 +587,10 @@ struct LoraModel : public GGMLRunner {
                     }
 
                     lora_down_name = lora_pre[type] + key + lora_downs[type] + ".weight";
-                    alpha_name     = lora_pre[type] + key + ".alpha";
-                    scale_name     = lora_pre[type] + key + ".scale";
+                    lora_mid_name  = lora_pre[type] + key + ".lora_mid.weight";
+
+                    alpha_name = lora_pre[type] + key + ".alpha";
+                    scale_name = lora_pre[type] + key + ".scale";
 
                     if (lora_tensors.find(lora_up_name) != lora_tensors.end()) {
                         lora_up = lora_tensors[lora_up_name];
@@ -594,6 +599,12 @@ struct LoraModel : public GGMLRunner {
                     if (lora_tensors.find(lora_down_name) != lora_tensors.end()) {
                         lora_down = lora_tensors[lora_down_name];
                     }
+
+                    if (lora_tensors.find(lora_mid_name) != lora_tensors.end()) {
+                        lora_mid = lora_tensors[lora_mid_name];
+                        applied_lora_tensors.insert(lora_mid_name);
+                    }
+
                     applied_lora_tensors.insert(lora_up_name);
                     applied_lora_tensors.insert(lora_down_name);
                     applied_lora_tensors.insert(alpha_name);
@@ -625,9 +636,20 @@ struct LoraModel : public GGMLRunner {
 
                 // ggml_mul_mat requires tensor b transposed
                 lora_down                  = ggml_cont(compute_ctx, ggml_transpose(compute_ctx, lora_down));
-                struct ggml_tensor* updown = ggml_mul_mat(compute_ctx, lora_up, lora_down);
-                updown                     = ggml_cont(compute_ctx, ggml_transpose(compute_ctx, updown));
-                updown                     = ggml_reshape(compute_ctx, updown, weight);
+                struct ggml_tensor* updown = NULL;
+                if (lora_mid == NULL) {
+                    updown = ggml_mul_mat(compute_ctx, lora_up, lora_down);
+                    updown = ggml_cont(compute_ctx, ggml_transpose(compute_ctx, updown));
+                } else {
+                    // undoing tucker decomposition for conv layers.
+                    // lora_mid  has shape (3,    3,   Rank, Rank)
+                    // lora_down has shape (Rank, In,  1,    1)
+                    // lora_up   has shape (Rank, Out, 1,    1)
+                    // conv layer shape is (3,    3,   Out,  In)
+                    updown = ggml_mul_n_mode(compute_ctx, ggml_mul_n_mode(compute_ctx, lora_mid, lora_down, 3), lora_up, 2);
+                    updown = ggml_cont(compute_ctx, updown);
+                }
+                updown = ggml_reshape(compute_ctx, updown, weight);
                 GGML_ASSERT(ggml_nelements(updown) == ggml_nelements(weight));
                 updown = ggml_scale_inplace(compute_ctx, updown, scale_value);
                 ggml_tensor* final_weight;