Skip to content

Commit

Permalink
test-backend-ops : cleanup, add moe test for batches
Browse files Browse the repository at this point in the history
  • Loading branch information
@slaren
slaren committed Dec 10, 2023
1 parent 54ba263 commit 54d254b
Showing 1 changed file with 32 additions and 35 deletions.
67 changes: 32 additions & 35 deletions tests/test-backend-ops.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,8 +20,6 @@ static void init_tensor_uniform(ggml_tensor * tensor, float min = -1.0f, float m
size_t size = ggml_nelements(tensor);
std::vector<float> data(size);

std::random_device rd;

#if 0
std::default_random_engine generator(rd());
std::uniform_real_distribution<float> distribution(min, max);
Expand All @@ -31,6 +29,7 @@ static void init_tensor_uniform(ggml_tensor * tensor, float min = -1.0f, float m
}
#endif
auto init_thread = [&](size_t start, size_t end) {
std::random_device rd;
std::default_random_engine generator(rd());
std::uniform_real_distribution<float> distribution(min, max);

Expand Down Expand Up @@ -341,13 +340,6 @@ struct test_case {
}
}

//if (t1->op == GGML_OP_SOFT_MAX) {
// printf("[%s] ", ggml_op_desc(t1));
// for (int i = 0; i < f1.size(); i++) {
// printf("(%x, %x) ", *(uint32_t*)&f1[i], *(uint32_t*)&f2[i]);
// }
// printf("\n");
//}
double err = nmse(f1.data(), f2.data(), f1.size());
if (err > ud->max_err) {
printf("[%s] NMSE = %f ", ggml_op_desc(t1), err);
Expand Down Expand Up @@ -447,8 +439,9 @@ struct test_case {
return size;
};
for (int i = 0; i < gf->n_nodes; i++) {
if (ggml_is_view_op(gf->nodes[i]->op) || gf->nodes[i] == out)
if (ggml_is_view_op(gf->nodes[i]->op) || gf->nodes[i] == out) {
continue;
}
mem += tensor_op_size(gf->nodes[i]);
}

Expand Down Expand Up @@ -1137,23 +1130,26 @@ struct test_sum_rows : public test_case {
}
};

// Mixtral MOE
struct test_moe : public test_case {
const int n_experts = 8;
const int n_experts_per_tok = 2;
const int n_tokens = 1;
const int n_embd = 4096;
const int n_ff = 14336;
const int n_experts;
const int n_experts_per_tok;
const int n_tokens;
const int n_embd;
const int n_ff;

std::string op_desc(ggml_tensor * t) override {
return "MOE";

GGML_UNUSED(t);
}

std::string vars() override {
return VARS_TO_STR5(n_experts, n_experts_per_tok, n_tokens, n_embd, n_ff);
}

test_moe() {
test_moe(int n_experts = 8, int n_experts_per_tok = 2, int n_tokens = 1, int n_embd = 4096, int n_ff = 14336)
: n_experts(n_experts), n_experts_per_tok(n_experts_per_tok), n_tokens(n_tokens), n_embd(n_embd), n_ff(n_ff) {
}

ggml_tensor * build_graph(ggml_context * ctx) override {
Expand All @@ -1171,24 +1167,20 @@ struct test_moe : public test_case {

ggml_tensor * cur = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, n_tokens);

ggml_tensor * logits = ggml_mul_mat(ctx, ffn_gate_inp, cur); // [n_tokens, num_experts]
ggml_tensor * probs = ggml_soft_max_ext(ctx, logits, nullptr, 1.0f/sqrtf(n_embd)); // [n_tokens, num_experts]
ggml_tensor * logits = ggml_mul_mat(ctx, ffn_gate_inp, cur);
ggml_tensor * probs = ggml_soft_max_ext(ctx, logits, nullptr, 1.0f/sqrtf(n_embd));

// select experts
ggml_tensor * selected_experts = ggml_top_k(ctx, probs, n_experts_per_tok); // [n_tokens, num_experts_per_tok]
ggml_tensor * selected_experts = ggml_top_k(ctx, probs, n_experts_per_tok);

ggml_tensor * weights = ggml_get_rows(ctx,
ggml_reshape_3d(ctx, probs, 1, n_experts, n_tokens), selected_experts);
printf("get rows args %ld %ld %ld %ld, %ld %ld %ld %ld\n",
weights->src[0]->ne[0], weights->src[0]->ne[1], weights->src[0]->ne[2], weights->src[0]->ne[3],
weights->src[1]->ne[0], weights->src[1]->ne[1], weights->src[1]->ne[2], weights->src[1]->ne[3]);


weights = ggml_reshape_2d(ctx, weights, n_experts_per_tok, n_tokens); // [n_tokens, num_experts_per_tok]
weights = ggml_reshape_2d(ctx, weights, n_experts_per_tok, n_tokens);

ggml_tensor * weights_sum = ggml_sum_rows(ctx, weights);

weights = ggml_div(ctx, weights, weights_sum); // [n_tokens, num_experts_per_tok]
weights = ggml_div(ctx, weights, weights_sum);

// compute expert outputs
ggml_tensor * moe_out = nullptr;
Expand All @@ -1202,9 +1194,9 @@ struct test_moe : public test_case {

cur_gate = ggml_silu(ctx, cur_gate);

cur_expert = ggml_mul(ctx, cur_up, cur_gate); // [n_tokens, n_embd]
cur_expert = ggml_mul(ctx, cur_up, cur_gate);

cur_expert = ggml_mul_mat_id(ctx, ffn_down_exp.data(), n_experts, selected_experts, i, cur_expert); // [n_tokens, n_embd]
cur_expert = ggml_mul_mat_id(ctx, ffn_down_exp.data(), n_experts, selected_experts, i, cur_expert);

cur_expert = ggml_mul(ctx, cur_expert,
ggml_view_2d(ctx, weights, 1, n_tokens, weights->nb[1], i*weights->nb[0]));
Expand Down Expand Up @@ -1240,8 +1232,6 @@ static bool test_backend(ggml_backend_t backend, test_mode mode, const char * op
GGML_TYPE_Q6_K
};

test_cases.emplace_back(new test_moe());

// unary ops
for (int op = 0; op < GGML_UNARY_OP_COUNT; op++) {
test_cases.emplace_back(new test_unary((ggml_unary_op) op));
Expand Down Expand Up @@ -1374,6 +1364,9 @@ static bool test_backend(ggml_backend_t backend, test_mode mode, const char * op

test_cases.emplace_back(new test_sum_rows());

test_cases.emplace_back(new test_moe(8, 2, 1, 4096, 14336));
test_cases.emplace_back(new test_moe(8, 2, 8, 4096, 14336));

// run tests
if (mode == MODE_TEST) {
ggml_backend_t backend_cpu = ggml_backend_cpu_init();
Expand All @@ -1389,14 +1382,17 @@ static bool test_backend(ggml_backend_t backend, test_mode mode, const char * op
ggml_backend_free(backend_cpu);

return n_ok == test_cases.size();
} else if (mode == MODE_PERF) {
}

if (mode == MODE_PERF) {
for (auto & test : test_cases) {
test->eval_perf(backend, op_name);
}
return true;
} else {
GGML_ASSERT(false);
}

GGML_ASSERT(false);
return false;
}

static void usage(char ** argv) {
Expand Down Expand Up @@ -1469,11 +1465,12 @@ int main(int argc, char ** argv) {
}

printf("%zu/%zu backends passed\n", n_ok, ggml_backend_reg_get_count());

if (n_ok != ggml_backend_reg_get_count()) {
printf("\033[1;31mFAIL\033[0m\n");
return 1;
} else {
printf("\033[1;32mOK\033[0m\n");
return 0;
}

printf("\033[1;32mOK\033[0m\n");
return 0;
}

0 comments on commit 54d254b

Please sign in to comment.