|
|
| |
|
|
| #include "llama.h" |
| #include "llama-adapter.h" |
| #include "llama-model.h" |
|
|
| #include "ggml.h" |
| #include "ggml-backend.h" |
| #include "gguf.h" |
|
|
| #include <algorithm> |
| #include <cmath> |
| #include <cstdint> |
| #include <cstdlib> |
| #include <cstring> |
| #include <filesystem> |
| #include <iomanip> |
| #include <iostream> |
| #include <stdexcept> |
| #include <string> |
| #include <vector> |
|
|
|
|
| static constexpr int64_t LORA_RANK = 4; |
| static constexpr float LORA_ALPHA = 8.0f; |
|
|
| static constexpr const char * SSM_TARGET = |
| "blk.0.ssm_alpha.weight"; |
|
|
| static constexpr const char * ATTN_TARGET = |
| "blk.11.attn_k.weight"; |
|
|
|
|
| static void require( |
| bool condition, |
| const std::string & message) { |
| if (!condition) { |
| throw std::runtime_error(message); |
| } |
| } |
|
|
|
|
| static std::string lowercase( |
| std::string value) { |
| std::transform( |
| value.begin(), |
| value.end(), |
| value.begin(), |
| [] (unsigned char character) { |
| return static_cast<char>( |
| std::tolower(character)); |
| }); |
|
|
| return value; |
| } |
|
|
|
|
| static float deterministic_signed( |
| size_t index, |
| uint32_t salt) { |
| uint32_t value = |
| static_cast<uint32_t>(index) |
| ^ salt |
| ^ 0x9e3779b9u; |
|
|
| value ^= value >> 16; |
| value *= 0x7feb352du; |
| value ^= value >> 15; |
| value *= 0x846ca68bu; |
| value ^= value >> 16; |
|
|
| const float unit = |
| static_cast<float>( |
| value & 0x00ffffffu) |
| / 16777215.0f; |
|
|
| return 2.0f * unit - 1.0f; |
| } |
|
|
|
|
| struct base_shape { |
| int64_t k = 0; |
| int64_t m = 0; |
| ggml_type type = GGML_TYPE_COUNT; |
| }; |
|
|
|
|
| static base_shape read_base_shape( |
| const std::string & model_path, |
| const std::string & tensor_name) { |
| ggml_context * metadata_context = |
| nullptr; |
|
|
| gguf_init_params params = { |
| true, |
| &metadata_context, |
| }; |
|
|
| gguf_context * gguf = |
| gguf_init_from_file( |
| model_path.c_str(), |
| params); |
|
|
| require( |
| gguf != nullptr, |
| "Could not open model GGUF metadata."); |
|
|
| require( |
| metadata_context != nullptr, |
| "GGUF metadata context is null."); |
|
|
| const int64_t tensor_id = |
| gguf_find_tensor( |
| gguf, |
| tensor_name.c_str()); |
|
|
| require( |
| tensor_id >= 0, |
| "Target tensor does not exist: " |
| + tensor_name); |
|
|
| ggml_tensor * tensor = |
| ggml_get_tensor( |
| metadata_context, |
| tensor_name.c_str()); |
|
|
| require( |
| tensor != nullptr, |
| "Target metadata tensor is null."); |
|
|
| require( |
| ggml_n_dims(tensor) == 2, |
| "Target tensor is not two-dimensional."); |
|
|
| base_shape shape = { |
| tensor->ne[0], |
| tensor->ne[1], |
| gguf_get_tensor_type( |
| gguf, |
| tensor_id), |
| }; |
|
|
| gguf_free(gguf); |
| ggml_free(metadata_context); |
|
|
| return shape; |
| } |
|
|
|
|
| static void generate_adapter( |
| const std::string & model_path, |
| const std::string & adapter_path, |
| const std::string & tensor_name, |
| uint32_t salt) { |
| const base_shape shape = |
| read_base_shape( |
| model_path, |
| tensor_name); |
|
|
| require( |
| shape.type == GGML_TYPE_Q1_0, |
| "Target is not Q1_0: " |
| + tensor_name); |
|
|
| require( |
| shape.k > 0 |
| && shape.m > 0 |
| && shape.k % 128 == 0, |
| "Target has invalid Q1 dimensions."); |
|
|
| const size_t a_elements = |
| static_cast<size_t>( |
| shape.k * LORA_RANK); |
|
|
| const size_t b_elements = |
| static_cast<size_t>( |
| LORA_RANK * shape.m); |
|
|
| const size_t tensor_bytes = |
| ( |
| a_elements |
| + b_elements |
| ) |
| * sizeof(float); |
|
|
| ggml_init_params tensor_params = { |
| |
| tensor_bytes |
| + 32 * ggml_tensor_overhead() |
| + 1024 * 1024, |
|
|
| |
| nullptr, |
|
|
| |
| false, |
| }; |
|
|
| ggml_context * tensor_context = |
| ggml_init(tensor_params); |
|
|
| require( |
| tensor_context != nullptr, |
| "Could not create adapter tensor context."); |
|
|
| ggml_tensor * tensor_a = |
| ggml_new_tensor_2d( |
| tensor_context, |
| GGML_TYPE_F32, |
| shape.k, |
| LORA_RANK); |
|
|
| ggml_tensor * tensor_b = |
| ggml_new_tensor_2d( |
| tensor_context, |
| GGML_TYPE_F32, |
| LORA_RANK, |
| shape.m); |
|
|
| const std::string name_a = |
| tensor_name + ".lora_a"; |
|
|
| const std::string name_b = |
| tensor_name + ".lora_b"; |
|
|
| ggml_set_name( |
| tensor_a, |
| name_a.c_str()); |
|
|
| ggml_set_name( |
| tensor_b, |
| name_b.c_str()); |
|
|
| float * data_a = |
| static_cast<float *>( |
| tensor_a->data); |
|
|
| float * data_b = |
| static_cast<float *>( |
| tensor_b->data); |
|
|
| require( |
| data_a != nullptr |
| && data_b != nullptr, |
| "Adapter tensor data is null."); |
|
|
| const float a_scale = |
| 0.050f |
| / std::sqrt( |
| static_cast<float>(shape.k)); |
|
|
| const float b_scale = |
| 0.050f |
| / std::sqrt( |
| static_cast<float>(LORA_RANK)); |
|
|
| for (size_t index = 0; |
| index < a_elements; |
| ++index) { |
| data_a[index] = |
| a_scale |
| * deterministic_signed( |
| index, |
| salt); |
| } |
|
|
| for (size_t index = 0; |
| index < b_elements; |
| ++index) { |
| data_b[index] = |
| b_scale |
| * deterministic_signed( |
| index, |
| salt ^ 0x68bc21ebu); |
| } |
|
|
| gguf_context * adapter_gguf = |
| gguf_init_empty(); |
|
|
| require( |
| adapter_gguf != nullptr, |
| "Could not create adapter GGUF context."); |
|
|
| gguf_set_val_str( |
| adapter_gguf, |
| "general.type", |
| "adapter"); |
|
|
| gguf_set_val_str( |
| adapter_gguf, |
| "general.architecture", |
| "qwen35"); |
|
|
| gguf_set_val_str( |
| adapter_gguf, |
| "general.name", |
| "Prism native Q1 loader block test"); |
|
|
| gguf_set_val_str( |
| adapter_gguf, |
| "adapter.type", |
| "lora"); |
|
|
| gguf_set_val_f32( |
| adapter_gguf, |
| "adapter.lora.alpha", |
| LORA_ALPHA); |
|
|
| gguf_add_tensor( |
| adapter_gguf, |
| tensor_a); |
|
|
| gguf_add_tensor( |
| adapter_gguf, |
| tensor_b); |
|
|
| const bool written = |
| gguf_write_to_file( |
| adapter_gguf, |
| adapter_path.c_str(), |
| false); |
|
|
| require( |
| written, |
| "Could not write adapter GGUF."); |
|
|
| gguf_free(adapter_gguf); |
| ggml_free(tensor_context); |
|
|
| require( |
| std::filesystem::exists( |
| adapter_path), |
| "Adapter GGUF was not created."); |
|
|
| std::cout |
| << "ADAPTER_GENERATED=" |
| << adapter_path |
| << "\n"; |
|
|
| std::cout |
| << "ADAPTER_TARGET=" |
| << tensor_name |
| << "\n"; |
|
|
| std::cout |
| << "ADAPTER_K=" |
| << shape.k |
| << "\n"; |
|
|
| std::cout |
| << "ADAPTER_M=" |
| << shape.m |
| << "\n"; |
| } |
|
|
|
|
| static void validate_adapter( |
| const llama_model * model, |
| const llama_adapter_lora * adapter, |
| const std::string & target_name, |
| const std::string & label) { |
| require( |
| adapter != nullptr, |
| label + " adapter is null."); |
|
|
| const auto iterator = |
| adapter->ab_map.find( |
| target_name); |
|
|
| require( |
| iterator |
| != adapter->ab_map.end(), |
| label + " adapter target is missing."); |
|
|
| const llama_adapter_lora_weight & pair = |
| iterator->second; |
|
|
| require( |
| pair.a != nullptr |
| && pair.b != nullptr, |
| label + " LoRA pair is incomplete."); |
|
|
| require( |
| pair.a->flags |
| & GGML_TENSOR_FLAG_PARAM, |
| label + " LoRA A is not a parameter."); |
|
|
| require( |
| pair.b->flags |
| & GGML_TENSOR_FLAG_PARAM, |
| label + " LoRA B is not a parameter."); |
|
|
| const ggml_tensor * base = |
| model->get_tensor( |
| target_name.c_str()); |
|
|
| require( |
| base != nullptr, |
| label + " base tensor is missing."); |
|
|
| require( |
| base->type == GGML_TYPE_Q1_0, |
| label + " base tensor is not Q1_0."); |
|
|
| require( |
| !( |
| base->flags |
| & GGML_TENSOR_FLAG_PARAM |
| ), |
| label + " base tensor became trainable."); |
|
|
| require( |
| pair.a->buffer != nullptr |
| && pair.b->buffer != nullptr, |
| label + " LoRA backend buffer is null."); |
|
|
| const std::string a_buffer = |
| ggml_backend_buffer_name( |
| pair.a->buffer); |
|
|
| const std::string b_buffer = |
| ggml_backend_buffer_name( |
| pair.b->buffer); |
|
|
| require( |
| lowercase(a_buffer).find("cuda") |
| != std::string::npos, |
| label + " LoRA A is not on CUDA: " |
| + a_buffer); |
|
|
| require( |
| lowercase(b_buffer).find("cuda") |
| != std::string::npos, |
| label + " LoRA B is not on CUDA: " |
| + b_buffer); |
|
|
| std::cout |
| << label |
| << "_ADAPTER_PARAMETER_COUNT=2\n"; |
|
|
| std::cout |
| << label |
| << "_LORA_A_PARAM=1\n"; |
|
|
| std::cout |
| << label |
| << "_LORA_B_PARAM=1\n"; |
|
|
| std::cout |
| << label |
| << "_BASE_PARAM=0\n"; |
|
|
| std::cout |
| << label |
| << "_BASE_TYPE=Q1_0\n"; |
|
|
| std::cout |
| << label |
| << "_LORA_A_BUFFER=" |
| << a_buffer |
| << "\n"; |
|
|
| std::cout |
| << label |
| << "_LORA_B_BUFFER=" |
| << b_buffer |
| << "\n"; |
| } |
|
|
|
|
| static std::vector<llama_token> tokenize( |
| const llama_model * model, |
| const std::string & text) { |
| const llama_vocab * vocab = |
| llama_model_get_vocab(model); |
|
|
| require( |
| vocab != nullptr, |
| "Vocabulary is null."); |
|
|
| const int required = |
| -llama_tokenize( |
| vocab, |
| text.c_str(), |
| text.size(), |
| nullptr, |
| 0, |
| true, |
| true); |
|
|
| require( |
| required > 0, |
| "Could not calculate token count."); |
|
|
| std::vector<llama_token> tokens( |
| static_cast<size_t>(required)); |
|
|
| const int written = |
| llama_tokenize( |
| vocab, |
| text.c_str(), |
| text.size(), |
| tokens.data(), |
| tokens.size(), |
| true, |
| true); |
|
|
| require( |
| written == required, |
| "Tokenization failed."); |
|
|
| return tokens; |
| } |
|
|
|
|
| static std::vector<float> evaluate_logits( |
| llama_model * model, |
| llama_adapter_lora * adapter, |
| const std::vector<llama_token> & tokens) { |
| llama_context_params context_params = |
| llama_context_default_params(); |
|
|
| context_params.n_ctx = 32; |
| context_params.n_batch = 16; |
| context_params.n_ubatch = 16; |
| context_params.flash_attn_type = |
| LLAMA_FLASH_ATTN_TYPE_DISABLED; |
| context_params.no_perf = true; |
|
|
| llama_context * context = |
| llama_init_from_model( |
| model, |
| context_params); |
|
|
| require( |
| context != nullptr, |
| "Could not create llama context."); |
|
|
| if (adapter != nullptr) { |
| llama_adapter_lora * adapters[] = { |
| adapter, |
| }; |
|
|
| float scales[] = { |
| 1.0f, |
| }; |
|
|
| const int set_result = |
| llama_set_adapters_lora( |
| context, |
| adapters, |
| 1, |
| scales); |
|
|
| require( |
| set_result == 0, |
| "Could not attach LoRA adapter."); |
| } |
|
|
| llama_batch batch = |
| llama_batch_get_one( |
| const_cast<llama_token *>( |
| tokens.data()), |
| static_cast<int32_t>( |
| tokens.size())); |
|
|
| const int decode_result = |
| llama_decode( |
| context, |
| batch); |
|
|
| require( |
| decode_result == 0, |
| "llama_decode failed: " |
| + std::to_string(decode_result)); |
|
|
| float * logits = |
| llama_get_logits_ith( |
| context, |
| -1); |
|
|
| require( |
| logits != nullptr, |
| "Final logits are null."); |
|
|
| const int64_t n_vocab = |
| model->vocab.n_tokens(); |
|
|
| require( |
| n_vocab > 0, |
| "Invalid vocabulary size."); |
|
|
| std::vector<float> result( |
| logits, |
| logits + n_vocab); |
|
|
| llama_free(context); |
|
|
| return result; |
| } |
|
|
|
|
| struct difference_metrics { |
| double max_abs = 0.0; |
| double mean_abs = 0.0; |
| double l2 = 0.0; |
| double checksum = 0.0; |
| }; |
|
|
|
|
| static difference_metrics compare_logits( |
| const std::vector<float> & baseline, |
| const std::vector<float> & adapted) { |
| require( |
| baseline.size() == adapted.size(), |
| "Logit vector size mismatch."); |
|
|
| difference_metrics metrics; |
|
|
| for (size_t index = 0; |
| index < baseline.size(); |
| ++index) { |
| require( |
| std::isfinite(baseline[index]) |
| && std::isfinite(adapted[index]), |
| "Non-finite model logit."); |
|
|
| const double difference = |
| static_cast<double>( |
| adapted[index]) |
| - static_cast<double>( |
| baseline[index]); |
|
|
| const double absolute = |
| std::abs(difference); |
|
|
| metrics.max_abs = |
| std::max( |
| metrics.max_abs, |
| absolute); |
|
|
| metrics.mean_abs += absolute; |
| metrics.l2 += difference * difference; |
|
|
| metrics.checksum += |
| difference |
| * static_cast<double>( |
| (index % 1009) + 1); |
| } |
|
|
| metrics.mean_abs /= |
| static_cast<double>( |
| baseline.size()); |
|
|
| metrics.l2 = |
| std::sqrt(metrics.l2); |
|
|
| return metrics; |
| } |
|
|
|
|
| int main( |
| int argc, |
| char ** argv) { |
| try { |
| require( |
| argc == 4, |
| "Usage: test-q1-lora-loader-blocks " |
| "MODEL SSM_ADAPTER ATTN_ADAPTER"); |
|
|
| #if defined(_WIN32) |
| _putenv_s( |
| "PRISM_Q1_LORA_TRAINING", |
| "1"); |
| #else |
| setenv( |
| "PRISM_Q1_LORA_TRAINING", |
| "1", |
| 1); |
| #endif |
|
|
| const std::string model_path = |
| argv[1]; |
|
|
| const std::string ssm_adapter_path = |
| argv[2]; |
|
|
| const std::string attn_adapter_path = |
| argv[3]; |
|
|
| std::cout |
| << std::fixed |
| << std::setprecision(12); |
|
|
| generate_adapter( |
| model_path, |
| ssm_adapter_path, |
| SSM_TARGET, |
| 0x10293847u); |
|
|
| generate_adapter( |
| model_path, |
| attn_adapter_path, |
| ATTN_TARGET, |
| 0xabcdef12u); |
|
|
| ggml_backend_load_all(); |
|
|
| llama_model_params model_params = |
| llama_model_default_params(); |
|
|
| model_params.n_gpu_layers = 999; |
| model_params.use_mmap = true; |
|
|
| llama_model * model = |
| llama_model_load_from_file( |
| model_path.c_str(), |
| model_params); |
|
|
| require( |
| model != nullptr, |
| "Could not load Bonsai model."); |
|
|
| require( |
| model->arch_name() == "qwen35", |
| "Loaded model is not qwen35."); |
|
|
| require( |
| model->hparams.n_layer() == 64, |
| "Expected 64 main blocks."); |
|
|
| require( |
| model->hparams.is_recr(0), |
| "Block 0 is not recurrent."); |
|
|
| require( |
| !model->hparams.is_recr(11), |
| "Block 11 is not full attention."); |
|
|
| require( |
| model->layers[0].ssm_alpha |
| == model->get_tensor( |
| SSM_TARGET), |
| "Block 0 SSM loader mapping mismatch."); |
|
|
| require( |
| model->layers[11].wk |
| == model->get_tensor( |
| ATTN_TARGET), |
| "Block 11 attention-K loader mapping mismatch."); |
|
|
| std::cout |
| << "MODEL_ARCHITECTURE=" |
| << model->arch_name() |
| << "\n"; |
|
|
| std::cout |
| << "MODEL_MAIN_BLOCK_COUNT=" |
| << model->hparams.n_layer() |
| << "\n"; |
|
|
| std::cout |
| << "BLOCK_SSM_INDEX=0\n"; |
|
|
| std::cout |
| << "BLOCK_SSM_RECURRENT=1\n"; |
|
|
| std::cout |
| << "BLOCK_SSM_TARGET=" |
| << SSM_TARGET |
| << "\n"; |
|
|
| std::cout |
| << "BLOCK_ATTN_INDEX=11\n"; |
|
|
| std::cout |
| << "BLOCK_ATTN_RECURRENT=0\n"; |
|
|
| std::cout |
| << "BLOCK_ATTN_TARGET=" |
| << ATTN_TARGET |
| << "\n"; |
|
|
| llama_adapter_lora * ssm_adapter = |
| llama_adapter_lora_init( |
| model, |
| ssm_adapter_path.c_str()); |
|
|
| llama_adapter_lora * attn_adapter = |
| llama_adapter_lora_init( |
| model, |
| attn_adapter_path.c_str()); |
|
|
| require( |
| ssm_adapter != nullptr, |
| "Could not load SSM adapter."); |
|
|
| require( |
| attn_adapter != nullptr, |
| "Could not load attention adapter."); |
|
|
| validate_adapter( |
| model, |
| ssm_adapter, |
| SSM_TARGET, |
| "SSM"); |
|
|
| validate_adapter( |
| model, |
| attn_adapter, |
| ATTN_TARGET, |
| "ATTN"); |
|
|
| const std::vector<llama_token> tokens = |
| tokenize( |
| model, |
| "One bit LoRA test."); |
|
|
| std::cout |
| << "MODEL_PROMPT_TOKEN_COUNT=" |
| << tokens.size() |
| << "\n"; |
|
|
| std::cout |
| << "BASE_FORWARD_BEGIN=1\n"; |
|
|
| const std::vector<float> base_logits = |
| evaluate_logits( |
| model, |
| nullptr, |
| tokens); |
|
|
| std::cout |
| << "BASE_FORWARD_STATUS=PASS\n"; |
|
|
| std::cout |
| << "SSM_FORWARD_BEGIN=1\n"; |
|
|
| const std::vector<float> ssm_logits = |
| evaluate_logits( |
| model, |
| ssm_adapter, |
| tokens); |
|
|
| std::cout |
| << "SSM_FORWARD_STATUS=PASS\n"; |
|
|
| std::cout |
| << "ATTN_FORWARD_BEGIN=1\n"; |
|
|
| const std::vector<float> attn_logits = |
| evaluate_logits( |
| model, |
| attn_adapter, |
| tokens); |
|
|
| std::cout |
| << "ATTN_FORWARD_STATUS=PASS\n"; |
|
|
| const difference_metrics ssm_difference = |
| compare_logits( |
| base_logits, |
| ssm_logits); |
|
|
| const difference_metrics attn_difference = |
| compare_logits( |
| base_logits, |
| attn_logits); |
|
|
| std::cout |
| << "SSM_LOGITS_MAX_ABS_DIFF=" |
| << ssm_difference.max_abs |
| << "\n"; |
|
|
| std::cout |
| << "SSM_LOGITS_MEAN_ABS_DIFF=" |
| << ssm_difference.mean_abs |
| << "\n"; |
|
|
| std::cout |
| << "SSM_LOGITS_L2_DIFF=" |
| << ssm_difference.l2 |
| << "\n"; |
|
|
| std::cout |
| << "SSM_LOGITS_DIFF_CHECKSUM=" |
| << ssm_difference.checksum |
| << "\n"; |
|
|
| std::cout |
| << "ATTN_LOGITS_MAX_ABS_DIFF=" |
| << attn_difference.max_abs |
| << "\n"; |
|
|
| std::cout |
| << "ATTN_LOGITS_MEAN_ABS_DIFF=" |
| << attn_difference.mean_abs |
| << "\n"; |
|
|
| std::cout |
| << "ATTN_LOGITS_L2_DIFF=" |
| << attn_difference.l2 |
| << "\n"; |
|
|
| std::cout |
| << "ATTN_LOGITS_DIFF_CHECKSUM=" |
| << attn_difference.checksum |
| << "\n"; |
|
|
| require( |
| ssm_difference.max_abs > 1.0e-8, |
| "SSM adapter did not change logits."); |
|
|
| require( |
| ssm_difference.l2 > 1.0e-8, |
| "SSM logit L2 difference is zero."); |
|
|
| require( |
| attn_difference.max_abs > 1.0e-8, |
| "Attention adapter did not change logits."); |
|
|
| require( |
| attn_difference.l2 > 1.0e-8, |
| "Attention logit L2 difference is zero."); |
|
|
| require( |
| std::abs( |
| ssm_difference.checksum) |
| > 1.0e-10, |
| "SSM difference checksum is zero."); |
|
|
| require( |
| std::abs( |
| attn_difference.checksum) |
| > 1.0e-10, |
| "Attention difference checksum is zero."); |
|
|
| std::cout |
| << "CHECK_REAL_QWEN35_LOADER_MAPPING=PASS\n"; |
|
|
| std::cout |
| << "CHECK_SSM_BLOCK_FULL_FORWARD=PASS\n"; |
|
|
| std::cout |
| << "CHECK_ATTN_BLOCK_FULL_FORWARD=PASS\n"; |
|
|
| std::cout |
| << "CHECK_SSM_LORA_PARAMETERS=PASS\n"; |
|
|
| std::cout |
| << "CHECK_ATTN_LORA_PARAMETERS=PASS\n"; |
|
|
| std::cout |
| << "CHECK_LORA_PARAMETERS_ON_CUDA=PASS\n"; |
|
|
| std::cout |
| << "CHECK_BASE_Q1_FROZEN=PASS\n"; |
|
|
| std::cout |
| << "CHECK_SSM_LOGITS_CHANGED=PASS\n"; |
|
|
| std::cout |
| << "CHECK_ATTN_LOGITS_CHANGED=PASS\n"; |
|
|
| std::cout |
| << "PERSISTENT_EXPANDED_WEIGHT_BYTES=0\n"; |
|
|
| std::cout |
| << "FINAL_STATUS=PASS\n"; |
|
|
| llama_adapter_lora_free( |
| ssm_adapter); |
|
|
| llama_adapter_lora_free( |
| attn_adapter); |
|
|
| llama_model_free(model); |
|
|
| return 0; |
|
|
| } catch ( |
| const std::exception & exception |
| ) { |
| std::cerr |
| << "FINAL_STATUS=FAIL\n"; |
|
|
| std::cerr |
| << "ERROR=" |
| << exception.what() |
| << "\n"; |
|
|
| return 1; |
| } |
| } |
|
|
