Brunobkr/llama.cpp_AlgMor24_github

1#include "common.h"2#include "log.h"3#include "ggml-backend.h"4#include "ggml.h"5#include "gguf.h"6#include "ggml-cpp.h"7#include "llama.h"8#include "llama-cpp.h"9 10// TODO: replace with #include "llama-ext.h" in the future11#include "../src/llama-arch.h"12#include "../src/llama-model-saver.h"13 14#include <cinttypes>15#include <cstddef>16#include <cstdio>17#include <cstring>18#include <cstdint>19#include <random>20#include <stdexcept>21#include <string>22#include <utility>23#include <vector>24 25// normalized mean squared error = mse(a, b) / mse(a, 0)26static double nmse(const std::vector<float> & a, const std::vector<float> & b) {27    GGML_ASSERT(a.size() == b.size());28    double mse_a_b = 0.0;29    double mse_a_0 = 0.0;30 31    for (size_t i = 0; i < a.size(); i++) {32        float a_i = a[i];33        float b_i = b[i];34 35        mse_a_b += (a_i - b_i) * (a_i - b_i);36        mse_a_0 += a_i * a_i;37    }38 39    return mse_a_b / mse_a_0;40}41 42static void set_tensor_data(struct ggml_tensor * tensor, void * userdata) {43    size_t seed = *(const size_t *) userdata;44    std::hash<std::string> hasher;45    seed ^= hasher(tensor->name);46    std::mt19937 gen(seed);47    std::normal_distribution<float> dis(0.0f, 1.0e-2f);48 49    const int64_t ne = ggml_nelements(tensor);50    if (tensor->type == GGML_TYPE_F32) {51        std::vector<float> tmp(ne);52        for (int64_t i = 0; i < ne; i++) {53            tmp[i] = dis(gen);54        }55        ggml_backend_tensor_set(tensor, tmp.data(), 0, ggml_nbytes(tensor));56    } else if (tensor->type == GGML_TYPE_F16) {57        std::vector<ggml_fp16_t> tmp(ne);58        for (int64_t i = 0; i < ne; i++) {59            tmp[i] = ggml_fp32_to_fp16(dis(gen));60        }61        ggml_backend_tensor_set(tensor, tmp.data(), 0, ggml_nbytes(tensor));62    } else {63        GGML_ABORT("fatal error");64    }65}66 67static void usage(char ** argv) {68    printf("Usage: %s [-a/--arch arch] [-s/--seed seed] [-v/--verbose]\n", argv[0]);69}70 71static std::vector<llama_token> get_tokens(const uint32_t n_tokens, const uint32_t n_vocab, const size_t seed){72    std::mt19937 gen(seed);73    std::uniform_int_distribution<> dis(0, n_vocab - 1);74    std::vector<llama_token> ret;75    ret.reserve(n_tokens);76    for (uint32_t i = 0; i < n_tokens; i++) {77        ret.push_back(dis(gen));78    }79    return ret;80}81 82static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {83    gguf_context_ptr ret(gguf_init_empty());84    llama_model_saver ms(arch, ret.get());85    const uint32_t n_ctx = 128;86 87    uint32_t n_vocab = 128;88    uint32_t n_embd  = 256;89    uint32_t n_head  = 2;90    uint32_t n_ff    = 384;91    uint32_t n_layer = 2;92    if (arch == LLM_ARCH_LLAMA4) {93        n_layer = 4; // hparams.n_no_rope_layer_step is hard-coded to 494    } else if (arch == LLM_ARCH_GEMMA4) {95        n_embd = 128;96        n_head = 2;97        n_ff   = 192;98        n_layer = 5; // need at least 5 for swa_pattern (every 5th is full_attention)99    } else if (arch == LLM_ARCH_GEMMA3N) {100        n_embd = 64;101        n_head = 1;102        n_ff   = 96;103        n_layer = 22; // hparams.n_layer_kv_from_start = 20 is hardcoded104    } else if (arch == LLM_ARCH_DEEPSEEK2105            || arch == LLM_ARCH_DEEPSEEK32106            || arch == LLM_ARCH_GLM_DSA107            || arch == LLM_ARCH_KIMI_LINEAR108            || arch == LLM_ARCH_MISTRAL4) {109        n_embd = 128;110        n_head = 1;111        n_ff   = 192;112    } else if (arch == LLM_ARCH_NEMOTRON_H || arch == LLM_ARCH_NEMOTRON_H_MOE) {113        n_layer = 3;114    } else if (arch == LLM_ARCH_CHAMELEON) {115        n_vocab = 10240;116    } else if (arch == LLM_ARCH_QWEN3TTS) {117        n_vocab = 4096; // must be >= the hard-coded codec head size (3072)118    }119 120    const uint32_t n_embd_head = n_embd / n_head;121 122    ms.add_kv(LLM_KV_GENERAL_ARCHITECTURE,      llm_arch_name(arch));123    ms.add_kv(LLM_KV_VOCAB_SIZE,                n_vocab);124    ms.add_kv(LLM_KV_CONTEXT_LENGTH,            n_ctx);125    ms.add_kv(LLM_KV_EMBEDDING_LENGTH,          n_embd);126    ms.add_kv(LLM_KV_FEATURES_LENGTH,           n_embd);127    ms.add_kv(LLM_KV_BLOCK_COUNT,               n_layer);128    ms.add_kv(LLM_KV_LEADING_DENSE_BLOCK_COUNT, uint32_t(1));129 130    if (arch == LLM_ARCH_NEMOTRON_H || arch == LLM_ARCH_NEMOTRON_H_MOE) {131        std::vector<uint32_t> n_ff_per_layer;132        n_ff_per_layer.reserve(n_layer);133        for (uint32_t il = 0; il < n_layer; il++) {134            n_ff_per_layer.push_back(il <= 1 ? 0 : n_ff);135        }136        ms.add_kv(LLM_KV_FEED_FORWARD_LENGTH, n_ff_per_layer);137    } else {138        ms.add_kv(LLM_KV_FEED_FORWARD_LENGTH, n_ff);139    }140 141    ms.add_kv(LLM_KV_USE_PARALLEL_RESIDUAL,   false);142    ms.add_kv(LLM_KV_LOGIT_SCALE,             1.0f);143    ms.add_kv(LLM_KV_TIME_MIX_EXTRA_DIM,      uint32_t(64));144    ms.add_kv(LLM_KV_TIME_DECAY_EXTRA_DIM,    uint32_t(128));145    ms.add_kv(LLM_KV_FULL_ATTENTION_INTERVAL, uint32_t(2));146 147    if (arch == LLM_ARCH_PLAMO2 || arch == LLM_ARCH_JAMBA || arch == LLM_ARCH_NEMOTRON_H || arch == LLM_ARCH_NEMOTRON_H_MOE ||148            arch == LLM_ARCH_GRANITE_HYBRID || arch == LLM_ARCH_LFM2 || arch == LLM_ARCH_LFM2MOE || arch == LLM_ARCH_KIMI_LINEAR) {149        GGML_ASSERT(n_layer >= 2);150        std::vector<uint32_t> n_head_per_layer;151        n_head_per_layer.reserve(n_layer);152        for (uint32_t il = 0; il < n_layer; il++) {153            n_head_per_layer.push_back(il == 1 ? 0 : n_head);154        }155        ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT, n_head_per_layer);156        ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT_KV, n_head_per_layer);157    } else {158        ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT, n_head);159        ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT_KV, n_head);160    }161 162    ms.add_kv(LLM_KV_ATTENTION_MAX_ALIBI_BIAS, 8.0f);163    if (arch == LLM_ARCH_DEEPSEEK2164            || arch == LLM_ARCH_DEEPSEEK32165            || arch == LLM_ARCH_GLM_DSA166            || arch == LLM_ARCH_KIMI_LINEAR167            || arch == LLM_ARCH_MISTRAL4) {168        ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH,       uint32_t(576));169        ms.add_kv(LLM_KV_ATTENTION_VALUE_LENGTH,     uint32_t(512));170        ms.add_kv(LLM_KV_ROPE_DIMENSION_COUNT,       uint32_t(64));171        ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH_MLA,   uint32_t(192));172        ms.add_kv(LLM_KV_ATTENTION_VALUE_LENGTH_MLA, uint32_t(128));173    } else if (arch == LLM_ARCH_MINIMAX_M3) {174        // partial rotary: n_rot must not exceed the indexer key length (64)175        ms.add_kv(LLM_KV_ROPE_DIMENSION_COUNT,       uint32_t(64));176    }177    ms.add_kv(LLM_KV_ATTENTION_CLAMP_KQV,              1.0f);178    ms.add_kv(LLM_KV_ATTENTION_LAYERNORM_EPS,          1e-5f);179    ms.add_kv(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS,      1e-5f);180    ms.add_kv(LLM_KV_ATTENTION_GROUPNORM_EPS,          1e-5f);181    ms.add_kv(LLM_KV_ATTENTION_GROUPNORM_GROUPS,       uint32_t(8));182    ms.add_kv(LLM_KV_ATTENTION_Q_LORA_RANK,            uint32_t(512));183    ms.add_kv(LLM_KV_ATTENTION_KV_LORA_RANK,           uint32_t(512));184    ms.add_kv(LLM_KV_ATTENTION_RELATIVE_BUCKETS_COUNT, uint32_t(8));185    ms.add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW,         n_ctx/8);186 187    if (arch == LLM_ARCH_GEMMA4) {188        ms.add_kv(LLM_KV_EMBEDDING_LENGTH_PER_LAYER,      n_embd/2);189        ms.add_kv(LLM_KV_ATTENTION_SHARED_KV_LAYERS,      uint32_t(0));190        ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH_SWA,        n_embd_head);191        ms.add_kv(LLM_KV_ATTENTION_VALUE_LENGTH_SWA,      n_embd_head);192        ms.add_kv(LLM_KV_ROPE_FREQ_BASE_SWA,              10000.0f);193        // SWA pattern: every 5th layer is full attention (matches E2B layer_types)194        ms.add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, uint32_t(5));195    } else if (arch == LLM_ARCH_COHERE2MOE || arch == LLM_ARCH_MIMO2 || arch == LLM_ARCH_STEP35 || arch == LLM_ARCH_MUSE_GLIMMER) {196        std::vector<uint32_t> pattern;197        pattern.reserve(n_layer);198        for (uint32_t il = 0; il < n_layer; il++) {199            pattern.push_back(il % 2);200        }201        ms.add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, pattern);202    } else {203        ms.add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, uint32_t(2));204    }205 206    // MSA requires one indexer head per GQA (KV) head, unlike the DSA archs where the207    // indexer head count is independent of the main attention head count.208    ms.add_kv(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT,   arch == LLM_ARCH_MINIMAX_M3 ? n_head : uint32_t(1));209    ms.add_kv(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH,   uint32_t(64));210    ms.add_kv(LLM_KV_ATTENTION_INDEXER_TOP_K,        uint32_t(8));211    ms.add_kv(LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE,   uint32_t(4));212    ms.add_kv(LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, uint32_t(1));213    ms.add_kv(LLM_KV_ROPE_DIMENSION_SECTIONS, std::vector<uint32_t>({n_embd_head/4, n_embd_head/4, n_embd_head/4, n_embd_head/4}));214    ms.add_kv(LLM_KV_TOKENIZER_MODEL,         "no_vocab");215    // ms.add_kv(LLM_KV_DENSE_2_FEAT_OUT,     n_embd);216    // ms.add_kv(LLM_KV_DENSE_3_FEAT_IN,      n_embd);217 218    if (moe) {219        ms.add_kv(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, n_ff);220        ms.add_kv(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, n_ff / 2);  // distinct from n_ff so a saver key-clobber surfaces on reload221        ms.add_kv(LLM_KV_INTERLEAVE_MOE_LAYER_STEP,  uint32_t(2));222        ms.add_kv(LLM_KV_EXPERT_COUNT,               uint32_t(2));223        ms.add_kv(LLM_KV_EXPERT_USED_COUNT,          uint32_t(1));224        ms.add_kv(LLM_KV_EXPERT_SHARED_COUNT,        uint32_t(1));225        ms.add_kv(LLM_KV_EXPERT_GATING_FUNC,         uint32_t(2)); // sigmoid226        ms.add_kv(LLM_KV_EXPERT_GROUP_SCALE,         1.0f);227        ms.add_kv(LLM_KV_EXPERTS_PER_GROUP,          uint32_t(1));228    }229 230    ms.add_kv(LLM_KV_POSNET_EMBEDDING_LENGTH,   n_embd);231    ms.add_kv(LLM_KV_POSNET_BLOCK_COUNT,        n_layer);232    ms.add_kv(LLM_KV_CONVNEXT_EMBEDDING_LENGTH, n_embd);233    ms.add_kv(LLM_KV_CONVNEXT_BLOCK_COUNT,      n_layer);234    ms.add_kv(LLM_KV_XIELU_ALPHA_N,             1.0f);235    ms.add_kv(LLM_KV_XIELU_ALPHA_P,             1.0f);236    ms.add_kv(LLM_KV_XIELU_BETA,                1.0f);237    ms.add_kv(LLM_KV_XIELU_EPS,                 1.0e-7f);238    ms.add_kv(LLM_KV_SSM_INNER_SIZE,            arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE ? 256 : 2*n_embd);239    ms.add_kv(LLM_KV_SSM_CONV_KERNEL,           uint32_t(4));240    ms.add_kv(LLM_KV_SSM_STATE_SIZE,            uint32_t(128));241    ms.add_kv(LLM_KV_SSM_TIME_STEP_RANK,        n_head);242    ms.add_kv(LLM_KV_SSM_GROUP_COUNT,           arch == LLM_ARCH_PLAMO2 ? 0 : uint32_t(2));243    ms.add_kv(LLM_KV_KDA_HEAD_DIM,              uint32_t(128));244    ms.add_kv(LLM_KV_WKV_HEAD_SIZE,             n_embd/n_head);245    ms.add_kv(LLM_KV_SHORTCONV_L_CACHE,         uint32_t(3));246 247    for (uint32_t il = 0; il < n_layer; il++) {248        ggml_tensor t;249        memset(&t, 0, sizeof(ggml_tensor));250        t.type = GGML_TYPE_F16;251        ggml_format_name(&t, "conv%" PRIu32 "d.weight", il);252        gguf_add_tensor(ms.gguf_ctx, &t);253        ggml_format_name(&t, "posnet.%" PRIu32 ".conv1.weight", il);254        gguf_add_tensor(ms.gguf_ctx, &t);255        ggml_format_name(&t, "posnet.%" PRIu32 ".conv2.weight", il);256        gguf_add_tensor(ms.gguf_ctx, &t);257        ggml_format_name(&t, "convnext.%" PRIu32 ".dw.weight", il);258        gguf_add_tensor(ms.gguf_ctx, &t);259    }260    return ret;261}262 263static bool silent_model_load_progress(float /*progress*/, void * /*user_data*/) {264    return true;265}266 267static std::pair<llama_model_ptr, llama_context_ptr> get_model_and_ctx(268        struct gguf_context * gguf_ctx, FILE * file, const size_t seed, const std::vector<ggml_backend_dev_t> & devs,269        const llama_split_mode split_mode = LLAMA_SPLIT_MODE_LAYER, bool encode = false) {270    GGML_ASSERT((gguf_ctx == nullptr) != (file == nullptr));271    llama_model_params model_params = llama_model_default_params();272    model_params.progress_callback = silent_model_load_progress;273    std::vector<ggml_backend_dev_t> devs_copy = devs;274    devs_copy.push_back(nullptr);275    model_params.devices = devs_copy.data();276    model_params.split_mode = split_mode;277 278    llama_context_params ctx_params = llama_context_default_params();279    ctx_params.n_ctx = 0;280    ctx_params.n_threads = 4;281    ctx_params.n_threads_batch = 4;282    if (!encode) {283        ctx_params.n_ubatch = 64;284    }285 286    size_t tmp = seed;287    llama_model_ptr model(gguf_ctx != nullptr ?288        llama_model_init_from_user(gguf_ctx, set_tensor_data, &tmp, model_params) :289        llama_model_load_from_file_ptr(file, model_params));290    if (!model) {291        throw std::runtime_error("failed to create llama model");292    }293    llama_context_ptr lctx(llama_init_from_model(model.get(), ctx_params));294    if (!lctx) {295        throw std::runtime_error("failed to create llama context");296    }297    return std::make_pair(std::move(model), std::move(lctx));298}299 300static std::vector<float> get_logits(301        llama_model * model, llama_context * lctx, const std::vector<llama_token> & tokens, bool encode = false) {302    const uint32_t n_vocab  = llama_vocab_n_tokens(llama_model_get_vocab(model));303    const uint32_t n_ctx    = llama_n_ctx(lctx);304    const uint32_t n_tokens = tokens.size();305    llama_batch batch = llama_batch_init(n_ctx, 0, 1);306    GGML_ASSERT(n_tokens <= n_ctx);307    for (uint32_t pos = 0; pos < n_tokens; pos++) {308        common_batch_add(batch, tokens[pos], pos, {0}, true);309    }310    batch.n_tokens = n_tokens;311    if (encode) {312        if (llama_encode(lctx, batch)) {313            llama_batch_free(batch);314            throw std::runtime_error("failed to encode batch");315        }316    }317    if (llama_decode(lctx, batch)) {318        llama_batch_free(batch);319        throw std::runtime_error("failed to decode batch");320    }321 322    std::vector<float> ret;323    ret.reserve(n_tokens*n_vocab);324    for (uint32_t i = 0; i < n_tokens; i++) {325        const float * logits_ith = llama_get_logits_ith(lctx, i);326        for (uint32_t j = 0; j < n_vocab; j++) {327            ret.push_back(logits_ith[j]);328        }329    }330    llama_batch_free(batch);331    return ret;332}333 334static bool moe_mandatory(const llm_arch arch) {335    switch (arch) {336        case LLM_ARCH_LLAMA4:337        case LLM_ARCH_COHERE2MOE:338        case LLM_ARCH_GROK:339        case LLM_ARCH_QWEN2MOE:340        case LLM_ARCH_QWEN3MOE:341        case LLM_ARCH_QWEN3NEXT:342        case LLM_ARCH_QWEN3VLMOE:343        case LLM_ARCH_QWEN35MOE:344        case LLM_ARCH_PHIMOE:345        case LLM_ARCH_DBRX:346        case LLM_ARCH_OLMOE:347        case LLM_ARCH_ARCTIC:348        case LLM_ARCH_DEEPSEEK:349        case LLM_ARCH_DEEPSEEK2:350        case LLM_ARCH_DEEPSEEK32:351        case LLM_ARCH_GLM4_MOE:352        case LLM_ARCH_GLM_DSA:353        case LLM_ARCH_EXAONE_MOE:354        case LLM_ARCH_BAILINGMOE:355        case LLM_ARCH_BAILINGMOE2:356        case LLM_ARCH_DOTS1:357        case LLM_ARCH_AFMOE:358        case LLM_ARCH_ERNIE4_5:359        case LLM_ARCH_ERNIE4_5_MOE:360        case LLM_ARCH_HUNYUAN_MOE:361        case LLM_ARCH_HY_V3:362        case LLM_ARCH_OPENAI_MOE:363        case LLM_ARCH_LFM2MOE:364        case LLM_ARCH_SMALLTHINKER:365        case LLM_ARCH_LLADA_MOE:366        case LLM_ARCH_GROVEMOE:367        case LLM_ARCH_MINIMAX_M2:368        case LLM_ARCH_MINIMAX_M3:369        case LLM_ARCH_RND1:370        case LLM_ARCH_PADDLEOCR:371        case LLM_ARCH_MIMO2:372        case LLM_ARCH_KIMI_LINEAR:373        case LLM_ARCH_STEP35:374        case LLM_ARCH_MISTRAL4:375        case LLM_ARCH_MELLUM:376        case LLM_ARCH_LAGUNA:377            return true;378        default:379            return false;380    }381}382 383static bool moe_implemented(const llm_arch arch) {384    if (moe_mandatory(arch)) {385        return true;386    }387    switch (arch) {388        case LLM_ARCH_LLAMA:389        case LLM_ARCH_REFACT:390        case LLM_ARCH_MINICPM:391        case LLM_ARCH_GRANITE:392        case LLM_ARCH_GRANITE_MOE:393        case LLM_ARCH_MISTRAL3:394        case LLM_ARCH_LLAMA_EMBED:395            return true;396        default:397            return false;398    }399}400 401static bool arch_supported(const llm_arch arch) {402    if (arch == LLM_ARCH_CLIP || arch == LLM_ARCH_GPTJ || arch == LLM_ARCH_UNKNOWN) {403        return false; // These models don't have usable implementations.404    }405    if (arch == LLM_ARCH_CHAMELEON) {406        return false; // Only half-implemented and to be removed in the future.407    }408    if (arch == LLM_ARCH_WAVTOKENIZER_DEC) {409        return false; // FIXME CUDA backend crashes.410    }411    if (arch == LLM_ARCH_GEMMA4 || arch == LLM_ARCH_GEMMA4_ASSISTANT) {412        return false; // FIXME @ngxson413    }414    if (arch == LLM_ARCH_GRANITE_SWITCH) {415        return false; // FIXME adapter fixture416    }417    if (arch == LLM_ARCH_LLAMA_EMBED || arch == LLM_ARCH_GEMMA_EMBEDDING || arch == LLM_ARCH_T5ENCODER) {418        return false; // FIXME Embedding (?) models produce inconsistent results.419    }420    if (arch == LLM_ARCH_RWKV6 || arch == LLM_ARCH_RWKV6QWEN2 || arch == LLM_ARCH_RWKV7 || arch == LLM_ARCH_ARWKV7) {421        return false; // FIXME RWKV models hang indefinitely.422    }423    if (arch == LLM_ARCH_BERT || arch == LLM_ARCH_MODERN_BERT || arch == LLM_ARCH_NOMIC_BERT || arch == LLM_ARCH_NOMIC_BERT_MOE ||424            arch == LLM_ARCH_NEO_BERT || arch == LLM_ARCH_JINA_BERT_V2 || arch == LLM_ARCH_JINA_BERT_V3 || arch == LLM_ARCH_EUROBERT) {425        return false; // TODO vocab426    }427    if (arch == LLM_ARCH_PLM) {428        return false; // TODO tensor shapes429    }430    if (arch == LLM_ARCH_DEEPSEEK2OCR) {431        return false;432    }433    if (arch == LLM_ARCH_DEEPSEEK4) {434        return false;435    }436 437    // FIXME: these hit scheduler/view-backed-output issues with WebGPU on CI.438#ifdef GGML_USE_WEBGPU439    if (arch == LLM_ARCH_DEEPSEEK32 || arch == LLM_ARCH_GLM_DSA) {440        return false;441    }442#endif // GGML_USE_WEBGPU443 444    // FIXME: jamba produces incorrect output (~0.55 NMSE vs CPU) on the HIP445    // backend on RDNA3.5 (gfx1151); the SSM kernels need investigation.446#ifdef GGML_USE_HIP447    if (arch == LLM_ARCH_JAMBA) {448        return false;449    }450#endif // GGML_USE_HIP451 452    return true;453}454 455static int save_models(const llm_arch target_arch, const size_t seed, const ggml_log_level log_level, const std::string & dir) {456    struct user_data_t {457        struct {458            ggml_log_callback callback;459            void * user_data;460        } original_logger;461        ggml_log_level min_level; // prints below this log level go to debug log462    };463    user_data_t ud;464    llama_log_get(&ud.original_logger.callback, &ud.original_logger.user_data);465    ud.min_level = log_level;466 467    llama_log_set([](ggml_log_level level, const char * text, void * user_data) {468        const user_data_t * ud = (const user_data_t *) user_data;469        const ggml_log_level level_eff = level >= ud->min_level ? level : GGML_LOG_LEVEL_DEBUG;470        ud->original_logger.callback(level_eff, text, ud->original_logger.user_data);471    }, &ud);472 473    for (const llm_arch & arch : llm_arch_all()) {474        if (arch == LLM_ARCH_UNKNOWN) {475            continue;476        }477        if (target_arch != LLM_ARCH_UNKNOWN && arch != target_arch) {478            continue;479        }480        if (arch == LLM_ARCH_GEMMA4 || arch == LLM_ARCH_GEMMA4_ASSISTANT) {481            continue; // FIXME: ISWA KV cache initialization needs more fixture params482        }483        if (arch == LLM_ARCH_EAGLE3 || arch == LLM_ARCH_DFLASH) {484            continue;485        }486        for (bool moe : {false, true}) {487            if (moe && !moe_implemented(arch)) {488                continue;489            }490            if (!moe && moe_mandatory(arch)) {491                continue;492            }493            if (!llama_model_saver_supports_arch(arch) || !arch_supported(arch)) {494                LOG_INF("%s: %s model (%s) is unsupported, skipping\n", __func__, llm_arch_name(arch), moe ? "MoE" : "dense");495                continue;496            }497            gguf_context_ptr gguf_ctx = get_gguf_ctx(arch, moe);498            auto model_and_ctx = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, {});499            const std::string path = dir + "/" + llm_arch_name(arch) + (moe ? "-moe.gguf" : "-dense.gguf");500            LOG_INF("%s: Saving %s model (%s) to %s...\n", __func__, llm_arch_name(arch), moe ? "MoE" : "dense", path.c_str());501            llama_model_save_to_file(model_and_ctx.first.get(), path.c_str());502        }503    }504    llama_log_set(ud.original_logger.callback, ud.original_logger.user_data);505    return 0;506}507 508static int test_backends(const llm_arch target_arch, const size_t seed, const ggml_log_level log_level) {509    struct user_data_t {510        struct {511            ggml_log_callback callback;512            void * user_data;513        } original_logger;514        ggml_log_level min_level; // prints below this log level go to debug log515    };516    user_data_t ud;517    llama_log_get(&ud.original_logger.callback, &ud.original_logger.user_data);518    ud.min_level = log_level;519 520    llama_log_set([](ggml_log_level level, const char * text, void * user_data) {521        const user_data_t * ud = (const user_data_t *) user_data;522        const ggml_log_level level_eff = level >= ud->min_level ? level : GGML_LOG_LEVEL_DEBUG;523        ud->original_logger.callback(level_eff, text, ud->original_logger.user_data);524    }, &ud);525 526    const std::vector<llama_token> tokens = get_tokens(128, 128, seed);527 528    struct device_config {529        std::vector<ggml_backend_dev_t> devs;530        std::string                     label;531        llama_split_mode                split_mode;532 533        device_config(std::vector<ggml_backend_dev_t> devs, std::string name, llama_split_mode split_mode)534            : devs(std::move(devs)), label(std::move(name)), split_mode(split_mode) {}535    };536 537    std::vector<device_config> dev_configs;538    size_t max_device_label_length = 4;539    {540        std::vector<ggml_backend_dev_t> devices_meta;541        {542            const size_t device_count = ggml_backend_dev_count();543            for (size_t i = 0; i < device_count; i++) {544                ggml_backend_dev_t dev = ggml_backend_dev_get(i);545                dev_configs.emplace_back(std::vector<ggml_backend_dev_t>{dev}, ggml_backend_dev_description(dev), LLAMA_SPLIT_MODE_LAYER);546                max_device_label_length = std::max(max_device_label_length, dev_configs.back().label.length());547 548                // cpu-based devices cannot be used in tensor split mode549                if (ggml_backend_dev_buffer_type(dev) != ggml_backend_cpu_buffer_type()) {550                    devices_meta.push_back(dev);551                }552            }553        }554 555        dev_configs.emplace_back(devices_meta, "Meta", LLAMA_SPLIT_MODE_TENSOR);556    }557 558    size_t max_arch_name_length = 0;559    for (const llm_arch & arch : llm_arch_all()) {560        max_arch_name_length = std::max(max_arch_name_length, strlen(llm_arch_name(arch)));561    }562 563    const std::string template_header  = std::string("|%" + std::to_string(max_arch_name_length) + "s|%") + std::to_string(max_device_label_length) + "s|%6s|%15s|%9s|\n";564    const std::string template_row_cfg = std::string("|%" + std::to_string(max_arch_name_length) + "s|%") + std::to_string(max_device_label_length) + "s|%6s|";565    const std::string template_row_res = "%15s %10s|%20s|\n";566 567    bool all_ok = true;568    common_log_flush(common_log_main());569    printf(template_header.c_str(), "Model arch.", "Device", "Config", "NMSE vs. CPU", "Roundtrip");570    printf("|");571    for (size_t i = 0; i < max_arch_name_length; i++) {572        printf("-");573    }574    printf("|");575    for (size_t i = 0; i < max_device_label_length; i++) {576        printf("-");577    }578    printf("|------|---------------|---------|\n");579    for (const llm_arch & arch : llm_arch_all()) {580        if (arch == LLM_ARCH_UNKNOWN) {581            continue;582        }583        if (target_arch != LLM_ARCH_UNKNOWN && arch != target_arch) {584            continue;585        }586        if (arch == LLM_ARCH_GEMMA4 || arch == LLM_ARCH_GEMMA4_ASSISTANT) {587            continue; // FIXME: ISWA KV cache initialization needs more fixture params588        }589        if (arch == LLM_ARCH_EAGLE3 || arch == LLM_ARCH_DFLASH) {590            continue;591        }592 593        const bool encode = arch == LLM_ARCH_T5 || arch == LLM_ARCH_DREAM || arch == LLM_ARCH_LLADA || arch == LLM_ARCH_LLADA_MOE || arch == LLM_ARCH_RND1;594        for (bool moe : {false, true}) {595            if (moe && !moe_implemented(arch)) {596                continue;597            }598            if (!moe && moe_mandatory(arch)) {599                continue;600            }601            const std::string config_name = moe ? "MoE" : "Dense";602            gguf_context_ptr gguf_ctx = get_gguf_ctx(arch, moe);603            std::pair<llama_model_ptr, llama_context_ptr> model_and_ctx_cpu;604            std::vector<float> logits_cpu;605            for (device_config & dc : dev_configs) {606                // print test config first; should anything fail during model loading or inference, at least we know which test case caused it607                printf(template_row_cfg.c_str(),608                    llm_arch_name(arch), dc.label.c_str(), config_name.c_str());609                fflush(stdout);610 611                std::pair<llama_model_ptr, llama_context_ptr> model_and_ctx_dev;612                std::vector<float> logits_dev;613                std::string status_nmse      = "\033[1;33mSKIP\033[0m";614                std::string status_roundtrip = "\033[1;33mSKIP\033[0m";615                char nmse_str[12] = {0};616                bool skip = !arch_supported(arch) || (dc.split_mode == LLAMA_SPLIT_MODE_TENSOR && dc.devs.empty());617                if (!skip) {618                    if (logits_cpu.empty()) {619                        model_and_ctx_cpu = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, {}, LLAMA_SPLIT_MODE_LAYER, encode);620                        logits_cpu = get_logits(model_and_ctx_cpu.first.get(), model_and_ctx_cpu.second.get(), tokens, encode);621                    }622                    if (dc.split_mode != LLAMA_SPLIT_MODE_TENSOR || llm_arch_supports_sm_tensor(arch)) {623                        model_and_ctx_dev = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, dc.devs, dc.split_mode, encode);624                        logits_dev = get_logits(model_and_ctx_dev.first.get(), model_and_ctx_dev.second.get(), tokens, encode);625                        const double nmse_val = nmse(logits_cpu, logits_dev);626                        snprintf(nmse_str, sizeof(nmse_str), "(%.2e)", nmse_val);627                        status_nmse = "\033[1;32mOK\033[0m";628                        if (nmse_val > 1e-4) {629                            all_ok = false;630                            status_nmse = "\033[1;31mFAIL\033[0m";631                        }632                    }633 634                    FILE * file = tmpfile(); // Can be null on Windows without administrator privileges.635                    // FIXME: when adding a tensor to a gguf_context a copy is made, this changes the pointer which the meta backend636                    //     in turn uses to map the tensors to their simple equivalents - this is fundamentally incompatible637                    if (file != nullptr && llama_model_saver_supports_arch(arch) && dc.split_mode != LLAMA_SPLIT_MODE_TENSOR) {638                        GGML_ASSERT(model_and_ctx_dev.first && model_and_ctx_dev.second);639                        llama_model_saver ms = llama_model_saver(model_and_ctx_dev.first.get());640                        ms.add_kv_from_model();641                        ms.add_tensors_from_model();642                        ms.save(file);643                        rewind(file);644 645                        auto model_and_ctx_roundtrip = get_model_and_ctx(nullptr, file, seed, dc.devs, dc.split_mode, encode);646                        const std::vector<float> logits_roundtrip = get_logits(647                            model_and_ctx_roundtrip.first.get(), model_and_ctx_roundtrip.second.get(), tokens, encode);648                        status_roundtrip = "\033[1;32mOK\033[0m";649                        GGML_ASSERT(logits_roundtrip.size() == logits_dev.size());650                        for (size_t i = 0; i < logits_roundtrip.size(); i++) {651                            if (logits_roundtrip[i] != logits_dev[i]) {652                                all_ok = false;653                                status_roundtrip = "\033[1;31mFAIL\033[0m";654                                break;655                            }656                        }657                    }658                }659 660                // log the results for this test case661                printf(template_row_res.c_str(),662                    status_nmse.c_str(), nmse_str, status_roundtrip.c_str());663            }664        }665    }666    llama_log_set(ud.original_logger.callback, ud.original_logger.user_data);667    return all_ok ? 0 : 1;668}669 670int main(int argc, char ** argv) {671    // FIXME these tests are disabled in the CI for macOS-latest-cmake-arm64 because they are segfaulting672    common_init();673    std::random_device rd;674 675    llm_arch arch = LLM_ARCH_UNKNOWN;676    size_t seed = rd();677    ggml_log_level log_level = GGML_LOG_LEVEL_ERROR;678    std::string out;679 680    for (int i = 1; i < argc; i++) {681        if (strcmp(argv[i], "-a") == 0 || strcmp(argv[i], "--arch") == 0) {682            if (i + 1 < argc) {683                const std::string arch_name = argv[++i];684                arch = llm_arch_from_string(arch_name);685                if (arch == LLM_ARCH_UNKNOWN) {686                    LOG_ERR("%s: unkown LLM architecture: %s\n", __func__, arch_name.c_str());687                    return 1;688                }689            } else {690                usage(argv);691                return 1;692            }693        }694        if (strcmp(argv[i], "-s") == 0 || strcmp(argv[i], "--seed") == 0) {695            if (i + 1 < argc) {696                seed = std::stoull(argv[++i]);697            } else {698                usage(argv);699                return 1;700            }701        }702        if (strcmp(argv[i], "-v") == 0 || strcmp(argv[i], "--verbose") == 0) {703            log_level = GGML_LOG_LEVEL_INFO;704            continue;705        }706        if (strcmp(argv[i], "-o") == 0 || strcmp(argv[i], "--out") == 0) {707            if (i + 1 < argc) {708                out = argv[++i];709            } else {710                usage(argv);711                return 1;712            }713        }714    }715    printf("%s: using seed %zu\n", __func__, seed);716 717    try {718        if (!out.empty()) {719            return save_models(arch, seed, log_level, out);720        }721        return test_backends(arch, seed, log_level);722    } catch (const std::exception & err) {723        fprintf(stderr, "encountered runtime error: %s\n", err.what());724        return -1;725    }726}727