406 lines
12 KiB
C
406 lines
12 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
#include "../../lib/test.h"
|
|
#include "../lora.h"
|
|
|
|
#include <string.h>
|
|
#include <stdlib.h>
|
|
#include <unistd.h>
|
|
|
|
/* ============================================================= */
|
|
/* Config Tests */
|
|
/* ============================================================= */
|
|
|
|
static void test_config_default(void) {
|
|
fl_config cfg = fl_config_default();
|
|
|
|
ASSERT_INT_EQ(8, cfg.rank);
|
|
ASSERT_FLT_NEAR(16.0f, cfg.alpha, 0.001f);
|
|
ASSERT_FLT_NEAR(0.0f, cfg.dropout, 0.001f);
|
|
ASSERT_FALSE(cfg.fan_in_fan_out);
|
|
}
|
|
|
|
/* ============================================================= */
|
|
/* Layer Creation Tests */
|
|
/* ============================================================= */
|
|
|
|
static void test_layer_create(void) {
|
|
fl_config cfg = fl_config_default();
|
|
fl_layer* layer = fl_layer_new("test.layer", 64, 128, &cfg);
|
|
|
|
ASSERT_NOT_NULL(layer);
|
|
ASSERT_STR_EQ("test.layer", layer->name);
|
|
ASSERT_INT_EQ(64, layer->d_in);
|
|
ASSERT_INT_EQ(128, layer->d_out);
|
|
ASSERT_INT_EQ(8, layer->rank);
|
|
ASSERT_FLT_NEAR(2.0f, layer->scale, 0.001f); /* 16/8 = 2 */
|
|
|
|
/* Check tensor shapes */
|
|
ASSERT_NOT_NULL(layer->lora_A);
|
|
ASSERT_NOT_NULL(layer->lora_B);
|
|
|
|
ASSERT_INT_EQ(2, layer->lora_A->ndim);
|
|
ASSERT_INT_EQ(8, (int)layer->lora_A->shape[0]); /* rank */
|
|
ASSERT_INT_EQ(64, (int)layer->lora_A->shape[1]); /* d_in */
|
|
|
|
ASSERT_INT_EQ(2, layer->lora_B->ndim);
|
|
ASSERT_INT_EQ(128, (int)layer->lora_B->shape[0]); /* d_out */
|
|
ASSERT_INT_EQ(8, (int)layer->lora_B->shape[1]); /* rank */
|
|
|
|
fl_layer_free(layer);
|
|
}
|
|
|
|
static void test_layer_custom_rank(void) {
|
|
fl_config cfg = {
|
|
.rank = 16,
|
|
.alpha = 32.0f,
|
|
.dropout = 0.0f,
|
|
.fan_in_fan_out = false
|
|
};
|
|
|
|
fl_layer* layer = fl_layer_new("custom", 256, 512, &cfg);
|
|
|
|
ASSERT_NOT_NULL(layer);
|
|
ASSERT_INT_EQ(16, layer->rank);
|
|
ASSERT_FLT_NEAR(2.0f, layer->scale, 0.001f); /* 32/16 = 2 */
|
|
|
|
ASSERT_INT_EQ(16, (int)layer->lora_A->shape[0]);
|
|
ASSERT_INT_EQ(256, (int)layer->lora_A->shape[1]);
|
|
ASSERT_INT_EQ(512, (int)layer->lora_B->shape[0]);
|
|
ASSERT_INT_EQ(16, (int)layer->lora_B->shape[1]);
|
|
|
|
fl_layer_free(layer);
|
|
}
|
|
|
|
static void test_layer_from_tensors(void) {
|
|
/* Create A and B matrices */
|
|
int64_t shape_A[] = {4, 32}; /* rank=4, d_in=32 */
|
|
int64_t shape_B[] = {64, 4}; /* d_out=64, rank=4 */
|
|
|
|
tm_tensor* A = tm_ones(shape_A, 2, TM_F32);
|
|
tm_tensor* B = tm_ones(shape_B, 2, TM_F32);
|
|
|
|
fl_layer* layer = fl_layer_from_tensors("from_tensors", A, B, 8.0f);
|
|
|
|
ASSERT_NOT_NULL(layer);
|
|
ASSERT_INT_EQ(32, layer->d_in);
|
|
ASSERT_INT_EQ(64, layer->d_out);
|
|
ASSERT_INT_EQ(4, layer->rank);
|
|
ASSERT_FLT_NEAR(2.0f, layer->scale, 0.001f); /* 8/4 = 2 */
|
|
|
|
tm_release(A);
|
|
tm_release(B);
|
|
fl_layer_free(layer);
|
|
}
|
|
|
|
/* ============================================================= */
|
|
/* Forward Pass Tests */
|
|
/* ============================================================= */
|
|
|
|
static void test_layer_forward_shape(void) {
|
|
fl_config cfg = fl_config_default();
|
|
fl_layer* layer = fl_layer_new("test", 32, 64, &cfg);
|
|
|
|
/* Create input: (batch_size=4, d_in=32) */
|
|
int64_t shape[] = {4, 32};
|
|
tm_tensor* x = tm_ones(shape, 2, TM_F32);
|
|
|
|
tm_tensor* out = fl_layer_forward(layer, x);
|
|
|
|
ASSERT_NOT_NULL(out);
|
|
ASSERT_INT_EQ(2, out->ndim);
|
|
ASSERT_INT_EQ(4, (int)out->shape[0]); /* batch_size */
|
|
ASSERT_INT_EQ(64, (int)out->shape[1]); /* d_out */
|
|
|
|
tm_release(x);
|
|
tm_release(out);
|
|
fl_layer_free(layer);
|
|
}
|
|
|
|
static void test_layer_forward_values(void) {
|
|
/* Create layer with known values */
|
|
int64_t shape_A[] = {2, 4}; /* rank=2, d_in=4 */
|
|
int64_t shape_B[] = {3, 2}; /* d_out=3, rank=2 */
|
|
|
|
tm_tensor* A = tm_ones(shape_A, 2, TM_F32);
|
|
tm_tensor* B = tm_ones(shape_B, 2, TM_F32);
|
|
|
|
fl_layer* layer = fl_layer_from_tensors("known", A, B, 2.0f); /* scale = 2/2 = 1 */
|
|
|
|
/* Input of all ones: (1, 4) */
|
|
int64_t shape_x[] = {1, 4};
|
|
tm_tensor* x = tm_ones(shape_x, 2, TM_F32);
|
|
|
|
tm_tensor* out = fl_layer_forward(layer, x);
|
|
|
|
ASSERT_NOT_NULL(out);
|
|
/* x @ A^T = [1,1,1,1] @ [[1,1,1,1],[1,1,1,1]]^T = [4, 4]
|
|
* [4,4] @ B^T = [4,4] @ [[1,1],[1,1],[1,1]]^T = [8, 8, 8]
|
|
* scale = 1, so output = [8, 8, 8]
|
|
*/
|
|
float* data = tm_data_f32(out);
|
|
ASSERT_FLT_NEAR(8.0f, data[0], 0.001f);
|
|
ASSERT_FLT_NEAR(8.0f, data[1], 0.001f);
|
|
ASSERT_FLT_NEAR(8.0f, data[2], 0.001f);
|
|
|
|
tm_release(A);
|
|
tm_release(B);
|
|
tm_release(x);
|
|
tm_release(out);
|
|
fl_layer_free(layer);
|
|
}
|
|
|
|
/* ============================================================= */
|
|
/* Merge Tests */
|
|
/* ============================================================= */
|
|
|
|
static void test_layer_merge(void) {
|
|
/* Create layer with known values */
|
|
int64_t shape_A[] = {2, 4}; /* rank=2, d_in=4 */
|
|
int64_t shape_B[] = {3, 2}; /* d_out=3, rank=2 */
|
|
|
|
tm_tensor* A = tm_ones(shape_A, 2, TM_F32);
|
|
tm_tensor* B = tm_ones(shape_B, 2, TM_F32);
|
|
|
|
fl_layer* layer = fl_layer_from_tensors("merge", A, B, 2.0f); /* scale = 1 */
|
|
|
|
/* Base weight matrix: (d_out=3, d_in=4) all zeros */
|
|
int64_t shape_W[] = {3, 4};
|
|
tm_tensor* W = tm_zeros(shape_W, 2, TM_F32);
|
|
|
|
tm_tensor* merged = fl_layer_get_merged(layer, W);
|
|
|
|
ASSERT_NOT_NULL(merged);
|
|
ASSERT_INT_EQ(3, (int)merged->shape[0]);
|
|
ASSERT_INT_EQ(4, (int)merged->shape[1]);
|
|
|
|
/* B @ A = [[1,1],[1,1],[1,1]] @ [[1,1,1,1],[1,1,1,1]] = [[2,2,2,2],[2,2,2,2],[2,2,2,2]] */
|
|
/* scale = 1, W = 0, so merged = scale * B @ A = [[2,2,2,2], ...] */
|
|
float* data = tm_data_f32(merged);
|
|
for (int i = 0; i < 12; i++) {
|
|
ASSERT_FLT_NEAR(2.0f, data[i], 0.001f);
|
|
}
|
|
|
|
tm_release(A);
|
|
tm_release(B);
|
|
tm_release(W);
|
|
tm_release(merged);
|
|
fl_layer_free(layer);
|
|
}
|
|
|
|
static void test_layer_merge_adds_to_weight(void) {
|
|
int64_t shape_A[] = {1, 2}; /* rank=1, d_in=2 */
|
|
int64_t shape_B[] = {2, 1}; /* d_out=2, rank=1 */
|
|
|
|
tm_tensor* A = tm_ones(shape_A, 2, TM_F32);
|
|
tm_tensor* B = tm_ones(shape_B, 2, TM_F32);
|
|
|
|
fl_layer* layer = fl_layer_from_tensors("add", A, B, 1.0f); /* scale = 1 */
|
|
|
|
/* Base weight: all 5s */
|
|
int64_t shape_W[] = {2, 2};
|
|
tm_tensor* W = tm_full(shape_W, 2, TM_F32, 5.0);
|
|
|
|
tm_tensor* merged = fl_layer_get_merged(layer, W);
|
|
|
|
/* B @ A = [[1]] @ [[1,1]] = [[1,1],[1,1]]
|
|
* scale = 1, W = 5, merged = W + scale * BA = 5 + 1 = 6
|
|
*/
|
|
float* data = tm_data_f32(merged);
|
|
ASSERT_FLT_NEAR(6.0f, data[0], 0.001f);
|
|
ASSERT_FLT_NEAR(6.0f, data[1], 0.001f);
|
|
ASSERT_FLT_NEAR(6.0f, data[2], 0.001f);
|
|
ASSERT_FLT_NEAR(6.0f, data[3], 0.001f);
|
|
|
|
tm_release(A);
|
|
tm_release(B);
|
|
tm_release(W);
|
|
tm_release(merged);
|
|
fl_layer_free(layer);
|
|
}
|
|
|
|
/* ============================================================= */
|
|
/* Adapter Tests */
|
|
/* ============================================================= */
|
|
|
|
static void test_adapter_create(void) {
|
|
fl_adapter* adapter = fl_adapter_new();
|
|
|
|
ASSERT_NOT_NULL(adapter);
|
|
ASSERT_INT_EQ(0, fl_adapter_size(adapter));
|
|
|
|
fl_adapter_free(adapter);
|
|
}
|
|
|
|
static void test_adapter_add_layers(void) {
|
|
fl_adapter* adapter = fl_adapter_new();
|
|
fl_config cfg = fl_config_default();
|
|
|
|
fl_layer* layer1 = fl_layer_new("layer1", 32, 64, &cfg);
|
|
fl_layer* layer2 = fl_layer_new("layer2", 64, 128, &cfg);
|
|
|
|
fl_adapter_add(adapter, layer1);
|
|
fl_adapter_add(adapter, layer2);
|
|
|
|
ASSERT_INT_EQ(2, fl_adapter_size(adapter));
|
|
|
|
fl_layer* found = fl_adapter_get(adapter, "layer1");
|
|
ASSERT_NOT_NULL(found);
|
|
ASSERT_STR_EQ("layer1", found->name);
|
|
|
|
found = fl_adapter_get(adapter, "layer2");
|
|
ASSERT_NOT_NULL(found);
|
|
ASSERT_STR_EQ("layer2", found->name);
|
|
|
|
ASSERT_NULL(fl_adapter_get(adapter, "nonexistent"));
|
|
|
|
fl_adapter_free(adapter); /* Also frees layers */
|
|
}
|
|
|
|
static void test_adapter_get_by_index(void) {
|
|
fl_adapter* adapter = fl_adapter_new();
|
|
fl_config cfg = fl_config_default();
|
|
|
|
fl_layer* layer = fl_layer_new("indexed", 32, 64, &cfg);
|
|
fl_adapter_add(adapter, layer);
|
|
|
|
fl_layer* found = fl_adapter_get_idx(adapter, 0);
|
|
ASSERT_NOT_NULL(found);
|
|
ASSERT_STR_EQ("indexed", found->name);
|
|
|
|
ASSERT_NULL(fl_adapter_get_idx(adapter, 1));
|
|
ASSERT_NULL(fl_adapter_get_idx(adapter, -1));
|
|
|
|
fl_adapter_free(adapter);
|
|
}
|
|
|
|
static void test_adapter_num_params(void) {
|
|
fl_adapter* adapter = fl_adapter_new();
|
|
|
|
fl_config cfg = {.rank = 4, .alpha = 8.0f, .dropout = 0.0f, .fan_in_fan_out = false};
|
|
|
|
/* Layer 1: A is (4 x 32) = 128, B is (64 x 4) = 256, total = 384 */
|
|
fl_layer* layer1 = fl_layer_new("l1", 32, 64, &cfg);
|
|
fl_adapter_add(adapter, layer1);
|
|
|
|
/* Layer 2: A is (4 x 64) = 256, B is (128 x 4) = 512, total = 768 */
|
|
fl_layer* layer2 = fl_layer_new("l2", 64, 128, &cfg);
|
|
fl_adapter_add(adapter, layer2);
|
|
|
|
/* Total: 384 + 768 = 1152 */
|
|
ASSERT_INT_EQ(1152, (int)fl_adapter_num_params(adapter));
|
|
|
|
fl_adapter_free(adapter);
|
|
}
|
|
|
|
/* ============================================================= */
|
|
/* Adapter I/O Tests */
|
|
/* ============================================================= */
|
|
|
|
static void test_adapter_save_and_load(void) {
|
|
const char* path = "/tmp/test_foundry_lora.safetensors";
|
|
|
|
/* Create adapter with known values */
|
|
fl_adapter* adapter = fl_adapter_new();
|
|
fl_config cfg = {.rank = 2, .alpha = 4.0f, .dropout = 0.0f, .fan_in_fan_out = false};
|
|
|
|
fl_layer* layer = fl_layer_new("test.proj", 8, 16, &cfg);
|
|
fl_adapter_add(adapter, layer);
|
|
|
|
/* Save */
|
|
bool saved = fl_adapter_save(adapter, path);
|
|
ASSERT_TRUE(saved);
|
|
|
|
fl_adapter_free(adapter);
|
|
|
|
/* Load */
|
|
fl_adapter* loaded = fl_adapter_load(path, &cfg);
|
|
ASSERT_NOT_NULL(loaded);
|
|
ASSERT_INT_EQ(1, fl_adapter_size(loaded));
|
|
|
|
fl_layer* loaded_layer = fl_adapter_get(loaded, "test.proj");
|
|
ASSERT_NOT_NULL(loaded_layer);
|
|
ASSERT_INT_EQ(8, loaded_layer->d_in);
|
|
ASSERT_INT_EQ(16, loaded_layer->d_out);
|
|
ASSERT_INT_EQ(2, loaded_layer->rank);
|
|
|
|
fl_adapter_free(loaded);
|
|
unlink(path);
|
|
}
|
|
|
|
/* ============================================================= */
|
|
/* Edge Cases */
|
|
/* ============================================================= */
|
|
|
|
static void test_null_safety(void) {
|
|
ASSERT_NULL(fl_layer_new(NULL, 32, 64, NULL));
|
|
ASSERT_NULL(fl_layer_new("test", 0, 64, NULL));
|
|
ASSERT_NULL(fl_layer_new("test", 32, 0, NULL));
|
|
ASSERT_NULL(fl_layer_from_tensors(NULL, NULL, NULL, 1.0f));
|
|
ASSERT_NULL(fl_layer_forward(NULL, NULL));
|
|
|
|
ASSERT_INT_EQ(0, fl_adapter_size(NULL));
|
|
ASSERT_NULL(fl_adapter_get(NULL, "test"));
|
|
ASSERT_NULL(fl_adapter_get_idx(NULL, 0));
|
|
ASSERT_INT_EQ(0, (int)fl_adapter_num_params(NULL));
|
|
|
|
fl_layer_free(NULL); /* Should not crash */
|
|
fl_adapter_free(NULL); /* Should not crash */
|
|
|
|
ASSERT_TRUE(1);
|
|
}
|
|
|
|
static void test_layer_b_initialized_to_zero(void) {
|
|
fl_config cfg = fl_config_default();
|
|
fl_layer* layer = fl_layer_new("test", 16, 32, &cfg);
|
|
|
|
/* B should be zeros initially (standard LoRA init) */
|
|
float* data = tm_data_f32(layer->lora_B);
|
|
for (int64_t i = 0; i < layer->lora_B->numel; i++) {
|
|
ASSERT_FLT_NEAR(0.0f, data[i], 0.0001f);
|
|
}
|
|
|
|
fl_layer_free(layer);
|
|
}
|
|
|
|
/* ============================================================= */
|
|
/* Main */
|
|
/* ============================================================= */
|
|
|
|
int main(void) {
|
|
printf("\n Running lora tests...\n\n");
|
|
|
|
/* Config tests */
|
|
RUN_TEST(config_default);
|
|
|
|
/* Layer creation tests */
|
|
RUN_TEST(layer_create);
|
|
RUN_TEST(layer_custom_rank);
|
|
RUN_TEST(layer_from_tensors);
|
|
|
|
/* Forward pass tests */
|
|
RUN_TEST(layer_forward_shape);
|
|
RUN_TEST(layer_forward_values);
|
|
|
|
/* Merge tests */
|
|
RUN_TEST(layer_merge);
|
|
RUN_TEST(layer_merge_adds_to_weight);
|
|
|
|
/* Adapter tests */
|
|
RUN_TEST(adapter_create);
|
|
RUN_TEST(adapter_add_layers);
|
|
RUN_TEST(adapter_get_by_index);
|
|
RUN_TEST(adapter_num_params);
|
|
|
|
/* Adapter I/O tests */
|
|
RUN_TEST(adapter_save_and_load);
|
|
|
|
/* Edge cases */
|
|
RUN_TEST(null_safety);
|
|
RUN_TEST(layer_b_initialized_to_zero);
|
|
|
|
TEST_REPORT();
|
|
|
|
return __ft_failed > 0 ? 1 : 0;
|
|
}
|