test_problematic_tensors.cpp

#include "arg.h"
#include "common.h"
#include "log.h"
#include "llama.h"
#include "ggml.h"
#include "gguf.h"
#include "math.h"
#include "string.h"

#include <cstdio>
#include <string>
#include <vector>
#include <numeric>
#include <fstream>
#include <thread>
#include <future>
#include <cmath>

static void set_tensor_type(ggml_tensor * tensor, ggml_type type) { // adapted from gguf_set_tensor_type
    const size_t  type_size = ggml_type_size(type);
    const int64_t blck_size = ggml_blck_size(type);

    tensor->type = type;
    GGML_ASSERT(tensor->ne[0] % blck_size == 0 && "tensor row size not divisible by block size of new type");

    tensor->nb[0] = type_size;
    tensor->nb[1] = tensor->nb[0]*(tensor->ne[0]/blck_size);
    for (int i = 2; i < GGML_MAX_DIMS; i++) {
        tensor->nb[i] = tensor->nb[i - 1]*tensor->ne[i - 1];
    }
}

static void dequantize(ggml_tensor * tensor) { // adapted from llama_tensor_dequantize_impl
    int64_t nelements = ggml_nelements(tensor);
    std::vector<float> output(nelements);
    float * f32_output = (float *) output.data();

    const ggml_type_traits * qtype = ggml_get_type_traits(tensor->type);

    uint8_t * data = (uint8_t *) tensor->data;
    std::vector<float> cdata;
    if ((tensor->buffer && !ggml_backend_buffer_is_host(tensor->buffer))) {
        auto n_bytes = ggml_nbytes(tensor);
        cdata.resize(n_bytes);
        ggml_backend_tensor_get(tensor, cdata.data(), 0, n_bytes);
        data = (uint8_t *) cdata.data();
    }

    if (tensor->type == GGML_TYPE_F16) {
        ggml_fp16_to_fp32_row((ggml_fp16_t *) data, f32_output, nelements);
    } else if (tensor->type == GGML_TYPE_BF16) {
        ggml_bf16_to_fp32_row((ggml_bf16_t *) data, f32_output, nelements);
    } else if (ggml_is_quantized(tensor->type)) {
        qtype->to_float(data, f32_output, nelements);
    } else {
        GGML_ABORT("fatal error"); // unreachable
    }
    
    set_tensor_type(tensor, GGML_TYPE_F32);
    float * new_data = (float *) malloc(output.size() * sizeof(float));
    memcpy(new_data, output.data(), output.size() * sizeof(float));
    tensor->data = new_data;
    double sum = 0.0f;
    float min = ((float *) tensor->data)[0];
    float max = ((float *) tensor->data)[0];
    for (int64_t i = 0; i < ggml_nelements(tensor); i++) {
        float elt = ((float *) tensor->data)[i];
        if (isnan(elt) || isinf(elt)) {
            GGML_ABORT("NaN or Inf found at position %ld", i);
        }
        sum += elt;
        if (elt < min) min = elt;
        if (elt > max) max = elt;
    }
    printf("\nSanity check: dequantized tensor has sum = %.8f, min = %.8f, max = %.8f\n", sum, min, max);
}

static void quantize(ggml_tensor * tensor, const float * source_data, ggml_type type) {
    int64_t nelements = ggml_nelements(tensor);
    
    size_t blck_size = tensor->ne[0];
    size_t n_blocks = tensor->ne[1];
    size_t n_experts = tensor->ne[2];
    
    size_t expert_size = ggml_row_size(type, n_blocks * blck_size);
    std::vector<uint8_t> dataq(ggml_row_size(type, n_blocks * blck_size * n_experts));
    
    printf("Quantizing to %s", ggml_type_name(type));
    for (size_t i = 0; i < n_experts; i++) {
        printf(".");
        ggml_quantize_chunk(type, source_data + (n_blocks * blck_size) * i, dataq.data() + expert_size * i, 0, n_blocks, blck_size, nullptr);
    }
    printf(" DONE\n");
    set_tensor_type(tensor, type);
    ggml_backend_tensor_set(tensor, dataq.data(), 0, dataq.size());
}

int main(int argc, char ** argv) {

    llama_log_set(nullptr, nullptr);
    llama_backend_init();
    ggml_backend_load_all_from_path("build/bin");

    // Initialize GGML context
    struct ggml_init_params params = {
        /*.mem_size   =*/ 10 * ggml_tensor_overhead() + ggml_graph_overhead(),
        /*.mem_buffer =*/ NULL,
        /*.no_alloc   =*/ true,
    };

    ggml_context * gctx = ggml_init(params);
    ggml_context * gctx_cpu = ggml_init(params);
    ggml_context * wctx = nullptr;
    ggml_context * nctx = nullptr;
    ggml_context * ictx = nullptr;
    struct gguf_init_params wparams = {
        /*.no_alloc = */ false,
        /*.ctx      = */ &wctx,
    }; 
    struct gguf_init_params nparams = {
        /*.no_alloc = */ false,
        /*.ctx      = */ &nctx,
    }; 
    struct gguf_init_params iparams = {
        /*.no_alloc = */ false,
        /*.ctx      = */ &ictx,
    }; 
    gguf_context * wgctx = gguf_init_from_file("problem-tensors-weights.gguf", wparams);
    gguf_context * ngctx = gguf_init_from_file("problem-tensors-norm.gguf", nparams);
    gguf_context * igctx = gguf_init_from_file("problem-tensors-ids.gguf", iparams);
    
    ggml_tensor * weights = ggml_get_next_tensor(wctx, ggml_get_first_tensor(wctx));
    ggml_tensor * norm = ggml_get_next_tensor(nctx, ggml_get_first_tensor(nctx));
    ggml_tensor * ids = ggml_get_next_tensor(ictx, ggml_get_first_tensor(ictx));

    ggml_context * gctx_cpu_comp = ggml_init(params);
    struct ggml_cgraph * gf_cpu = ggml_new_graph(gctx_cpu_comp);
    ggml_tensor * mul_mat_id_cpu = ggml_mul_mat_id(gctx_cpu, weights, norm, ids);
    ggml_build_forward_expand(gf_cpu, mul_mat_id_cpu);

    ggml_backend_t cpu = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr);

    ggml_gallocr_t allocr = ggml_gallocr_new(ggml_backend_get_default_buffer_type(cpu));
    ggml_gallocr_alloc_graph(allocr, gf_cpu);

    ggml_backend_graph_compute(cpu, gf_cpu);

    double sum_cpu = 0.0f;
    float max_cpu = ((float *) mul_mat_id_cpu->data)[0];
    float min_cpu = ((float *) mul_mat_id_cpu->data)[0];
    for (uint64_t i = 0; i < ggml_nelements(mul_mat_id_cpu); i++) {
        float elt = ((float *) mul_mat_id_cpu->data)[i];
        sum_cpu += elt;
        max_cpu = elt > max_cpu ? elt : max_cpu;
        min_cpu = elt < min_cpu ? elt : min_cpu;
    }
    printf("\nCPU sum of matmul: %.8f, max: %.8f, min: %.8f, nelements: %lu\n\n", sum_cpu, max_cpu, min_cpu, ggml_nelements(mul_mat_id_cpu));
    
    struct ggml_cgraph * gf = ggml_new_graph(gctx);
    
    ggml_tensor * w_cuda = ggml_new_tensor_4d(gctx, weights->type, weights->ne[0], weights->ne[1], weights->ne[2], weights->ne[3]);
    ggml_tensor * n_cuda = ggml_new_tensor_4d(gctx, norm->type, norm->ne[0], norm->ne[1], norm->ne[2], norm->ne[3]);
    ggml_tensor * i_cuda = ggml_new_tensor_4d(gctx, ids->type, ids->ne[0], ids->ne[1], ids->ne[2], ids->ne[3]);

    ggml_backend_t cuda = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_GPU, nullptr);
    ggml_backend_alloc_ctx_tensors(gctx, cuda);
    
    ggml_backend_tensor_set(w_cuda, weights->data, 0, ggml_nbytes(w_cuda));
    ggml_backend_tensor_set(n_cuda, norm->data, 0, ggml_nbytes(n_cuda));
    ggml_backend_tensor_set(i_cuda, ids->data, 0, ggml_nbytes(i_cuda));

    ggml_context * gctx_cuda_comp = ggml_init(params);
    struct ggml_cgraph * gf_cuda = ggml_new_graph(gctx_cuda_comp);
    ggml_tensor * mul_mat_id_cuda = ggml_mul_mat_id(gctx_cuda_comp, w_cuda, n_cuda, i_cuda);
    ggml_build_forward_expand(gf_cuda, mul_mat_id_cuda);

    ggml_gallocr_t cuda_allocr = ggml_gallocr_new(ggml_backend_get_default_buffer_type(cuda));
    ggml_gallocr_alloc_graph(cuda_allocr, gf_cuda);
    ggml_backend_graph_compute(cuda, gf_cuda);

    std::vector<float> vec;
    
    auto n_bytes = ggml_nbytes(mul_mat_id_cuda);
    vec.resize(n_bytes);
    ggml_backend_tensor_get(mul_mat_id_cuda, vec.data(), 0, n_bytes);
    double sum = 0.0f;
    float max = vec[0];
    float min = vec[0];
    float maxdiff = 0;
    uint64_t maxdiff_pos = -1;
    for (uint64_t i = 0; i < ggml_nelements(mul_mat_id_cuda); i++) {
        float elt = vec[i];
        float org_elt = ((float *) mul_mat_id_cpu->data)[i];
        float diff = fabs(elt - org_elt);
        if (diff > maxdiff) {
            maxdiff = diff;
            maxdiff_pos = i;
        }
        sum += elt;
        max = elt > max ? elt : max;
        min = elt < min ? elt : min;
    }
    printf("CUDA sum of matmul: %.8f, max: %.8f, min: %.8f, max diff: %.8f at pos %lu, nelements: %lu\n\n", sum, max, min, maxdiff, maxdiff_pos, ggml_nelements(mul_mat_id_cuda));
    ggml_gallocr_free(cuda_allocr);

    dequantize(weights);

    ggml_context * gctx_cpu_comp_deq = ggml_init(params);
    struct ggml_cgraph * gf_cpu_deq = ggml_new_graph(gctx_cpu_comp_deq);
    ggml_tensor * mul_mat_id_cpu_deq = ggml_mul_mat_id(gctx_cpu_comp_deq, weights, norm, ids);
    ggml_build_forward_expand(gf_cpu_deq, mul_mat_id_cpu_deq);

    ggml_gallocr_t allocr_deq = ggml_gallocr_new(ggml_backend_get_default_buffer_type(cpu));
    ggml_gallocr_alloc_graph(allocr_deq, gf_cpu_deq);

    ggml_backend_graph_compute(cpu, gf_cpu_deq);

    double sum_cpu_deq = 0.0f;
    float max_cpu_deq = ((float *) mul_mat_id_cpu_deq->data)[0];
    float min_cpu_deq = ((float *) mul_mat_id_cpu_deq->data)[0];
    for (uint64_t i = 0; i < ggml_nelements(mul_mat_id_cpu_deq); i++) {
        float elt = ((float *) mul_mat_id_cpu_deq->data)[i];
        sum_cpu_deq += elt;
        max_cpu_deq = elt > max_cpu_deq ? elt : max_cpu_deq;
        min_cpu_deq = elt < min_cpu_deq ? elt : min_cpu_deq;
    }
    printf("\nCPU sum of matmul (dequantized): %.8f, max: %.8f, min: %.8f, nelements: %lu\n\n", sum_cpu_deq, max_cpu_deq, min_cpu_deq, ggml_nelements(mul_mat_id_cpu_deq));
    ggml_gallocr_free(allocr_deq);

    ggml_context * gctx_cuda_comp_deq = ggml_init(params);
    ggml_context * gctx_cuda_dequant = ggml_init(params);
    
    struct ggml_cgraph * gf_cuda_deq = ggml_new_graph(gctx_cuda_comp_deq);
    ggml_tensor * w_cuda_deq = ggml_new_tensor_4d(gctx_cuda_comp_deq, GGML_TYPE_F32, weights->ne[0], weights->ne[1], weights->ne[2], weights->ne[3]);
    ggml_backend_alloc_ctx_tensors(gctx_cuda_comp_deq, cuda);
    ggml_backend_tensor_set(w_cuda_deq, weights->data, 0, ggml_nbytes(weights));

    ggml_tensor * mul_mat_id_cuda_deq = ggml_mul_mat_id(gctx_cuda_comp_deq, w_cuda_deq, n_cuda, i_cuda);
    ggml_build_forward_expand(gf_cuda_deq, mul_mat_id_cuda_deq);

    ggml_gallocr_t allocr_cuda_deq = ggml_gallocr_new(ggml_backend_get_default_buffer_type(cuda));
    ggml_gallocr_alloc_graph(allocr_cuda_deq, gf_cuda_deq);

    ggml_backend_graph_compute(cuda, gf_cuda_deq);

    std::vector<float> vec_deq(ggml_nbytes(mul_mat_id_cuda_deq));
    ggml_backend_tensor_get(mul_mat_id_cuda_deq, vec_deq.data(), 0, n_bytes);

    double sum_cuda_deq = 0.0f;
    float max_cuda_deq = vec_deq[0];
    float min_cuda_deq = vec_deq[0];
    for (uint64_t i = 0; i < vec_deq.size(); i++) {
        float elt = vec_deq[i];
        sum_cuda_deq += elt;
        max_cuda_deq = elt > max_cpu_deq ? elt : max_cpu_deq;
        min_cuda_deq = elt < min_cpu_deq ? elt : min_cpu_deq;
    }
    printf("\nCUDA sum of matmul (dequantized): %.8f, max: %.8f, min: %.8f, nelements: %lu\n\n", sum_cuda_deq, max_cuda_deq, min_cuda_deq, ggml_nelements(mul_mat_id_cuda_deq));
    ggml_gallocr_free(allocr_cuda_deq);
    ggml_free(gctx_cuda_comp_deq);
    ggml_free(gctx_cuda_dequant);

    /*ggml_type test_quantizations[] = { GGML_TYPE_IQ2_S, GGML_TYPE_IQ3_XXS, GGML_TYPE_IQ3_S, GGML_TYPE_IQ4_NL, GGML_TYPE_Q4_1 };
    for (int i = 0; i < 6; i++) {
        std::vector<float> qdata;
        {
            ggml_context * gctx_cuda_quant = ggml_init(params);
            ggml_context * gctx_cuda_req_comp = ggml_init(params);
            ggml_tensor * w_cuda_qt = ggml_new_tensor_4d(gctx_cuda_quant, test_quantizations[i], weights->ne[0], weights->ne[1], weights->ne[2], weights->ne[3]);
            ggml_backend_alloc_ctx_tensors(gctx_cuda_quant, cuda);
            quantize(w_cuda_qt, (const float *) weights->data, test_quantizations[i]);
            qdata.resize(ggml_nbytes(w_cuda_qt));
            ggml_backend_tensor_get(w_cuda_qt, qdata.data(), 0, qdata.size());

            struct ggml_cgraph * gf_cuda_req = ggml_new_graph(gctx_cuda_req_comp);
            ggml_tensor * mul_mat_id_cuda_req = ggml_mul_mat_id(gctx_cuda_comp, w_cuda_qt, n_cuda, i_cuda);
            ggml_build_forward_expand(gf_cuda_req, mul_mat_id_cuda_req);

            ggml_gallocr_t cuda_allocr_req = ggml_gallocr_new(ggml_backend_get_default_buffer_type(cuda));
            ggml_gallocr_alloc_graph(cuda_allocr_req, gf_cuda_req);
            ggml_backend_graph_compute(cuda, gf_cuda_req);

            std::vector<float> vec;
            
            auto n_bytes = ggml_nbytes(mul_mat_id_cuda_req);
            vec.resize(n_bytes);
            ggml_backend_tensor_get(mul_mat_id_cuda_req, vec.data(), 0, n_bytes);
            double sum = 0.0f;
            float max = vec[0];
            float min = vec[0];
            float maxdiff = 0;
            uint64_t maxdiff_pos = -1;
            for (uint64_t i = 0; i < ggml_nelements(mul_mat_id_cuda_req); i++) {
                float elt = vec[i];
                float org_elt = ((float *) mul_mat_id_cpu->data)[i];
                float diff = fabs(elt - org_elt);
                if (diff > maxdiff) {
                    maxdiff = diff;
                    maxdiff_pos = i;
                }
                sum += elt;
                max = elt > max ? elt : max;
                min = elt < min ? elt : min;
            }
            printf("CUDA sum of quant %s matmul: %.8f, max: %.8f, min: %.8f, max diff: %.8f at pos %lu, nelements: %lu\n\n", ggml_type_name(test_quantizations[i]), sum, max, min, maxdiff, maxdiff_pos, ggml_nelements(mul_mat_id_cuda_req));
            ggml_gallocr_free(cuda_allocr_req);
            ggml_free(gctx_cuda_quant);
            ggml_free(gctx_cuda_req_comp);
        }
        {
            ggml_context * gctx_cpu_quant = ggml_init(params);
            ggml_context * gctx_cpu_req_comp = ggml_init(params);
            ggml_tensor * w_cpu_qt = ggml_new_tensor_4d(gctx_cpu_quant, test_quantizations[i], weights->ne[0], weights->ne[1], weights->ne[2], weights->ne[3]);
            ggml_backend_alloc_ctx_tensors(gctx_cpu_quant, cpu);
            set_tensor_type(w_cpu_qt, test_quantizations[i]);
            w_cpu_qt->data = qdata.data();

            struct ggml_cgraph * gf_cpu_req = ggml_new_graph(gctx_cpu_req_comp);
            ggml_tensor * mul_mat_id_cpu_req = ggml_mul_mat_id(gctx_cpu_comp, w_cpu_qt, norm, ids);
            ggml_build_forward_expand(gf_cpu_req, mul_mat_id_cpu_req);

            ggml_gallocr_t cpu_allocr_req = ggml_gallocr_new(ggml_backend_get_default_buffer_type(cpu));
            ggml_gallocr_alloc_graph(cpu_allocr_req, gf_cpu_req);
            ggml_backend_graph_compute(cuda, gf_cpu_req);

            std::vector<float> vec;
            
            auto n_bytes = ggml_nbytes(mul_mat_id_cpu_req);
            vec.resize(n_bytes);
            ggml_backend_tensor_get(mul_mat_id_cpu_req, vec.data(), 0, n_bytes);
            double sum = 0.0f;
            float max = vec[0];
            float min = vec[0];
            float maxdiff = 0;
            uint64_t maxdiff_pos = -1;
            for (uint64_t i = 0; i < ggml_nelements(mul_mat_id_cpu_req); i++) {
                float elt = vec[i];
                float org_elt = ((float *) mul_mat_id_cpu->data)[i];
                float diff = fabs(elt - org_elt);
                if (diff > maxdiff) {
                    maxdiff = diff;
                    maxdiff_pos = i;
                }
                sum += elt;
                max = elt > max ? elt : max;
                min = elt < min ? elt : min;
            }
            printf("CPU sum of quant %s matmul: %.8f, max: %.8f, min: %.8f, max diff: %.8f at pos %lu, nelements: %lu\n\n", ggml_type_name(test_quantizations[i]), sum, max, min, maxdiff, maxdiff_pos, ggml_nelements(mul_mat_id_cpu_req));
            ggml_gallocr_free(cpu_allocr_req);
            ggml_free(gctx_cpu_quant);
            ggml_free(gctx_cpu_req_comp);
        }
    }*/
    return 0;
}