diff --git a/GGUF_ROUTE_B_QWEN38.md b/GGUF_ROUTE_B_QWEN38.md new file mode 100644 index 000000000..fae4ac2bc --- /dev/null +++ b/GGUF_ROUTE_B_QWEN38.md @@ -0,0 +1,191 @@ +# GGUF Route B for Qwen3.5 + +## Overview + +This integration runs selected GGUF block-quantized weights directly from an +InfiniLM checkpoint. The converter copies supported GGUF block bytes into +safetensors as `uint8` tensors and records their GGML type in +`quantization_config.ggml_types`. InfiniLM resolves each weight by checkpoint +name and dispatches it to InfiniCore's `linear_gguf` operator. + +The design keeps model-specific mapping in Python while making the C++ packed +Linear path reusable by other model integrations. + +## Supported scope + +- Model profile: Qwen3.5 / Qwen3.8 27B Route B mapping. +- Native packed Linear types: `Q8_0`, `Q4_K`, `Q5_K`, and `Q6_K`. +- Dense BF16 fallback for parameters without a native packed execution path, + including embeddings, output head, normalization/scalar tensors, and IQ4 + tensors. +- NVIDIA execution through InfiniCore `linear_gguf`. +- Decode/small-batch and prefill execution paths selected inside InfiniCore. + +The current implementation intentionally rejects tensor parallelism for packed +GGUF weights. It also does not provide native IQ4, packed embedding, or packed +output-head kernels. + +## Dependency + +The InfiniLM changes require the corresponding InfiniCore `linear_gguf` +operator and its supported GGML block decoders: + +- InfiniCore pull request: https://github.com/InfiniTensor/InfiniCore/pull/1545 + +Build and install that InfiniCore revision before building InfiniLM. + +## Checkpoint format + +Packed Linear weights use a `weight_bytes` suffix and shape +`[out_features, row_bytes]`, where: + +```text +row_bytes = in_features / block_size * type_size +``` + +The converter writes a top-level configuration: + +```json +{ + "quantization_config": { + "quant_method": "gguf", + "key_prefix": "model.language_model.", + "ggml_types": { + "model.language_model.layers.0.mlp.gate_proj.weight_bytes": 14, + "model.language_model.layers.0.input_layernorm.weight": "dense_bf16" + }, + "activation_vperm": [] + } +} +``` + +`ggml_types` keys are exact safetensors parameter names. Values are GGML type +ids or `"dense_bf16"`. Runtime lookup requires exactly one packed or dense +candidate and fails on missing or ambiguous metadata. + +## Conversion + +The converter depends on Python packages used by the project plus +`gguf-py`. Install `gguf-py` or point `LLAMA_CPP_DIR` at a llama.cpp +checkout: + +```bash +export LLAMA_CPP_DIR=/path/to/llama.cpp +python3 scripts/gguf_to_infinilm.py \ + --gguf /path/to/model.gguf \ + --out /path/to/infinilm-checkpoint \ + --tokenizer-dir /path/to/tokenizer-config \ + --verify sample +``` + +`--tokenizer-dir` is optional. Vocabulary and merges are exported from GGUF; +the directory only supplies auxiliary files such as +`tokenizer_config.json` and a chat template. + +Useful options: + +- `--dry-run`: validate metadata, shapes, orientation, and packed row sizes + without writing tensors. +- `--layers N`: create a loadable checkpoint containing the first N layers. +- `--verify {off,sample,all}`: control post-write verification. +- `--skip-pack`: keep existing tensor shards while refreshing configuration, + tokenizer files, and verification. +- `--emit-dense-ref PATH`: create a fully dequantized BF16 reference for + numerical comparison. + +The converter is intentionally fail-closed. A model whose dimensions differ +from the Qwen3.8 27B profile requires a new mapping profile instead of silently +reusing incompatible shapes. + +## Mapping and transforms + +`scripts/gguf_mapping.py` is the single source of truth for: + +- GGUF-to-InfiniLM parameter names and shapes; +- packed versus dense storage; +- fused tensor slices; +- conversion-time value-head permutations; +- runtime activation-permutation metadata; +- generated InfiniLM model configuration. + +`scripts/gguf_transforms.py` contains the shared NumPy transformations. GGUF +Qwen conversion stores selected value heads in tiled `[value][key]` order, +while InfiniLM uses grouped `[key][value]` order. Complete packed rows can be +permuted during conversion without modifying block bytes. + +Some output-projection transformations affect columns instead of rows. Moving +packed columns across quantization blocks would require requantization, so the +converter emits `activation_vperm` rules and the runtime applies the +equivalent grouped-to-tiled permutation to the input activation. + +## Runtime integration + +`GGUFBlockQuantization` provides: + +- name-aware parameter layout selection; +- exact GGML type resolution; +- independent buffers for fused Linear shards; +- per-shard dispatch when fused projections use different GGML types; +- dense BF16 execution through the regular Linear operator; +- packed execution through `linear_gguf`; +- validation for dtype, contiguity, block divisibility, bias, and unsupported + tensor-parallel configurations. + +Linear constructors pass a checkpoint stem to the quantization layer. Fused +projections retain one stem per shard so Q/K/V or gate/up components can resolve +different source types and concatenate their outputs in the original order. + +Qwen3.5 model changes supply these stems for attention, MLP, and gated-delta-net +projections. The Python remap also avoids applying a second normalization +`+1` adjustment because llama.cpp already bakes that offset into GGUF. + +## Validation performed + +The submitted branch has been validated with: + +- repository formatting checks; +- a successful InfiniLM extension build; +- the official single-request test; +- the official offline benchmark; +- a local fixed MMLU-format smoke test; +- the official service test with 64/64 successful requests; +- end-to-end Qwen3.8 27B packed-checkpoint loading and generation. + +Observed offline performance on the validation machine was approximately: + +```text +decode throughput: 5.33 tokens/s +prefill throughput: 6.1 tokens/s +time to first token: 10.49 s +``` + +These numbers establish functionality, not a portable performance claim. They +were not collected as a controlled comparison against llama.cpp with identical +prompts, context lengths, sampling, and device settings. + +## Known limitations + +- Packed GGUF tensor parallelism is not implemented. +- IQ4 tensors, embeddings, and the output head use dense BF16 fallbacks. +- Strict token-for-token agreement with llama.cpp is not guaranteed; numerical + comparisons are the appropriate correctness criterion for quantized kernels. +- A full external MMLU dataset run was unavailable on the validation machine; + only the local MMLU-format execution path was exercised. +- Upstream CI and maintainer review remain authoritative for merge readiness. + +## Extending Route B to another model + +Most C++ work is reusable. A new model integration should: + +1. Define a model-specific mapping profile with exact checkpoint names, logical + shapes, fused slices, and required transforms. +2. Generate exact `ggml_types` entries and any activation-permutation rules. +3. Pass checkpoint stems from each model Linear constructor. +4. Add native block types to InfiniCore only when the model uses unsupported + GGML formats; otherwise reuse `linear_gguf`. +5. Validate conversion with dry-run, exact key/shape/dtype checks, packed-row + byte preservation, a loadable small-layer checkpoint, and end-to-end output. +6. Benchmark correctness and performance separately with controlled settings. + +This separation keeps GGUF storage and dispatch generic while isolating +architecture-specific tensor naming and permutation rules in the converter. diff --git a/csrc/config/quant_config.cpp b/csrc/config/quant_config.cpp index e58966d89..7a25457b8 100644 --- a/csrc/config/quant_config.cpp +++ b/csrc/config/quant_config.cpp @@ -20,6 +20,12 @@ QuantConfig::get_quantization_method() const { return std::make_shared(quantization_config); } else if (quant_method == "gptq") { return std::make_shared(quantization_config); + } else if (quant_method == "fp8") { + return std::make_shared(quantization_config); + } else if (quant_method == "gguf") { + // Route B keeps the original GGUF block bytes on the device and + // delegates decoding and multiplication to InfiniCore. + return std::make_shared(quantization_config); } else if (quant_method == "quark") { return std::make_shared(quantization_config); } else { diff --git a/csrc/engine/rank_worker.cpp b/csrc/engine/rank_worker.cpp index 0fa0a84cf..7936f88d7 100644 --- a/csrc/engine/rank_worker.cpp +++ b/csrc/engine/rank_worker.cpp @@ -2,11 +2,48 @@ #include "../models/model_factory.hpp" #include "infinicore/ops.hpp" #include "infinicore/ops/distributed/send_recv.hpp" +#include +#include #include #include namespace infinilm::engine { +namespace { + +infinicore::Tensor negative_infinity_cpu(infinicore::DataType dtype) { + auto scalar = infinicore::Tensor::empty( + {1}, dtype, infinicore::Device::cpu()); + switch (dtype) { + case infinicore::DataType::F16: { + const uint16_t value = 0xfc00U; + std::memcpy(scalar->data(), &value, sizeof(value)); + break; + } + case infinicore::DataType::BF16: { + const uint16_t value = 0xff80U; + std::memcpy(scalar->data(), &value, sizeof(value)); + break; + } + case infinicore::DataType::F32: { + const uint32_t value = 0xff800000U; + std::memcpy(scalar->data(), &value, sizeof(value)); + break; + } + case infinicore::DataType::F64: { + const uint64_t value = 0xfff0000000000000ULL; + std::memcpy(scalar->data(), &value, sizeof(value)); + break; + } + default: + throw std::runtime_error( + "suppressed_token_ids requires floating-point logits"); + } + return scalar; +} + +} // namespace + RankWorker::RankWorker( std::shared_ptr infinilm_config, const distributed::RankInfo &rank_info, @@ -484,12 +521,45 @@ void RankWorker::thread_loop() { const size_t n_out = sample_all_positions ? static_cast(input_offsets[n_req]) : n_req; auto output_ids{infinicore::Tensor::empty({n_out}, infinicore::DataType::I64, rank_info_.device)}; + const auto &suppressed = local_args.suppressed_token_ids; + if (!suppressed.empty() && suppressed.size() != n_req) { + throw std::runtime_error( + "suppressed_token_ids must contain one list per request"); + } + infinicore::Tensor neg_inf_cpu; + infinicore::Tensor neg_inf_device; + if (!suppressed.empty()) { + neg_inf_cpu = negative_infinity_cpu(logits->dtype()); + neg_inf_device = neg_inf_cpu->to(rank_info_.device); + } + + size_t req_idx = 0; + for (size_t i{0}; i < n_out; ++i) { size_t score_idx = i; if (!sample_all_positions && !logits_are_last_token_only) { score_idx = static_cast(input_offsets[i + 1] - 1); } auto score{logits->view({logits_positions, vocab_size})->narrow({{0, score_idx, 1}})->view({vocab_size})}; + if (sample_all_positions) { + while (req_idx + 1 < n_req + && i >= static_cast(input_offsets[req_idx + 1])) { + ++req_idx; + } + } else { + req_idx = i; + } + if (!suppressed.empty()) { + for (const int64_t token_id : suppressed[req_idx]) { + if (token_id < 0 + || static_cast(token_id) >= vocab_size) { + throw std::runtime_error( + "suppressed token ID is outside the vocabulary"); + } + score->narrow({{0, static_cast(token_id), 1}}) + ->copy_from(neg_inf_device); + } + } auto out{output_ids->narrow({{0, i, 1}})->view({})}; float random_val = std::uniform_real_distribution(0, 1)(rng_); infinicore::op::random_sample_( diff --git a/csrc/engine/rank_worker.hpp b/csrc/engine/rank_worker.hpp index d396ef6f1..6a5830c0f 100644 --- a/csrc/engine/rank_worker.hpp +++ b/csrc/engine/rank_worker.hpp @@ -11,6 +11,7 @@ #include #include +#include #include #include #include @@ -73,6 +74,9 @@ class RankWorker { /// Sample logits at every packed input position instead of one token per request. bool sample_all_positions{false}; + /// Token IDs excluded from sampling for each request in the batch. + std::vector> suppressed_token_ids{}; + float temperature{1}; int top_k{50}; diff --git a/csrc/layers/causal_lm_templates/text_model.hpp b/csrc/layers/causal_lm_templates/text_model.hpp index 49b60d7f4..a979b913c 100644 --- a/csrc/layers/causal_lm_templates/text_model.hpp +++ b/csrc/layers/causal_lm_templates/text_model.hpp @@ -6,11 +6,15 @@ #include "infinicore/nn/embedding.hpp" #include "infinicore/nn/rmsnorm.hpp" #include "infinicore/ops.hpp" +#include "infinicore/ops/add_rms_norm.hpp" +#include "infinicore/ops/cast.hpp" #include "infinicore/ops/distributed/allgather.hpp" #include "infinicore/ops/distributed/send_recv.hpp" #include "infinicore/tensor.hpp" +#include #include #include +#include namespace infinilm::layers::causal_lm_templates { @@ -79,7 +83,8 @@ class TextModel : public infinicore::nn::Module { return hidden_states; } - norm_->forward_inplace(hidden_states, residual); + dump_pre_final_norm_if_requested(hidden_states, residual); + final_norm_inplace(hidden_states, residual); return hidden_states; } @@ -119,7 +124,8 @@ class TextModel : public infinicore::nn::Module { return hidden_states; } - norm_->forward_inplace(hidden_states, residual); + dump_pre_final_norm_if_requested(hidden_states, residual); + final_norm_inplace(hidden_states, residual); return hidden_states; } @@ -136,6 +142,22 @@ class TextModel : public infinicore::nn::Module { INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, norm); private: + void final_norm_inplace(infinicore::Tensor &hidden_states, + infinicore::Tensor &residual) const { + const char *env = std::getenv("INFINILM_FINAL_NORM_FP32_FUSED"); + const bool enabled = env != nullptr && env[0] != '\0' && std::string(env) != "0"; + if (!enabled) { + norm_->forward_inplace(hidden_states, residual); + return; + } + auto y32 = infinicore::Tensor::empty(hidden_states->shape(), infinicore::DataType::F32, hidden_states->device()); + auto sum32 = infinicore::Tensor::empty(residual->shape(), infinicore::DataType::F32, residual->device()); + infinicore::op::add_rms_norm_(y32, sum32, hidden_states, residual, norm_->weight(), + static_cast(norm_->eps())); + hidden_states = y32; + residual = sum32; + } + bool is_first_pp_stage() const { return pp_stage_ == 0; } bool is_last_pp_stage() const { return pp_stage_ + 1 == pp_size_; } @@ -206,6 +228,24 @@ class TextModel : public infinicore::nn::Module { return infinicore::op::add(residual, hidden_states); } + void dump_pre_final_norm_if_requested( + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual) const { + const char *dump_dir = std::getenv("INFINILM_FINAL_PRENORM_DUMP_DIR"); + if (dump_dir == nullptr || dump_dir[0] == '\0') { + return; + } + const char *dump_numel = std::getenv("INFINILM_FINAL_PRENORM_DUMP_NUMEL"); + if (dump_numel != nullptr && dump_numel[0] != '\0' + && hidden_states->numel() + != std::strtoull(dump_numel, nullptr, 10)) { + return; + } + auto pre_norm = materialize_hidden_states(hidden_states, residual); + pre_norm->debug( + std::string(dump_dir) + "/infini_pre_final_norm.bin"); + } + infinicore::DataType dtype_{infinicore::DataType::F32}; size_t hidden_size_{0}; size_t pp_size_{1}; diff --git a/csrc/layers/linear/base_linear.cpp b/csrc/layers/linear/base_linear.cpp index dc4c77f62..10c51363b 100644 --- a/csrc/layers/linear/base_linear.cpp +++ b/csrc/layers/linear/base_linear.cpp @@ -9,19 +9,24 @@ BaseLinear::BaseLinear(size_t in_features, size_t out_features, bool bias, const infinicore::DataType &dtype, const infinicore::Device &device, int split_dim, int tp_rank, int tp_size, - int tp_num_heads) + int tp_num_heads, const std::string &stem) : in_features_(in_features), out_features_(out_features), has_bias_(bias), dtype_(dtype), split_dim_(split_dim), + stem_(stem), quantization_(quantization) { device_ = device; auto layout = quantization_->get_param_layout( in_features, out_features, split_dim, tp_rank, tp_size, - tp_num_heads, dtype, bias); + tp_num_heads, dtype, bias, stem); + + // An empty layout marks a fused group whose shards are allocated by the + // derived class through init_fused_shards(). + sharded_ = layout.empty(); for (const auto &desc : layout) { infinicore::nn::Parameter param( @@ -33,6 +38,10 @@ BaseLinear::BaseLinear(size_t in_features, size_t out_features, } infinicore::Tensor BaseLinear::compute_linear(infinicore::Tensor &input) const { + if (sharded_ && parameters_.empty()) { + throw std::runtime_error( + "BaseLinear::compute_linear: fused quantization shards were not registered"); + } // Build params map from direct parameters only (not state_dict which uses a // static local and is not thread-safe across RankWorker threads). infinilm::quantization::ParamsMap params; @@ -40,7 +49,7 @@ infinicore::Tensor BaseLinear::compute_linear(infinicore::Tensor &input) const { params[name] = static_cast(param); } - return quantization_->forward(params, input, has_bias_, alpha_); + return quantization_->forward(params, input, has_bias_, alpha_, stem_, shard_stems_); } infinicore::Tensor BaseLinear::compute_linear_allreduce( @@ -161,4 +170,35 @@ std::vector BaseLinear::split_params( parameters_, splits, split_dim_, tp_rank, tp_size, tp_num_heads); } +std::vector BaseLinear::init_fused_shards( + const std::vector &shards) { + std::vector registered; + shard_stems_.clear(); + shard_stems_.reserve(shards.size()); + for (size_t i = 0; i < shards.size(); ++i) { + const auto &sh = shards[i]; + // The index is shared by the shard parameter key and shard_stems_. + shard_stems_.push_back(sh.stem); + // Each shard is a complete column-parallel parameter. GGUF currently + // supports tp_size == 1 and rejects tensor-parallel execution. + auto layout = quantization_->get_param_layout( + in_features_, sh.out_features, split_dim_, 0, 1, -1, dtype_, false, sh.stem); + if (layout.empty()) { + throw std::runtime_error( + "BaseLinear::init_fused_shards: shard '" + sh.stem + "' returned an empty layout"); + } + for (const auto &desc : layout) { + infinicore::nn::Parameter param( + desc.shape, desc.dtype, device_, desc.split_dim, 0, 1, 0); + // The "shard." prefix preserves concatenation order in forward(). + this->register_parameter( + std::string(infinilm::quantization::GGUFBlockQuantization::SHARD_PREFIX) + std::to_string(i) + "." + desc.name, + param); + registered.push_back({sh.name + "." + desc.name, std::move(param)}); + } + } + sharded_ = true; + return registered; +} + } // namespace infinilm::nn diff --git a/csrc/layers/linear/base_linear.hpp b/csrc/layers/linear/base_linear.hpp index 8b452544e..b8d742a88 100644 --- a/csrc/layers/linear/base_linear.hpp +++ b/csrc/layers/linear/base_linear.hpp @@ -18,7 +18,8 @@ class BaseLinear : public infinicore::nn::Module { const infinicore::DataType &dtype = infinicore::DataType::F32, const infinicore::Device &device = infinicore::Device(), int split_dim = -1, int tp_rank = 0, int tp_size = 1, - int tp_num_heads = -1); + int tp_num_heads = -1, + const std::string &stem = ""); // Forward pass: output = input @ weight.T + bias infinicore::Tensor forward(infinicore::Tensor &input) const; @@ -53,6 +54,23 @@ class BaseLinear : public infinicore::nn::Module { const std::vector &splits, int tp_rank, int tp_size, int tp_num_heads) const; + // One shard of a fused linear, for schemes that cannot share a single fused + // buffer (GGUF block quantization: row_bytes differs per shard type). + struct FusedShard { + std::string name; // Name used when registering with the parent module. + size_t out_features; // Logical output rows for this shard. + std::string stem; // Checkpoint stem used for quantization lookup. + }; + + // Allocate one buffer per fused Linear shard. Local parameter keys use + // "shard.", while returned names use "." for + // registration with the parent module. This path is selected when the + // quantization scheme returns an empty layout for the fused group. + // shard_stems_ preserves each checkpoint stem for per-shard lookup during + // forward execution. + std::vector init_fused_shards( + const std::vector &shards); + // Allow subclasses to access the raw parameters map const infinicore::nn::Parameter &get_parameter_ref(const std::string &name) const; @@ -67,6 +85,12 @@ class BaseLinear : public infinicore::nn::Module { infinicore::DataType dtype_; int split_dim_ = -1; float alpha_ = 1.0f; + std::string stem_; // Checkpoint tensor path used by name-based quantization lookup. + // Per-shard checkpoint stems recorded by init_fused_shards. The index + // matches the i in the corresponding "shard.*" parameter key. + // This vector is empty for non-fused Linear layers. + std::vector shard_stems_; + bool sharded_ = false; // Fused layout whose parameters live in shard.* buffers. std::shared_ptr quantization_; }; diff --git a/csrc/layers/linear/fused_linear.cpp b/csrc/layers/linear/fused_linear.cpp index 1bcb96c94..ee89783d6 100644 --- a/csrc/layers/linear/fused_linear.cpp +++ b/csrc/layers/linear/fused_linear.cpp @@ -20,7 +20,7 @@ QKVParallelLinear::QKVParallelLinear(size_t hidden_size, num_q_head, num_kv_head, num_kv_head, bias, bias, bias, quantization, - dtype, device, rank_info) {} + dtype, device, rank_info, "") {} QKVParallelLinear::QKVParallelLinear(size_t hidden_size, size_t q_dim, size_t k_dim, size_t v_dim, @@ -29,16 +29,19 @@ QKVParallelLinear::QKVParallelLinear(size_t hidden_size, std::shared_ptr quantization, const infinicore::DataType &dtype, const infinicore::Device &device, - engine::distributed::RankInfo rank_info) + engine::distributed::RankInfo rank_info, + const std::string &stem) : infinilm::nn::ColumnParallelLinear( - hidden_size, - calculate_out_feature_size(num_q_head, q_dim, num_k_head, k_dim, num_v_head, v_dim, rank_info), - quantization == nullptr ? std::make_shared() : quantization, - (q_bias || k_bias || v_bias), - dtype, - device, - rank_info.tp_rank, - rank_info.tp_size), + hidden_size, + calculate_out_feature_size(num_q_head, q_dim, num_k_head, k_dim, num_v_head, v_dim, rank_info), + quantization == nullptr ? std::make_shared() : quantization, + (q_bias || k_bias || v_bias), + dtype, + device, + rank_info.tp_rank, + rank_info.tp_size, + -1, + stem), q_dim_(q_dim), k_dim_(k_dim), v_dim_(v_dim), @@ -83,8 +86,9 @@ QKVParallelLinear::QKVParallelLinear(size_t hidden_size, bool bias, const infinicore::DataType &dtype, const infinicore::Device &device, - engine::distributed::RankInfo rank_info) - : QKVParallelLinear(hidden_size, head_dim, head_dim, head_dim, num_q_head, num_kv_head, num_kv_head, bias, bias, bias, q_name, k_name, v_name, register_fn, quantization, dtype, device, rank_info) { + engine::distributed::RankInfo rank_info, + const std::string &prefix) + : QKVParallelLinear(hidden_size, head_dim, head_dim, head_dim, num_q_head, num_kv_head, num_kv_head, bias, bias, bias, q_name, k_name, v_name, register_fn, quantization, dtype, device, rank_info, prefix) { } QKVParallelLinear::QKVParallelLinear(size_t hidden_size, @@ -96,15 +100,41 @@ QKVParallelLinear::QKVParallelLinear(size_t hidden_size, std::shared_ptr quantization, const infinicore::DataType &dtype, const infinicore::Device &device, - engine::distributed::RankInfo rank_info) - : QKVParallelLinear(hidden_size, q_dim, k_dim, v_dim, num_q_head, num_k_head, num_v_head, q_bias, k_bias, v_bias, quantization, dtype, device, rank_info) { + engine::distributed::RankInfo rank_info, + const std::string &prefix) + : QKVParallelLinear(hidden_size, q_dim, k_dim, v_dim, num_q_head, num_k_head, num_v_head, q_bias, k_bias, v_bias, quantization, dtype, device, rank_info, prefix) { register_fn_ = register_fn; - split_infos_ = { - {q_name, 0, q_out_size_, 0}, - {k_name, q_out_size_, k_out_size_, num_k_head_}, - {v_name, q_out_size_ + k_out_size_, v_out_size_, num_v_head_}, - }; - auto params = this->split_params(split_infos_, tp_rank_, tp_size_, num_k_head_); + if (this->sharded_) { + // GGUF Q, K, and V shards may use different GGML types, so there is + // no fused buffer to narrow. Each shard owns a separate buffer and + // its stem identifies the corresponding checkpoint tensor. + if (prefix.empty()) { + throw std::runtime_error( + "QKVParallelLinear requires `prefix` when the quantization " + "scheme resolves layouts by checkpoint tensor name."); + } + shard_specs_ = { + {q_name, q_out_size_, prefix + "." + q_name + "."}, + {k_name, k_out_size_, prefix + "." + k_name + "."}, + {v_name, v_out_size_, prefix + "." + v_name + "."}, + }; + } else { + split_infos_ = { + {q_name, 0, q_out_size_, 0}, + {k_name, q_out_size_, k_out_size_, num_k_head_}, + {v_name, q_out_size_ + k_out_size_, v_out_size_, num_v_head_}, + }; + } + register_fused_params(); +} + +void QKVParallelLinear::register_fused_params() { + if (!register_fn_) { + return; + } + auto params = this->sharded_ + ? this->init_fused_shards(shard_specs_) + : this->split_params(split_infos_, tp_rank_, tp_size_, num_k_head_); for (auto &sp : params) { register_fn_(sp.full_name, std::move(sp.param)); } @@ -112,11 +142,11 @@ QKVParallelLinear::QKVParallelLinear(size_t hidden_size, void QKVParallelLinear::process_weights_after_loading() { BaseLinear::process_weights_after_loading(); + // `split_infos_` is empty for sharded quantization layouts because the + // shard parameters are the load targets. Reallocation would discard the + // bytes that were already loaded. if (register_fn_ && !split_infos_.empty()) { - auto params = this->split_params(split_infos_, tp_rank_, tp_size_, num_k_head_); - for (auto &sp : params) { - register_fn_(sp.full_name, std::move(sp.param)); - } + register_fused_params(); } } @@ -125,23 +155,27 @@ void QKVParallelLinear::process_weights_after_loading() { // --------------------------------------------------------- GateUpParallelLinear::GateUpParallelLinear(size_t hidden_size, size_t intermediate_size, std::shared_ptr quantization, bool bias, const infinicore::DataType &dtype, const infinicore::Device &device, - engine::distributed::RankInfo rank_info) - : GateUpParallelLinear(hidden_size, intermediate_size, bias, bias, quantization, dtype, device, rank_info) { + engine::distributed::RankInfo rank_info, + const std::string &stem) + : GateUpParallelLinear(hidden_size, intermediate_size, bias, bias, quantization, dtype, device, rank_info, stem) { } GateUpParallelLinear::GateUpParallelLinear(size_t hidden_size, size_t intermediate_size, bool gate_bias, bool up_bias, std::shared_ptr quantization, const infinicore::DataType &dtype, const infinicore::Device &device, - engine::distributed::RankInfo rank_info) + engine::distributed::RankInfo rank_info, + const std::string &stem) : infinilm::nn::ColumnParallelLinear( - hidden_size, - intermediate_size * 2, - quantization == nullptr ? std::make_shared() : quantization, - gate_bias || up_bias, - dtype, - device, - rank_info.tp_rank, - rank_info.tp_size), + hidden_size, + intermediate_size * 2, + quantization == nullptr ? std::make_shared() : quantization, + gate_bias || up_bias, + dtype, + device, + rank_info.tp_rank, + rank_info.tp_size, + -1, + stem), gate_bias_(gate_bias), up_bias_(up_bias) { if (gate_bias_ != up_bias_) { @@ -166,19 +200,45 @@ GateUpParallelLinear::GateUpParallelLinear(size_t hidden_size, size_t intermedia std::shared_ptr quantization, bool bias, const infinicore::DataType &dtype, const infinicore::Device &device, - engine::distributed::RankInfo rank_info) - : GateUpParallelLinear(hidden_size, intermediate_size, quantization, bias, dtype, device, rank_info) { - const std::string &key_name = parameters_.count("qweight") ? "qweight" : "weight"; - const auto &key_param = get_parameter_ref(key_name); - int fused_dim = this->get_quantization()->get_fused_split_dim(); - size_t logical_output = this->get_quantization()->get_logical_dim_size(key_param->size(fused_dim)); - size_t half_size = logical_output / 2; + engine::distributed::RankInfo rank_info, + const std::string &prefix) + : GateUpParallelLinear(hidden_size, intermediate_size, quantization, bias, dtype, device, rank_info, prefix) { register_fn_ = register_fn; - split_infos_ = { - {gate_name, 0, half_size}, - {up_name, half_size, half_size}, - }; - auto params = this->split_params(split_infos_, tp_rank_, tp_size_, -1); + if (this->sharded_) { + // GGUF gate and up shards may use different types and row sizes, so + // they cannot share a fused buffer. Each shard owns its buffer and + // independently resolves whether it is quantized or dense. + if (prefix.empty()) { + throw std::runtime_error( + "GateUpParallelLinear requires `prefix` when the " + "quantization scheme resolves layouts by checkpoint tensor name."); + } + const size_t half = intermediate_size / tp_size_; + shard_specs_ = { + {gate_name, half, prefix + "." + gate_name + "."}, + {up_name, half, prefix + "." + up_name + "."}, + }; + } else { + const std::string &key_name = parameters_.count("qweight") ? "qweight" : "weight"; + const auto &key_param = get_parameter_ref(key_name); + int fused_dim = this->get_quantization()->get_fused_split_dim(); + size_t logical_output = this->get_quantization()->get_logical_dim_size(key_param->size(fused_dim)); + size_t half_size = logical_output / 2; + split_infos_ = { + {gate_name, 0, half_size}, + {up_name, half_size, half_size}, + }; + } + register_fused_params(); +} + +void GateUpParallelLinear::register_fused_params() { + if (!register_fn_) { + return; + } + auto params = this->sharded_ + ? this->init_fused_shards(shard_specs_) + : this->split_params(split_infos_, tp_rank_, tp_size_, -1); for (auto &sp : params) { register_fn_(sp.full_name, std::move(sp.param)); } @@ -186,11 +246,10 @@ GateUpParallelLinear::GateUpParallelLinear(size_t hidden_size, size_t intermedia void GateUpParallelLinear::process_weights_after_loading() { BaseLinear::process_weights_after_loading(); + // As in `QKVParallelLinear`, sharded layouts have no `split_infos_` and + // must not repeat the split allocation. if (register_fn_ && !split_infos_.empty()) { - auto params = this->split_params(split_infos_, tp_rank_, tp_size_, -1); - for (auto &sp : params) { - register_fn_(sp.full_name, std::move(sp.param)); - } + register_fused_params(); } } diff --git a/csrc/layers/linear/fused_linear.hpp b/csrc/layers/linear/fused_linear.hpp index 8773a081c..edb81c899 100644 --- a/csrc/layers/linear/fused_linear.hpp +++ b/csrc/layers/linear/fused_linear.hpp @@ -16,7 +16,8 @@ class QKVParallelLinear : public infinilm::nn::ColumnParallelLinear { std::shared_ptr quantization = nullptr, const infinicore::DataType &dtype = infinicore::DataType::F32, const infinicore::Device &device = infinicore::Device(), - engine::distributed::RankInfo rank_info = engine::distributed::RankInfo()); + engine::distributed::RankInfo rank_info = engine::distributed::RankInfo(), + const std::string &stem = ""); explicit QKVParallelLinear(size_t hidden_size, size_t head_dim, @@ -36,7 +37,8 @@ class QKVParallelLinear : public infinilm::nn::ColumnParallelLinear { std::shared_ptr quantization = nullptr, const infinicore::DataType &dtype = infinicore::DataType::F32, const infinicore::Device &device = infinicore::Device(), - engine::distributed::RankInfo rank_info = engine::distributed::RankInfo()); + engine::distributed::RankInfo rank_info = engine::distributed::RankInfo(), + const std::string &prefix = ""); QKVParallelLinear(size_t hidden_size, size_t head_dim, @@ -47,7 +49,8 @@ class QKVParallelLinear : public infinilm::nn::ColumnParallelLinear { bool bias = false, const infinicore::DataType &dtype = infinicore::DataType::F32, const infinicore::Device &device = infinicore::Device(), - engine::distributed::RankInfo rank_info = engine::distributed::RankInfo()); + engine::distributed::RankInfo rank_info = engine::distributed::RankInfo(), + const std::string &prefix = ""); void process_weights_after_loading() override; @@ -91,6 +94,13 @@ class QKVParallelLinear : public infinilm::nn::ColumnParallelLinear { size_t num_kv_head_replicas_ = 1; RegisterParamFn register_fn_; std::vector split_infos_; + // Used by layouts such as GGUF that allocate one buffer per shard. This + // is mutually exclusive with `split_infos_`; see `sharded_`. + std::vector shard_specs_; + + // Pass each shard parameter to `register_fn`, whether it is a narrowed + // view or an independent buffer. + void register_fused_params(); }; class GateUpParallelLinear : public infinilm::nn::ColumnParallelLinear { @@ -100,12 +110,14 @@ class GateUpParallelLinear : public infinilm::nn::ColumnParallelLinear { bool bias = false, const infinicore::DataType &dtype = infinicore::DataType::F32, const infinicore::Device &device = infinicore::Device(), - engine::distributed::RankInfo rank_info = engine::distributed::RankInfo()); + engine::distributed::RankInfo rank_info = engine::distributed::RankInfo(), + const std::string &stem = ""); GateUpParallelLinear(size_t hidden_size, size_t intermediate_size, bool gate_bias, bool up_bias, std::shared_ptr quantization = nullptr, const infinicore::DataType &dtype = infinicore::DataType::F32, const infinicore::Device &device = infinicore::Device(), - engine::distributed::RankInfo rank_info = engine::distributed::RankInfo()); + engine::distributed::RankInfo rank_info = engine::distributed::RankInfo(), + const std::string &stem = ""); GateUpParallelLinear(size_t hidden_size, size_t intermediate_size, const std::string &gate_name, const std::string &up_name, @@ -114,7 +126,8 @@ class GateUpParallelLinear : public infinilm::nn::ColumnParallelLinear { bool bias = false, const infinicore::DataType &dtype = infinicore::DataType::F32, const infinicore::Device &device = infinicore::Device(), - engine::distributed::RankInfo rank_info = engine::distributed::RankInfo()); + engine::distributed::RankInfo rank_info = engine::distributed::RankInfo(), + const std::string &prefix = ""); void process_weights_after_loading() override; @@ -128,6 +141,9 @@ class GateUpParallelLinear : public infinilm::nn::ColumnParallelLinear { bool up_bias_; RegisterParamFn register_fn_; std::vector split_infos_; + std::vector shard_specs_; + + void register_fused_params(); }; } // namespace infinilm::layers::linear diff --git a/csrc/layers/linear/linear.cpp b/csrc/layers/linear/linear.cpp index f24496700..eacf0a187 100644 --- a/csrc/layers/linear/linear.cpp +++ b/csrc/layers/linear/linear.cpp @@ -13,8 +13,9 @@ Linear::Linear(size_t in_features, size_t out_features, bool bias, Linear::Linear(size_t in_features, size_t out_features, std::shared_ptr quantization, - bool bias, const infinicore::DataType &dtype, const infinicore::Device &device) - : BaseLinear(in_features, out_features, quantization, bias, dtype, device, -1, 0, 1) { + bool bias, const infinicore::DataType &dtype, const infinicore::Device &device, + const std::string &stem) + : BaseLinear(in_features, out_features, quantization, bias, dtype, device, -1, 0, 1, -1, stem) { } infinicore::Tensor Linear::forward(infinicore::Tensor &input) const { @@ -42,9 +43,9 @@ ColumnParallelLinear::ColumnParallelLinear(size_t in_features, size_t out_featur std::shared_ptr quantization, bool bias, const infinicore::DataType &dtype, const infinicore::Device &device, infinicore::Size tp_rank, infinicore::Size tp_size, - int tp_num_heads) + int tp_num_heads, const std::string &stem) : BaseLinear(in_features, out_features, quantization, bias, dtype, device, - 0, tp_rank, tp_size, tp_num_heads), + 0, tp_rank, tp_size, tp_num_heads, stem), tp_rank_(tp_rank), tp_size_(tp_size) { } @@ -74,9 +75,9 @@ RowParallelLinear::RowParallelLinear(size_t in_features, size_t out_features, std::shared_ptr quantization, bool bias, const infinicore::DataType &dtype, const infinicore::Device &device, infinicore::Size tp_rank, infinicore::Size tp_size, - infinicclComm_t communicator) + infinicclComm_t communicator, const std::string &stem) : BaseLinear(in_features, out_features, quantization, bias, dtype, device, - 1, tp_rank, tp_size), + 1, tp_rank, tp_size, -1, stem), tp_rank_(tp_rank), tp_size_(tp_size), communicator_(communicator) { } diff --git a/csrc/layers/linear/linear.hpp b/csrc/layers/linear/linear.hpp index 566cee77c..9f1046097 100644 --- a/csrc/layers/linear/linear.hpp +++ b/csrc/layers/linear/linear.hpp @@ -21,7 +21,8 @@ class Linear : public BaseLinear { std::shared_ptr quantization, bool bias = true, const infinicore::DataType &dtype = infinicore::DataType::F32, - const infinicore::Device &device = infinicore::Device()); + const infinicore::Device &device = infinicore::Device(), + const std::string &stem = ""); infinicore::Tensor forward(infinicore::Tensor &input) const; std::string extra_repr() const; @@ -42,7 +43,8 @@ class ColumnParallelLinear : public BaseLinear { const infinicore::DataType &dtype = infinicore::DataType::F32, const infinicore::Device &device = infinicore::Device(), infinicore::Size tp_rank = 0, infinicore::Size tp_size = 1, - int tp_num_heads = -1); + int tp_num_heads = -1, + const std::string &stem = ""); infinicore::Tensor forward(infinicore::Tensor &input) const; std::string extra_repr() const; @@ -67,7 +69,8 @@ class RowParallelLinear : public BaseLinear { const infinicore::DataType &dtype = infinicore::DataType::F32, const infinicore::Device &device = infinicore::Device(), infinicore::Size tp_rank = 0, infinicore::Size tp_size = 1, - infinicclComm_t communicator = nullptr); + infinicclComm_t communicator = nullptr, + const std::string &stem = ""); infinicore::Tensor forward(infinicore::Tensor &input) const; std::string extra_repr() const; diff --git a/csrc/layers/mlp/mlp.cpp b/csrc/layers/mlp/mlp.cpp index f7604c505..3a8d5f255 100644 --- a/csrc/layers/mlp/mlp.cpp +++ b/csrc/layers/mlp/mlp.cpp @@ -5,7 +5,8 @@ namespace infinilm::layers::mlp { MLP::MLP(std::shared_ptr model_config, - const infinicore::Device &device) { + const infinicore::Device &device, + const std::string &prefix) { const auto &dtype{model_config->get_dtype()}; hidden_size_ = model_config->get("hidden_size"); @@ -20,10 +21,11 @@ MLP::MLP(std::shared_ptr model_config, auto register_fn = [this](const std::string &n, infinicore::nn::Parameter p) { this->register_parameter(n, std::move(p)); }; gate_up_proj_ = std::make_shared( hidden_size_, intermediate_size_, "gate_proj", "up_proj", register_fn, - quantization_method, use_bias_, dtype, device, rank_info); + quantization_method, use_bias_, dtype, device, rank_info, prefix); down_proj_ = this->register_module( "down_proj", intermediate_size_, hidden_size_, quantization_method, - use_bias_, dtype, device, tp_rank, tp_size, rank_info.comm); + use_bias_, dtype, device, tp_rank, tp_size, rank_info.comm, + prefix.empty() ? std::string() : prefix + ".down_proj."); } infinicore::Tensor MLP::forward(const infinicore::Tensor &hidden_states) const { diff --git a/csrc/layers/mlp/mlp.hpp b/csrc/layers/mlp/mlp.hpp index abd81bd88..d929b1cd5 100644 --- a/csrc/layers/mlp/mlp.hpp +++ b/csrc/layers/mlp/mlp.hpp @@ -24,9 +24,13 @@ class MLP : public infinicore::nn::Module { * * @param model_config: Model configuration. * @param device Device to create tensors on + * @param prefix Checkpoint path for this layer, such as + * `layers.0.mlp`. It is used only by quantization schemes, such as + * GGUF, that resolve layouts by tensor name. */ MLP(std::shared_ptr model_config, - const infinicore::Device &device); + const infinicore::Device &device, + const std::string &prefix = ""); /** * @brief Forward pass: compute MLP output diff --git a/csrc/layers/quantization/base_quantization.hpp b/csrc/layers/quantization/base_quantization.hpp index 4b17cc949..9bc70eed5 100644 --- a/csrc/layers/quantization/base_quantization.hpp +++ b/csrc/layers/quantization/base_quantization.hpp @@ -38,7 +38,7 @@ struct SplitParam { class BaseQuantization : public std::enable_shared_from_this { public: - explicit BaseQuantization(const nlohmann::json &quant_config) : quant_config_(quant_config){}; + explicit BaseQuantization(const nlohmann::json &quant_config) : quant_config_(quant_config) {}; virtual ~BaseQuantization() = default; const nlohmann::json &get_config() const { return quant_config_; } @@ -62,6 +62,67 @@ class BaseQuantization : public std::enable_shared_from_this { float alpha = 1.0f) const = 0; + // ---- Name-aware variants ------------------------------------------------- + // Some schemes decide a parameter's layout from the *checkpoint tensor name* + // instead of from the module's (in_features, out_features) pair: GGUF block + // quantization has one ggml type per tensor, and `row_bytes` is a function of + // that type, so nothing can be derived from the logical shape alone. + // + // `stem` is the checkpoint path of the weight *without* the final tensor-name + // component, relative to quantization_config.key_prefix, and it always keeps + // the trailing separator: + // "layers.0.mlp.gate_proj." -> that one weight + // "layers.0.self_attn." -> ditto (probes weight / weight_bytes) + // "layers.0.self_attn" (no trailing '.') -> a fused linear: this scheme owns + // no buffer, each shard has its own checkpoint entry and is registered + // separately (see BaseLinear::init_fused_shards). + // An empty stem is always an error for such schemes. + // + // The default implementations forward to the name-less versions, so the + // existing quantization classes need no change. + virtual std::vector get_param_layout( + size_t in_features, size_t out_features, + int split_dim, int tp_rank, int tp_size, + int tp_num_heads, + const infinicore::DataType &dtype, + bool bias, + const std::string &stem) const { + return get_param_layout(in_features, out_features, split_dim, tp_rank, + tp_size, tp_num_heads, dtype, bias); + } + + virtual infinicore::Tensor forward( + const ParamsMap ¶ms, + const infinicore::Tensor &input, + bool has_bias, + float alpha, + const std::string &stem) const { + return forward(params, input, has_bias, alpha); + } + + // A fused linear's `stem` is only the *group* name (no trailing '.'), which is + // not enough for name-driven schemes: each shard has its own checkpoint entry + // and its own format (measured on Qwen3.8: 0 of 16 full-attn groups share one + // ggml type across q/k/v, 4 of 32 FFN groups across gate/up). + // + // `shard_stems[i]` is the checkpoint stem of the shard behind parameter key + // "shard." — both are produced by the same loop in + // BaseLinear::init_fused_shards, so index correspondence is a local invariant, + // not a cross-phase assumption. Empty vector = not a fused linear. + // + // Default forwards to the stem-only version, so quantization classes without + // per-shard formats need no change. + virtual infinicore::Tensor forward( + const ParamsMap ¶ms, + const infinicore::Tensor &input, + bool has_bias, + float alpha, + const std::string &stem, + const std::vector &shard_stems) const { + (void)shard_stems; + return forward(params, input, has_bias, alpha, stem); + } + virtual infinicore::Tensor forward_allreduce( const ParamsMap ¶ms, const infinicore::Tensor &input, diff --git a/csrc/layers/quantization/fp8.cpp b/csrc/layers/quantization/fp8.cpp new file mode 100644 index 000000000..475f5652e --- /dev/null +++ b/csrc/layers/quantization/fp8.cpp @@ -0,0 +1,177 @@ +#include "fp8.hpp" +#include "none_quantization.hpp" + +#include +#include +#include + +#include +#include + +namespace infinilm::quantization { + +std::vector FP8Quantization::get_param_layout( + size_t in_features, size_t out_features, + int split_dim, int tp_rank, int tp_size, + int /*tp_num_heads*/, + const infinicore::DataType &dtype, + bool bias) const { + + std::vector descs; + + // Weight: FP8 (E4M3) format - keep as F8, do NOT convert to BF16 + descs.push_back({"weight", {out_features, in_features}, infinicore::DataType::F8, split_dim, tp_rank, tp_size}); + + // Per-block weight scale (inverse): BF16, shape = [ceil(N/128), ceil(K/128)] + size_t num_out_blocks = (out_features + BLOCK_SIZE - 1) / BLOCK_SIZE; + size_t num_in_blocks = (in_features + BLOCK_SIZE - 1) / BLOCK_SIZE; + descs.push_back({"weight_scale_inv", {num_out_blocks, num_in_blocks}, infinicore::DataType::F32, split_dim, tp_rank, tp_size}); + + if (bias) { + descs.push_back({"bias", {out_features}, dtype, split_dim >= 0 ? 0 : -1, split_dim >= 0 ? tp_rank : 0, split_dim >= 0 ? tp_size : 1}); + } + return descs; +} + +infinicore::Tensor FP8Quantization::forward( + const ParamsMap ¶ms, + const infinicore::Tensor &input, + bool has_bias, + float /*alpha*/) const { + + auto weight_it = params.find("weight"); + auto scale_it = params.find("weight_scale_inv"); + auto bias_it = params.find("bias"); + + if (weight_it == params.end()) { + throw std::runtime_error("FP8Quantization::forward: weight not found"); + } + if (scale_it == params.end()) { + throw std::runtime_error("FP8Quantization::forward: weight_scale_inv not found"); + } + + auto weight = weight_it->second; + auto scale = scale_it->second; + + // Ensure input, weight, and scale are contiguous + // (split_params creates narrow views that may not be contiguous) + auto x = input->is_contiguous() ? input : input->contiguous(); + auto w = weight->is_contiguous() ? weight : weight->contiguous(); + auto s = scale->is_contiguous() ? scale : scale->contiguous(); + + // Get dimensions + auto x_shape = x->shape(); + size_t ndim = x_shape.size(); + size_t K = x_shape[ndim - 1]; // last dim is always feature dim + // M = product of all leading dims + size_t M = 1; + for (size_t i = 0; i < ndim - 1; i++) { + M *= x_shape[i]; + } + auto w_shape = w->shape(); + size_t N = w_shape[0]; + + // Flatten input to 2D [M, K] and ensure contiguous + auto flat = x->view({M, K}); + flat = flat->is_contiguous() ? flat : flat->contiguous(); + + // Allocate output [M, N] + auto out = infinicore::Tensor::empty( + {M, N}, input->dtype(), input->device()); + + // Call block-FP8 linear: BF16 input x F8 weight + block scale -> BF16 output + infinicore::op::block_fp8_linear_( + out, flat, w, s); + + if (has_bias && bias_it != params.end()) { + auto bias = bias_it->second; + auto bias_broadcast = bias->view({1, N}); + infinicore::op::add_(out, out, bias_broadcast); + } + + // Reshape output to match input's leading dims with N + std::vector out_shape(x_shape.begin(), x_shape.end() - 1); + out_shape.push_back(N); + return out->view(out_shape); +} + +std::vector FP8Quantization::split_params( + const std::unordered_map ¶ms, + const std::vector &splits, + int narrow_dim, + int tp_rank, int tp_size, int /*tp_num_heads*/) const { + + std::vector result; + auto weight_it = params.find("weight"); + auto scale_it = params.find("weight_scale_inv"); + auto bias_it = params.find("bias"); + + for (const auto &s : splits) { + result.push_back({s.prefix + ".weight", + infinicore::nn::Parameter( + weight_it->second->narrow({{static_cast(narrow_dim), s.start, s.size}}), + narrow_dim, tp_rank, tp_size, s.num_shards)}); + + if (scale_it != params.end()) { + size_t scale_start = s.start / BLOCK_SIZE; + size_t scale_size = (s.size + BLOCK_SIZE - 1) / BLOCK_SIZE; + result.push_back({s.prefix + ".weight_scale_inv", + infinicore::nn::Parameter( + scale_it->second->narrow({{static_cast(narrow_dim), scale_start, scale_size}}), + narrow_dim, tp_rank, tp_size, s.num_shards)}); + } + + if (bias_it != params.end()) { + result.push_back({s.prefix + ".bias", + infinicore::nn::Parameter( + bias_it->second->narrow({{0, s.start, s.size}}), + 0, tp_rank, tp_size, s.num_shards)}); + } + } + return result; +} + +std::shared_ptr FP8Quantization::process_weights_after_loading( + ParamsMap ¶ms, + const infinicore::Device &device, + int /*split_dim*/) const { + + auto weight_it = params.find("weight"); + auto scale_it = params.find("weight_scale_inv"); + + if (weight_it == params.end()) { + return nullptr; + } + if (scale_it == params.end()) { + spdlog::debug("FP8: no weight_scale_inv found, skipping"); + return nullptr; + } + + auto weight = weight_it->second; + auto scale = scale_it->second; + + size_t out_features = weight->shape()[0]; + size_t in_features = weight->shape()[1]; + + size_t num_out_blocks = (out_features + BLOCK_SIZE - 1) / BLOCK_SIZE; + size_t num_in_blocks = (in_features + BLOCK_SIZE - 1) / BLOCK_SIZE; + + auto scale_shape = scale->shape(); + if (scale_shape.size() != 2 || scale_shape[0] != num_out_blocks || scale_shape[1] != num_in_blocks) { + throw std::runtime_error("FP8Quantization: weight_scale_inv shape mismatch"); + } + + // Keep weight as FP8 (1 byte/element), just ensure contiguous + params["weight"] = weight->contiguous(); + + // Scale is already FP32 (converted on Python side during loading) + params["weight_scale_inv"] = scale->contiguous(); + + spdlog::debug("FP8: kept weight as F8, scale cast to F32, shape=[{}, {}]", + out_features, in_features); + + // Return nullptr to continue using FP8Quantization (not NoneQuantization) + return nullptr; +} + +} // namespace infinilm::quantization diff --git a/csrc/layers/quantization/fp8.hpp b/csrc/layers/quantization/fp8.hpp new file mode 100644 index 000000000..7897852ba --- /dev/null +++ b/csrc/layers/quantization/fp8.hpp @@ -0,0 +1,44 @@ +#pragma once + +#include "base_quantization.hpp" + +namespace infinilm::quantization { + +class FP8Quantization : public BaseQuantization { +public: + explicit FP8Quantization(const nlohmann::json &quant_config) + : BaseQuantization(quant_config) {} + + QuantScheme get_quant_scheme() const override { + return QuantScheme::FP8_W8A8; + } + + std::vector get_param_layout( + size_t in_features, size_t out_features, + int split_dim, int tp_rank, int tp_size, + int tp_num_heads, + const infinicore::DataType &dtype, + bool bias) const override; + + infinicore::Tensor forward( + const ParamsMap ¶ms, + const infinicore::Tensor &input, + bool has_bias, + float alpha = 1.0f) const override; + + std::vector split_params( + const std::unordered_map ¶ms, + const std::vector &splits, + int narrow_dim, + int tp_rank, int tp_size, int tp_num_heads) const override; + + std::shared_ptr process_weights_after_loading( + ParamsMap ¶ms, + const infinicore::Device &device, + int split_dim = -1) const override; + +private: + static constexpr size_t BLOCK_SIZE = 128; +}; + +} // namespace infinilm::quantization diff --git a/csrc/layers/quantization/gguf.cpp b/csrc/layers/quantization/gguf.cpp new file mode 100644 index 000000000..9c52e2990 --- /dev/null +++ b/csrc/layers/quantization/gguf.cpp @@ -0,0 +1,590 @@ +#include "gguf.hpp" + +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace infinilm::quantization { + +namespace { + +// GGML block metadata. A packed row contains +// in_features / block_size * type_size bytes. Unsupported types must be +// converted to dense BF16 instead of relying on a guessed runtime stride. +struct GgmlBlock { + int64_t id; + const char *name; + size_t block_size; + size_t type_size; +}; + +constexpr GgmlBlock GGML_BLOCKS[] = { + {8, "Q8_0", 32, 34}, + {12, "Q4_K", 256, 144}, + {13, "Q5_K", 256, 176}, + {14, "Q6_K", 256, 210}, +}; + +const GgmlBlock *ggml_block(int64_t id) { + for (const auto &b : GGML_BLOCKS) { + if (b.id == id) { + return &b; + } + } + return nullptr; +} + +std::string supported_types() { + std::string s; + for (const auto &b : GGML_BLOCKS) { + if (!s.empty()) { + s += "/"; + } + s += b.name; + } + return s; +} + +constexpr const char *DENSE_MARK = "dense_bf16"; + +bool env_enabled(const char *name) { + const char *value = std::getenv(name); + return value != nullptr && value[0] != '\0' && std::strcmp(value, "0") != 0; +} + +bool use_f32_decode_output(const std::string &table_key, size_t m_count) { + if (!env_enabled("INFINI_GGUF_F32_DECODE_OUT") || m_count > 16) { + return false; + } + const char *match = std::getenv("INFINI_GGUF_F32_DECODE_OUT_MATCH"); + return match == nullptr || match[0] == '\0' + || table_key.find(match) != std::string::npos; +} + +} // namespace + +GGUFBlockQuantization::GGUFBlockQuantization(const nlohmann::json &quant_config) + : BaseQuantization(quant_config) { + if (!quant_config_.is_object() || !quant_config_.contains("ggml_types")) { + throw std::runtime_error( + "GGUFBlockQuantization: quantization_config is missing ggml_types"); + } + key_prefix_ = get_or("key_prefix", ""); + + const auto &table = quant_config_.at("ggml_types"); + if (!table.is_object() || table.empty()) { + throw std::runtime_error("GGUFBlockQuantization: ggml_types is empty"); + } + + size_t n_blob = 0; + size_t n_dense = 0; + size_t n_outside = 0; + for (const auto &kv : table.items()) { + const std::string &name = kv.key(); + // Keep keys outside key_prefix unchanged. This includes root-level + // tensors such as lm_head.weight and preserves exact checkpoint names. + std::string key = name; + if (!key_prefix_.empty() && name.compare(0, key_prefix_.size(), key_prefix_) == 0) { + key = name.substr(key_prefix_.size()); + } else { + ++n_outside; + } + + int64_t id = DENSE_BF16; + if (kv.value().is_string()) { + const std::string v = kv.value().get(); + if (v != DENSE_MARK) { + throw std::runtime_error( + "GGUFBlockQuantization: value '" + v + "' for '" + name + + "' is neither an integer type id nor \"" + DENSE_MARK + "\""); + } + ++n_dense; + } else { + if (!kv.value().is_number_integer()) { + throw std::runtime_error( + "GGUFBlockQuantization: value for '" + name + "' is not an integer ggml type id"); + } + id = kv.value().get(); + if (id == DENSE_BF16) { + throw std::runtime_error( + "GGUFBlockQuantization: type id for '" + name + "' conflicts with dense sentinel -1"); + } + if (!ggml_block(id)) { + throw std::runtime_error( + "GGUFBlockQuantization: '" + name + "' uses unsupported ggml type id=" + + std::to_string(id) + " (supported: " + supported_types() + + "); unsupported types must be converted to dense BF16"); + } + ++n_blob; + } + + if (!types_.emplace(std::move(key), TypeEntry{id, name}).second) { + throw std::runtime_error( + "GGUFBlockQuantization: duplicate key after removing key_prefix: '" + name + "'"); + } + } + + // Activation value-head permutation is required even when the rule list is + // empty. Missing metadata could silently misalign activations and columns. + if (!quant_config_.contains("activation_vperm")) { + throw std::runtime_error( + "GGUFBlockQuantization: quantization_config is missing activation_vperm; " + "refresh config.json with the converter --skip-pack option"); + } + { + const auto &rules = quant_config_.at("activation_vperm"); + if (!rules.is_array()) { + throw std::runtime_error("GGUFBlockQuantization: activation_vperm must be an array, got " + std::string(rules.type_name())); + } + for (const auto &j : rules) { + if (!j.is_object()) { + throw std::runtime_error("GGUFBlockQuantization: activation_vperm entry must be an object"); + } + ActVPerm r; + for (const char *key : {"suffix", "num_k_heads", "num_v_per_k", "head_dim"}) { + if (!j.contains(key)) { + throw std::runtime_error("GGUFBlockQuantization: activation_vperm entry is missing '" + std::string(key) + "'"); + } + } + r.suffix = j.at("suffix").get(); + r.n_k = j.at("num_k_heads").get(); + r.r = j.at("num_v_per_k").get(); + r.hd = j.at("head_dim").get(); + if (r.suffix.empty() || r.suffix.back() != '.' || !r.n_k || !r.r || !r.hd) { + throw std::runtime_error( + "GGUFBlockQuantization: invalid activation_vperm entry: suffix='" + + r.suffix + "' must end with '.', and dimensions must be positive (got " + + std::to_string(r.n_k) + "/" + std::to_string(r.r) + "/" + + std::to_string(r.hd) + ")"); + } + if (std::any_of(vperm_.begin(), vperm_.end(), + [&r](const ActVPerm &e) { return e.suffix == r.suffix; })) { + throw std::runtime_error("GGUFBlockQuantization: duplicate activation_vperm suffix '" + r.suffix + "'"); + } + vperm_.push_back(std::move(r)); + } + } + + static std::atomic config_logged{false}; + if (!config_logged.exchange(true, std::memory_order_relaxed)) { + spdlog::info( + "GGUF block quantization: {} entries (blob {} / dense {} / outside prefix {}), key_prefix='{}'{}", + types_.size(), n_blob, n_dense, n_outside, key_prefix_, + key_prefix_.empty() + ? " (not set; table keys are relative safetensors names)" + : " (for example, root-level lm_head entries)"); + + std::string vs; + for (const auto &r : vperm_) { + if (!vs.empty()) { + vs += ", "; + } + vs += r.suffix + "=" + std::to_string(r.n_k) + "x" + std::to_string(r.r) + "x" + std::to_string(r.hd); + } + spdlog::info("GGUF block quantization: {} activation V-head permutation rules (grouped->tiled): {}", + vperm_.size(), vs.empty() ? "none" : vs); + } +} + +GGUFBlockQuantization::~GGUFBlockQuantization() { + if (n_blob_ + n_dense_ + n_group_ > 0) { + spdlog::debug("GGUF block quantization: layout matches blob {} / dense {} / fused group {}", + n_blob_, n_dense_, n_group_); + } +} + +bool GGUFBlockQuantization::is_known_type(int64_t type_id) { + return ggml_block(type_id) != nullptr; +} + +std::string GGUFBlockQuantization::describe(const std::string &stem) const { + // Restore the absolute checkpoint name for searchable diagnostics. + return (stem.empty() ? std::string("") : key_prefix_ + stem); +} + +int64_t GGUFBlockQuantization::resolve(const std::string &stem, std::string *matched_key) const { + const std::string blob_key = stem + BLOB_SUFFIX; + const std::string dense_key = stem + DENSE_SUFFIX; + const auto blob_it = types_.find(blob_key); + const auto dense_it = types_.find(dense_key); + const int hits = (blob_it != types_.end()) + (dense_it != types_.end()); + + // Require exactly one packed or dense candidate. Falling back on missing + // metadata could load successfully while producing incorrect output. + if (hits != 1) { + throw std::runtime_error( + "GGUFBlockQuantization: stem '" + describe(stem) + "' matched " + + std::to_string(hits) + " type-table candidates; expected exactly one of '" + + blob_key + "' or '" + dense_key + "' among " + + std::to_string(types_.size()) + " entries"); + } + const auto &hit = blob_it != types_.end() ? *blob_it : *dense_it; + if (matched_key) { + *matched_key = hit.second.name; + } + return hit.second.id; +} + +bool GGUFBlockQuantization::has_group(const std::string &group_stem) const { + const std::string head = group_stem + "."; + return std::any_of(types_.begin(), types_.end(), [&head](const auto &kv) { + return kv.first.compare(0, head.size(), head) == 0; + }); +} + +size_t GGUFBlockQuantization::row_bytes(size_t in_features, int64_t type_id) const { + const GgmlBlock *b = ggml_block(type_id); + if (!b) { + throw std::runtime_error( + "GGUFBlockQuantization: unsupported ggml type id=" + std::to_string(type_id) + + " (supported: " + supported_types() + ")"); + } + if (in_features % b->block_size != 0) { + throw std::runtime_error( + "GGUFBlockQuantization: in_features=" + std::to_string(in_features) + + " is not divisible by " + b->name + " block size " + + std::to_string(b->block_size)); + } + return in_features / b->block_size * b->type_size; +} + +const GGUFBlockQuantization::ActVPerm *GGUFBlockQuantization::vperm_rule( + const std::string &stem) const { + for (const auto &r : vperm_) { + if (stem.size() >= r.suffix.size() && stem.compare(stem.size() - r.suffix.size(), r.suffix.size(), r.suffix) == 0) { + return &r; + } + } + return nullptr; +} + +infinicore::Tensor GGUFBlockQuantization::gather_grouped_to_tiled( + const ActVPerm &rule, const infinicore::Tensor &input, const std::string &name) { + const auto shape = input->shape(); + const size_t ndim = shape.size(); + if (ndim < 2) { + throw std::runtime_error( + "GGUFBlockQuantization: activation for " + name + " has rank=" + + std::to_string(ndim) + "; expected at least [..., in_features]"); + } + const size_t K = shape[ndim - 1]; + const size_t want = rule.n_k * rule.r * rule.hd; + if (K != want) { + throw std::runtime_error( + "GGUFBlockQuantization: activation last dimension for " + name + " is " + + std::to_string(K) + ", expected num_k_heads*num_v_per_k*head_dim=" + + std::to_string(want) + "; head permutation cannot be applied to a shard"); + } + // [..., n_k, r, hd] -> [..., r, n_k, hd], grouped to tiled order. + const size_t k_axis = ndim - 1; + infinicore::Shape grouped(shape.begin(), shape.end() - 1); + grouped.insert(grouped.end(), {rule.n_k, rule.r, rule.hd}); + infinicore::Shape order; + order.reserve(grouped.size()); + for (size_t a = 0; a + 1 < ndim; ++a) { + order.push_back(a); + } + order.insert(order.end(), {k_axis + 1, k_axis, k_axis + 2}); + + auto x = input->is_contiguous() ? input : input->contiguous(); + return x->view(grouped)->permute(order)->contiguous()->view(shape); +} + +std::vector GGUFBlockQuantization::get_param_layout( + size_t, size_t, int, int, int, int, + const infinicore::DataType &, bool) const { + throw std::runtime_error( + "GGUFBlockQuantization: get_param_layout requires a checkpoint stem to resolve the ggml type"); +} + +std::vector GGUFBlockQuantization::get_param_layout( + size_t in_features, size_t out_features, + int split_dim, int tp_rank, int tp_size, + int tp_num_heads, + const infinicore::DataType &dtype, + bool bias, + const std::string &stem) const { + (void)tp_num_heads; + + if (stem.empty()) { + throw std::runtime_error( + "GGUFBlockQuantization: missing checkpoint stem while constructing Linear (in=" + + std::to_string(in_features) + ", out=" + std::to_string(out_features) + ")"); + } + if (tp_size != 1 || tp_rank != 0) { + throw std::runtime_error( + "GGUFBlockQuantization: tensor parallelism is not supported for packed GGUF weights: " + + describe(stem)); + } + if (bias) { + throw std::runtime_error( + "GGUFBlockQuantization: GGUF checkpoint has no bias tensor for " + describe(stem)); + } + + // A stem without a trailing '.' identifies a fused Linear group. Its + // individual shard buffers are allocated by BaseLinear::init_fused_shards. + if (stem.back() != '.') { + if (!has_group(stem)) { + throw std::runtime_error( + "GGUFBlockQuantization: fused-group stem '" + stem + + "' has no '" + stem + "..*' entry in the type table"); + } + ++n_group_; + return {}; + } + + const int64_t id = resolve(stem); + if (id == DENSE_BF16) { + ++n_dense_; + // The converter stored this tensor as dense BF16; use regular GEMM. + return {{"weight", {out_features, in_features}, dtype, split_dim, tp_rank, tp_size}}; + } + + ++n_blob_; + const size_t rb = row_bytes(in_features, id); + return {{{BLOB_SUFFIX}, {out_features, rb}, infinicore::DataType::U8, split_dim, tp_rank, tp_size}}; +} + +infinicore::Tensor GGUFBlockQuantization::forward( + const ParamsMap &, const infinicore::Tensor &, bool, float) const { + throw std::runtime_error( + "GGUFBlockQuantization: forward requires a checkpoint stem; fused Linear also requires shard_stems"); +} + +infinicore::Tensor GGUFBlockQuantization::forward_shard( + const std::string &suffix, + const infinicore::Tensor &weight, + const infinicore::Tensor &input, + float alpha, + int64_t type_id, + const std::string &table_key) const { + if (suffix == DENSE_SUFFIX) { + // The suffix selected by get_param_layout must agree with the type table. + if (type_id != DENSE_BF16) { + throw std::runtime_error( + "GGUFBlockQuantization: " + table_key + " has parameter suffix " + + DENSE_SUFFIX + " but type table reports ggml type id=" + + std::to_string(type_id)); + } + auto x = input->is_contiguous() ? input : input->contiguous(); + auto w = weight->is_contiguous() ? weight : weight->contiguous(); + return infinicore::op::linear(x, w, std::nullopt, alpha); + } + if (suffix == BLOB_SUFFIX) { + // Pass packed block bytes directly to the kernel. Never reinterpret + // them as BF16 through a dense fallback. + if (alpha != 1.0F) { + throw std::runtime_error( + "linear_gguf: alpha=" + std::to_string(alpha) + + " is unsupported for packed GGUF weights: " + table_key); + } + auto x = input->is_contiguous() ? input : input->contiguous(); + auto w = weight->is_contiguous() ? weight : weight->contiguous(); + + const auto x_shape = x->shape(); + const size_t ndim = x_shape.size(); + const size_t K = x_shape[ndim - 1]; + size_t M = 1; + for (size_t i = 0; i + 1 < ndim; ++i) { + M *= x_shape[i]; + } + const size_t N = static_cast(w->size(0)); + // linear_gguf selects the decode or prefill kernel using its shared + // kMaxDecodeM threshold, so this layer does not duplicate that limit. + + auto flat = x->view({M, K}); + flat = flat->is_contiguous() ? flat : flat->contiguous(); + const bool f32_decode_out = use_f32_decode_output(table_key, M); + const auto out_dtype = f32_decode_out + ? infinicore::DataType::F32 + : input->dtype(); + auto out = infinicore::Tensor::empty({M, N}, out_dtype, input->device()); + // Log the first packed invocation as a lightweight wiring diagnostic. + static std::atomic blob_calls{0}; + if (blob_calls.fetch_add(1) == 0) { + spdlog::info( + "linear_gguf: first packed forward {} -- M={} N={} K={} ggml_type={} row_bytes={}", + table_key, M, N, K, type_id, w->size(1)); + } + if (f32_decode_out) { + static std::atomic f32_calls{0}; + if (f32_calls.fetch_add(1) == 0) { + spdlog::warn( + "linear_gguf: experimental F32 decode output enabled; first match {} -- M={} N={} K={}", + table_key, M, N, K); + } + } + infinicore::op::linear_gguf_(out, flat, w, type_id); + + std::vector out_shape(x_shape.begin(), x_shape.end() - 1); + out_shape.push_back(N); + return out->view(out_shape); + } + throw std::runtime_error( + "GGUFBlockQuantization: parameter suffix '" + suffix + "' for " + table_key + + " is neither " + DENSE_SUFFIX + " nor " + BLOB_SUFFIX); +} + +infinicore::Tensor GGUFBlockQuantization::forward( + const ParamsMap ¶ms, + const infinicore::Tensor &input, + bool has_bias, + float alpha, + const std::string &stem) const { + // The overload below handles fused Linear layers with per-shard stems. + return forward(params, input, has_bias, alpha, stem, {}); +} + +infinicore::Tensor GGUFBlockQuantization::forward( + const ParamsMap ¶ms, + const infinicore::Tensor &input, + bool has_bias, + float alpha, + const std::string &stem, + const std::vector &shard_stems) const { + if (has_bias) { + throw std::runtime_error( + "GGUFBlockQuantization: bias is not supported (" + describe(stem) + ")"); + } + + // Apply activation permutation before either packed or dense execution; + // both layouts preserve the source GGUF column order. + infinicore::Tensor x = input; + const ActVPerm *rule = vperm_rule(stem); + if (!shard_stems.empty()) { + for (const auto &s : shard_stems) { + if (vperm_rule(s)) { + throw std::runtime_error( + "GGUFBlockQuantization: shard '" + describe(s) + "' in fused group '" + + describe(stem) + "' matches an activation-permutation rule, but one input " + "cannot be permuted independently for each shard"); + } + } + } else if (rule) { + x = gather_grouped_to_tiled(*rule, input, describe(stem)); + // Log the first permutation as a lightweight wiring diagnostic. + static std::atomic vperm_applied{0}; + if (vperm_applied.fetch_add(1) == 0) { + spdlog::info( + "linear_gguf: first activation V-head permutation {} -- grouped->tiled {}x{}x{}", + describe(stem), rule->n_k, rule->r, rule->hd); + } + } + + // A non-fused layer owns exactly one weight or weight_bytes parameter. + if (shard_stems.empty()) { + if (params.size() != 1) { + throw std::runtime_error( + "GGUFBlockQuantization: " + describe(stem) + " has " + + std::to_string(params.size()) + + " parameters but no shard_stems; BaseLinear::compute_linear did not pass them"); + } + const auto &kv = *params.begin(); + std::string table_key; + const int64_t id = resolve(stem, &table_key); + return forward_shard(kv.first, kv.second, x, alpha, id, table_key); + } + + // Fused parameters use shard., where i is their order along the + // output dimension and matches SplitInfo. Resolve each shard independently + // and concatenate outputs in that order. + if (shard_stems.size() != params.size()) { + throw std::runtime_error( + "GGUFBlockQuantization: " + describe(stem) + " has " + + std::to_string(params.size()) + " shard parameters but received " + + std::to_string(shard_stems.size()) + + " shard stems; both must be created by BaseLinear::init_fused_shards"); + } + std::vector> parts; + for (const auto &kv : params) { + if (kv.first.compare(0, std::string(SHARD_PREFIX).size(), SHARD_PREFIX) != 0) { + throw std::runtime_error( + "GGUFBlockQuantization: fused Linear parameter '" + kv.first + + "' does not match " + SHARD_PREFIX + ". (" + describe(stem) + ")"); + } + const size_t dot = kv.first.find('.'); + if (dot == std::string::npos) { + throw std::runtime_error( + "GGUFBlockQuantization: fused Linear parameter '" + kv.first + "' is missing '.'"); + } + const size_t idx = std::stoul(kv.first.substr(std::string(SHARD_PREFIX).size(), + dot - std::string(SHARD_PREFIX).size())); + if (idx >= shard_stems.size()) { + throw std::runtime_error( + "GGUFBlockQuantization: shard index in parameter '" + kv.first + + "' is outside shard_stems (" + describe(stem) + ")"); + } + std::string table_key; + const int64_t id = resolve(shard_stems[idx], &table_key); + parts.emplace_back(idx, forward_shard(kv.first.substr(dot + 1), kv.second, x, alpha, + id, table_key)); + } + std::sort(parts.begin(), parts.end(), + [](const auto &a, const auto &b) { return a.first < b.first; }); + + std::vector outs; + outs.reserve(parts.size()); + for (auto &p : parts) { + outs.push_back(p.second); + } + const auto shape = input->shape(); + return infinicore::op::cat(outs, static_cast(shape.size()) - 1); +} + +std::vector GGUFBlockQuantization::split_params( + const std::unordered_map ¶ms, + const std::vector &splits, + int, int, int, int) const { + // Fused GGUF shards already have independent buffers. Rename + // shard. to . without slicing data. + std::vector result; + for (size_t i = 0; i < splits.size(); ++i) { + const std::string head = std::string(SHARD_PREFIX) + std::to_string(i) + "."; + for (const auto &kv : params) { + if (kv.first.compare(0, head.size(), head) != 0) { + continue; + } + result.push_back({splits[i].prefix + "." + kv.first.substr(head.size()), + infinicore::nn::Parameter(kv.second)}); + } + } + if (result.size() != splits.size()) { + throw std::runtime_error( + "GGUFBlockQuantization::split_params: expected " + + std::to_string(splits.size()) + " shard parameters but matched " + + std::to_string(result.size()) + + "; fused GGUF Linear must use BaseLinear::init_fused_shards"); + } + return result; +} + +std::shared_ptr GGUFBlockQuantization::process_weights_after_loading( + ParamsMap ¶ms, + const infinicore::Device &, + int) const { + for (auto &kv : params) { + const bool is_blob = kv.first.size() >= strlen(BLOB_SUFFIX) && kv.first.compare(kv.first.size() - strlen(BLOB_SUFFIX), strlen(BLOB_SUFFIX), BLOB_SUFFIX) == 0; + if (!is_blob) { + continue; + } + if (kv.second->dtype() != infinicore::DataType::U8) { + throw std::runtime_error( + "GGUFBlockQuantization: packed parameter '" + kv.first + "' must have U8 dtype"); + } + if (!kv.second->is_contiguous()) { + throw std::runtime_error( + "GGUFBlockQuantization: packed parameter '" + kv.first + "' must be contiguous"); + } + } + // Keep the quantization scheme and raw bytes unchanged. + return nullptr; +} + +} // namespace infinilm::quantization diff --git a/csrc/layers/quantization/gguf.hpp b/csrc/layers/quantization/gguf.hpp new file mode 100644 index 000000000..bd55cf485 --- /dev/null +++ b/csrc/layers/quantization/gguf.hpp @@ -0,0 +1,147 @@ +#pragma once + +#include "base_quantization.hpp" + +#include +#include +#include + +namespace infinilm::quantization { + +// GGUF block quantization stores the original block bytes in safetensors. +// Each row has row_bytes(in_features, type) bytes rather than in_features +// scalar elements, so the checkpoint tensor name determines the layout. +// +// config.json:quantization_config.ggml_types maps safetensors names to either +// a ggml type id or "dense_bf16" for tensors dequantized during conversion. +// key_prefix is removed once because nested modules do not know their absolute +// checkpoint path. +class GGUFBlockQuantization : public BaseQuantization { +public: + // Sentinel for entries stored as dense BF16 rather than GGUF blocks. + static constexpr int64_t DENSE_BF16 = -1; + // Checkpoint suffix for raw block data; shared with scripts/gguf_mapping.py. + static constexpr const char *BLOB_SUFFIX = "weight_bytes"; + static constexpr const char *DENSE_SUFFIX = "weight"; + // Parameter-key prefix used for fused Linear shards. + static constexpr const char *SHARD_PREFIX = "shard"; + + explicit GGUFBlockQuantization(const nlohmann::json &quant_config); + + ~GGUFBlockQuantization() override; + + QuantScheme get_quant_scheme() const override { + return QuantScheme::GGUF_BLOCK; + } + + // A stem is required to resolve the ggml type. + std::vector get_param_layout( + size_t in_features, size_t out_features, + int split_dim, int tp_rank, int tp_size, + int tp_num_heads, + const infinicore::DataType &dtype, + bool bias) const override; + + std::vector get_param_layout( + size_t in_features, size_t out_features, + int split_dim, int tp_rank, int tp_size, + int tp_num_heads, + const infinicore::DataType &dtype, + bool bias, + const std::string &stem) const override; + + infinicore::Tensor forward( + const ParamsMap ¶ms, + const infinicore::Tensor &input, + bool has_bias, + float alpha = 1.0f) const override; + + infinicore::Tensor forward( + const ParamsMap ¶ms, + const infinicore::Tensor &input, + bool has_bias, + float alpha, + const std::string &stem) const override; + + // Fused Linear entry point. Every shard resolves its own ggml type from + // its checkpoint stem; a shared group stem is insufficient. + infinicore::Tensor forward( + const ParamsMap ¶ms, + const infinicore::Tensor &input, + bool has_bias, + float alpha, + const std::string &stem, + const std::vector &shard_stems) const override; + + // Map shard names to .. Fused GGUF shards already have + // independent buffers, so this method does not slice or modify data. + std::vector split_params( + const std::unordered_map ¶ms, + const std::vector &splits, + int narrow_dim, + int tp_rank, int tp_size, int tp_num_heads) const override; + + // Validate raw block buffers without modifying their bytes. + std::shared_ptr process_weights_after_loading( + ParamsMap ¶ms, + const infinicore::Device &device, + int split_dim = -1) const override; + + // Resolve a stem to a ggml type id or DENSE_BF16. Missing and ambiguous + // matches fail closed instead of silently selecting a dense path. + // matched_key receives the actual normalized table key when requested. + int64_t resolve(const std::string &stem, std::string *matched_key = nullptr) const; + size_t row_bytes(size_t in_features, int64_t type_id) const; + bool has_group(const std::string &group_stem) const; + size_t table_size() const { return types_.size(); } + + static bool is_known_type(int64_t type_id); + +private: + // type_id selects the execution path; table_key identifies the exact + // checkpoint entry in diagnostics. + infinicore::Tensor forward_shard( + const std::string &suffix, + const infinicore::Tensor &weight, + const infinicore::Tensor &input, + float alpha, + int64_t type_id, + const std::string &table_key) const; + + // Runtime value-head permutation for weights whose columns were exported + // in tiled order while the GDN kernel produces grouped activations. The + // converter records rules in quantization_config.activation_vperm because + // permuting quantized columns would require requantization. + struct ActVPerm { + std::string suffix; // Suffix including the trailing '.', for example "linear_attn.out_proj.". + size_t n_k; // Number of key heads. + size_t r; // Value heads per key head. + size_t hd; // value head_dim + }; + + // Match by suffix because layer indices are not part of this local contract. + const ActVPerm *vperm_rule(const std::string &stem) const; + + // Convert [..., n_k*r*hd] grouped order to [..., r*n_k*hd] tiled order. + static infinicore::Tensor gather_grouped_to_tiled( + const ActVPerm &rule, const infinicore::Tensor &input, const std::string &name); + + std::string describe(const std::string &stem) const; + + // Keep the original config key for diagnostics even though lookups use a + // normalized key with key_prefix removed. + struct TypeEntry { + int64_t id; + std::string name; + }; + + std::unordered_map types_; + std::string key_prefix_; + std::vector vperm_; // Empty when the converted model needs no permutation. + // Mutable because layout queries are logically const. + mutable size_t n_blob_ = 0; + mutable size_t n_dense_ = 0; + mutable size_t n_group_ = 0; +}; + +} // namespace infinilm::quantization diff --git a/csrc/layers/quantization/quantization.hpp b/csrc/layers/quantization/quantization.hpp index 0cc9cd7e2..b75314e1b 100644 --- a/csrc/layers/quantization/quantization.hpp +++ b/csrc/layers/quantization/quantization.hpp @@ -4,6 +4,8 @@ #include "awq_marlin.hpp" #include "base_quantization.hpp" #include "compressed_tensors.hpp" +#include "fp8.hpp" +#include "gguf.hpp" #include "gptq.hpp" #include "gptq_marlin.hpp" #include "gptq_qy.hpp" diff --git a/csrc/layers/quantization/quantization_scheme.hpp b/csrc/layers/quantization/quantization_scheme.hpp index 455968a7a..37f414c6f 100644 --- a/csrc/layers/quantization/quantization_scheme.hpp +++ b/csrc/layers/quantization/quantization_scheme.hpp @@ -11,6 +11,9 @@ enum class QuantScheme { GPTQ_W4A16, GPTQ_MARLIN_W4A16, MXFP4_W4A16, + FP8_W8A16, + FP8_W8A8, + GGUF_BLOCK, }; enum class KVQuantAlgo { diff --git a/csrc/models/qwen3_5/qwen3_5_attention.cpp b/csrc/models/qwen3_5/qwen3_5_attention.cpp index e7a47f4f9..6371f6f65 100644 --- a/csrc/models/qwen3_5/qwen3_5_attention.cpp +++ b/csrc/models/qwen3_5/qwen3_5_attention.cpp @@ -6,12 +6,33 @@ #include "../../utils.hpp" #include #include +#include #include #include #include +#include #include namespace infinilm::models::qwen3_5 { +namespace { + +bool should_dump_attention(size_t layer_idx) { + const char *dump_dir = std::getenv("INFINILM_ATTENTION_DUMP_DIR"); + const char *target = std::getenv("INFINILM_ATTENTION_DUMP_LAYER"); + return dump_dir != nullptr && dump_dir[0] != '\0' + && target != nullptr && target[0] != '\0' + && layer_idx == std::strtoull(target, nullptr, 10); +} + +void dump_attention_tensor(const infinicore::Tensor &tensor, + const char *name, + size_t layer_idx) { + const char *dump_dir = std::getenv("INFINILM_ATTENTION_DUMP_DIR"); + tensor->debug(std::string(dump_dir) + "/infini_attention_" + name + "_" + + std::to_string(layer_idx) + ".bin"); +} + +} // namespace Qwen35Attention::Qwen35Attention(std::shared_ptr model_config, size_t layer_idx, @@ -44,13 +65,18 @@ Qwen35Attention::Qwen35Attention(std::shared_ptr auto quantization_method = model_config->get_quantization_method(); auto register_fn = [this](const std::string &n, infinicore::nn::Parameter p) { this->register_parameter(n, std::move(p)); }; + // Checkpoint path for this layer after removing + // `quantization_config.key_prefix`. It is used only by quantization + // schemes, such as GGUF, that resolve types by tensor name. + const std::string prefix = "layers." + std::to_string(layer_idx_) + ".self_attn"; qkv_proj_ = std::make_shared( hidden_size_, head_dim_, total_num_heads, total_num_kv_heads, "q_proj", "k_proj", "v_proj", register_fn, - quantization_method, use_bias, dtype, device, rank_info); + quantization_method, use_bias, dtype, device, rank_info, prefix); o_proj_ = this->register_module( "o_proj", total_num_heads * head_dim_, hidden_size_, quantization_method, - use_output_bias, dtype, device, tp_rank, tp_size, rank_info.comm); + use_output_bias, dtype, device, tp_rank, tp_size, rank_info.comm, + prefix + ".o_proj."); const auto &rope_params = model_config->get_config_json()["rope_parameters"]; const double partial_rotary_factor = rope_params["partial_rotary_factor"].get(); @@ -135,21 +161,50 @@ infinicore::Tensor Qwen35Attention::forward_paged_(const infinicore::Tensor &pos ASSERT_EQ(batch_size, 1); auto [q, gate, k, v] = qkv_proj_->forward_split(hidden_states_mutable); + const bool dump_attention = should_dump_attention(layer_idx_); + if (dump_attention) { + dump_attention_tensor(q, "q_raw", layer_idx_); + dump_attention_tensor(gate, "gate_raw", layer_idx_); + dump_attention_tensor(k, "k_raw", layer_idx_); + dump_attention_tensor(v, "v_raw", layer_idx_); + } auto q_reshaped = q->view({seq_len, num_attention_heads_, head_dim_}); auto k_reshaped = k->view({seq_len, num_key_value_heads_, head_dim_}); auto v_reshaped = v->view({seq_len, num_key_value_heads_, head_dim_}); q_reshaped = q_norm_->forward(q_reshaped); k_reshaped = k_norm_->forward(k_reshaped); + if (dump_attention) { + dump_attention_tensor(q_reshaped, "q_norm", layer_idx_); + dump_attention_tensor(k_reshaped, "k_norm", layer_idx_); + } auto pos_shape = position_ids->shape(); if (pos_shape.size() != 2 && pos_shape.size() != 1) { throw std::runtime_error("Unexpected position_ids shape"); } std::tie(q_reshaped, k_reshaped) = mrope_->forward(q_reshaped, k_reshaped, position_ids); + if (dump_attention) { + dump_attention_tensor(q_reshaped, "q_rope", layer_idx_); + dump_attention_tensor(k_reshaped, "k_rope", layer_idx_); + } auto attn_output = attn_->forward(q_reshaped, k_reshaped, v_reshaped); - attn_output = infinicore::op::mul(attn_output, infinicore::op::sigmoid(gate)->view(attn_output->shape())); - return o_proj_->forward(attn_output); + if (dump_attention) { + dump_attention_tensor(attn_output, "core_output", layer_idx_); + } + auto gate_sigmoid = infinicore::op::sigmoid(gate)->view(attn_output->shape()); + if (dump_attention) { + dump_attention_tensor(gate_sigmoid, "gate_sigmoid", layer_idx_); + } + attn_output = infinicore::op::mul(attn_output, gate_sigmoid); + if (dump_attention) { + dump_attention_tensor(attn_output, "gated_output", layer_idx_); + } + auto projected = o_proj_->forward(attn_output); + if (dump_attention) { + dump_attention_tensor(projected, "projected_output", layer_idx_); + } + return projected; } } // namespace infinilm::models::qwen3_5 diff --git a/csrc/models/qwen3_5/qwen3_5_decoderLayer.cpp b/csrc/models/qwen3_5/qwen3_5_decoderLayer.cpp index 70964bb69..810678f8f 100644 --- a/csrc/models/qwen3_5/qwen3_5_decoderLayer.cpp +++ b/csrc/models/qwen3_5/qwen3_5_decoderLayer.cpp @@ -1,10 +1,45 @@ #include "qwen3_5_decoderLayer.hpp" #include "infinicore/ops.hpp" +#include "infinicore/ops/add_rms_norm.hpp" +#include "infinicore/ops/cast.hpp" +#include #include #include #include namespace infinilm::models::qwen3_5 { +namespace { + +void dump_prefill_tensor(const infinicore::Tensor &tensor, + const std::string &filename) { + const char *dump_dir = std::getenv("INFINILM_LAYER_DUMP_DIR"); + if (dump_dir == nullptr || dump_dir[0] == '\0' || !tensor) { + return; + } + const char *dump_numel = std::getenv("INFINILM_LAYER_DUMP_NUMEL"); + if (dump_numel == nullptr || dump_numel[0] == '\0' + || tensor->numel() != std::strtoull(dump_numel, nullptr, 10)) { + return; + } + tensor->debug(std::string(dump_dir) + "/" + filename); +} + +bool should_dump_layer(size_t layer_idx) { + const char *first_n = std::getenv("INFINILM_LAYER_DUMP_FIRST_N"); + if (first_n != nullptr && first_n[0] != '\0' + && layer_idx < std::strtoull(first_n, nullptr, 10)) { + return true; + } + return (layer_idx + 1) % 8 == 0; +} + +bool should_dump_operators(size_t layer_idx) { + const char *target = std::getenv("INFINILM_OPERATOR_DUMP_LAYER"); + return target != nullptr && target[0] != '\0' + && layer_idx == std::strtoull(target, nullptr, 10); +} + +} // namespace Qwen35DecoderLayer::Qwen35DecoderLayer(std::shared_ptr model_config, size_t layer_idx, @@ -17,7 +52,9 @@ Qwen35DecoderLayer::Qwen35DecoderLayer(std::shared_ptr layer_types = model_config->get>("layer_types"); layer_type_ = layer_types[layer_idx]; @@ -33,15 +70,105 @@ Qwen35DecoderLayer::Qwen35DecoderLayer(std::shared_ptr Qwen35DecoderLayer::forward(const infinicore::Tensor &positions, infinicore::Tensor &hidden_states, infinicore::Tensor &residual) { - input_layernorm_->forward_inplace(hidden_states, residual); + if (layer_idx_ == 0) { + dump_prefill_tensor(hidden_states, "infini_embed.bin"); + } + if (residual + && hidden_states->dtype() == infinicore::DataType::F32 + && residual->dtype() == infinicore::DataType::BF16) { + auto y = infinicore::Tensor::empty( + hidden_states->shape(), infinicore::DataType::BF16, hidden_states->device()); + auto residual_out = infinicore::Tensor::empty( + residual->shape(), infinicore::DataType::BF16, residual->device()); + infinicore::op::add_rms_norm_( + y, residual_out, hidden_states, residual, + input_layernorm_->weight(), + static_cast(input_layernorm_->eps())); + hidden_states = y; + residual = residual_out; + } else { + input_layernorm_->forward_inplace(hidden_states, residual); + } if ("linear_attention" == layer_type_) { hidden_states = linear_attn_->forward(hidden_states); } else if ("full_attention" == layer_type_) { hidden_states = self_attn_->forward(positions, hidden_states); } - post_attention_layernorm_->forward_inplace(hidden_states, residual); + const char *fp32_fused_env = std::getenv("INFINILM_POST_NORM_FP32_FUSED"); + const bool fp32_fused = fp32_fused_env != nullptr && fp32_fused_env[0] != '\0' + && std::string(fp32_fused_env) != "0"; + const bool mixed_gguf_f32 = residual + && hidden_states->dtype() == infinicore::DataType::F32 + && residual->dtype() == infinicore::DataType::BF16; + if (mixed_gguf_f32) { + auto y = infinicore::Tensor::empty( + hidden_states->shape(), infinicore::DataType::BF16, hidden_states->device()); + auto residual_out = infinicore::Tensor::empty( + residual->shape(), infinicore::DataType::BF16, residual->device()); + infinicore::op::add_rms_norm_( + y, residual_out, hidden_states, residual, + post_attention_layernorm_->weight(), + static_cast(post_attention_layernorm_->eps())); + hidden_states = y; + residual = residual_out; + } else if (fp32_fused) { + auto a32 = infinicore::Tensor::empty(hidden_states->shape(), infinicore::DataType::F32, hidden_states->device()); + auto b32 = infinicore::Tensor::empty(residual->shape(), infinicore::DataType::F32, residual->device()); + infinicore::op::cast_(a32, hidden_states); + infinicore::op::cast_(b32, residual); + auto y32 = infinicore::Tensor::empty(hidden_states->shape(), infinicore::DataType::F32, hidden_states->device()); + auto r32 = infinicore::Tensor::empty(residual->shape(), infinicore::DataType::F32, residual->device()); + infinicore::op::add_rms_norm_(y32, r32, a32, b32, + post_attention_layernorm_->weight(), + static_cast(post_attention_layernorm_->eps())); + hidden_states = y32; + residual = r32; + } else { + post_attention_layernorm_->forward_inplace(hidden_states, residual); + } + if (should_dump_operators(layer_idx_)) { + dump_prefill_tensor(residual, + "infini_attn_residual_" + std::to_string(layer_idx_) + ".bin"); + dump_prefill_tensor(hidden_states, + "infini_attn_post_norm_" + std::to_string(layer_idx_) + ".bin"); + } + const char *fp32_mlp_env = std::getenv("INFINILM_POST_NORM_FP32_MLP"); + const bool fp32_mlp = fp32_mlp_env != nullptr && fp32_mlp_env[0] != '\0' + && std::string(fp32_mlp_env) != "0"; + if (fp32_mlp && !fp32_fused) { + auto fp32_hidden = infinicore::Tensor::empty( + hidden_states->shape(), infinicore::DataType::F32, hidden_states->device()); + infinicore::op::cast_(fp32_hidden, hidden_states); + hidden_states = fp32_hidden; + } hidden_states = mlp_->forward(hidden_states); + if (should_dump_operators(layer_idx_)) { + dump_prefill_tensor(hidden_states, + "infini_ffn_out_" + std::to_string(layer_idx_) + ".bin"); + } + if (fp32_mlp && !fp32_fused) { + auto bf16_hidden = infinicore::Tensor::empty( + hidden_states->shape(), infinicore::DataType::BF16, hidden_states->device()); + infinicore::op::cast_(bf16_hidden, hidden_states); + hidden_states = bf16_hidden; + } + if (should_dump_layer(layer_idx_)) { + auto materialized = residual ? infinicore::op::add(residual, hidden_states) + : hidden_states; + dump_prefill_tensor(materialized, + "infini_layer_" + std::to_string(layer_idx_) + "_post_ffn.bin"); + } + if (fp32_fused) { + auto bf16_hidden = infinicore::Tensor::empty( + hidden_states->shape(), infinicore::DataType::BF16, hidden_states->device()); + infinicore::op::cast_(bf16_hidden, hidden_states); + hidden_states = bf16_hidden; + auto bf16_residual = infinicore::Tensor::empty( + residual->shape(), infinicore::DataType::BF16, residual->device()); + infinicore::op::cast_(bf16_residual, residual); + residual = bf16_residual; + } return std::make_tuple(hidden_states, residual); } diff --git a/csrc/models/qwen3_5/qwen3_5_for_causal_lm.cpp b/csrc/models/qwen3_5/qwen3_5_for_causal_lm.cpp index 72fe1a87f..1105b2d24 100644 --- a/csrc/models/qwen3_5/qwen3_5_for_causal_lm.cpp +++ b/csrc/models/qwen3_5/qwen3_5_for_causal_lm.cpp @@ -1,12 +1,16 @@ #include "qwen3_5_for_causal_lm.hpp" #include "../models_registry.hpp" +#include "infinicore/ops/gemm.hpp" +#include #include #include #include namespace infinilm::models::qwen3_5 { +// TextModel diagnostic hooks are compiled into this Qwen3.5 translation unit. + Qwen35ForCausalLM::Qwen35ForCausalLM( std::shared_ptr model_config, const infinicore::Device &device) { @@ -14,6 +18,9 @@ Qwen35ForCausalLM::Qwen35ForCausalLM( const size_t hidden_size = model_config->get("hidden_size"); const size_t vocab_size = model_config->get("vocab_size"); const auto &dtype = model_config->get_dtype(); + fp32_lm_head_output_ = model_config->get_config_json().value( + "lm_head_output_dtype", std::string()) + == "float32"; INFINICORE_NN_MODULE_INIT(model, model_config, device); INFINICORE_NN_MODULE_INIT( @@ -23,7 +30,40 @@ Qwen35ForCausalLM::Qwen35ForCausalLM( InfinilmModel::Output Qwen35ForCausalLM::forward( const InfinilmModel::Input &input) const { auto hidden_states = model_->forward(input); - return {lm_head_->forward(hidden_states)}; + const char *dump_dir = std::getenv("INFINILM_LAYER_DUMP_DIR"); + const char *dump_numel = std::getenv("INFINILM_LAYER_DUMP_NUMEL"); + if (dump_dir != nullptr && dump_dir[0] != '\0' + && dump_numel != nullptr && dump_numel[0] != '\0' + && hidden_states->numel() + == std::strtoull(dump_numel, nullptr, 10)) { + hidden_states->debug( + std::string(dump_dir) + "/infini_result_norm.bin"); + } + infinicore::Tensor logits; + if (fp32_lm_head_output_) { + auto hidden = hidden_states->is_contiguous() + ? hidden_states + : hidden_states->contiguous(); + const size_t ndim = hidden->ndim(); + auto output_shape = hidden->shape(); + output_shape[ndim - 1] = lm_head_->out_features(); + logits = infinicore::Tensor::empty( + output_shape, infinicore::DataType::F32, hidden->device()); + size_t rows = 1; + for (size_t i = 0; i + 1 < ndim; ++i) { + rows *= hidden->shape()[i]; + } + auto weight = lm_head_->weight()->contiguous(); + infinicore::op::gemm_( + logits->view({rows, lm_head_->out_features()}), + hidden->view({rows, lm_head_->in_features()}), + weight->permute({1, 0}), + 1.0f, + 0.0f); + } else { + logits = lm_head_->forward(hidden_states); + } + return {logits, hidden_states}; } void Qwen35ForCausalLM::reset_cache( diff --git a/csrc/models/qwen3_5/qwen3_5_for_causal_lm.hpp b/csrc/models/qwen3_5/qwen3_5_for_causal_lm.hpp index 51211481f..e610f0598 100644 --- a/csrc/models/qwen3_5/qwen3_5_for_causal_lm.hpp +++ b/csrc/models/qwen3_5/qwen3_5_for_causal_lm.hpp @@ -16,6 +16,7 @@ class Qwen35ForCausalLM : public InfinilmModel { void reset_cache(const cache::CacheConfig *cache_config) override; protected: + bool fp32_lm_head_output_{false}; INFINICORE_NN_MODULE(Qwen35Model, model); INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, lm_head); }; diff --git a/csrc/models/qwen3_5/qwen3_5_fused_qkv_linear.cpp b/csrc/models/qwen3_5/qwen3_5_fused_qkv_linear.cpp index 65409bc18..1b1cdfb25 100644 --- a/csrc/models/qwen3_5/qwen3_5_fused_qkv_linear.cpp +++ b/csrc/models/qwen3_5/qwen3_5_fused_qkv_linear.cpp @@ -14,16 +14,19 @@ Qwen35FusedQKVLinear::Qwen35FusedQKVLinear(size_t hidden_size, bool bias, const infinicore::DataType &dtype, const infinicore::Device &device, - engine::distributed::RankInfo rank_info) + engine::distributed::RankInfo rank_info, + const std::string &prefix) : infinilm::layers::linear::ColumnParallelLinear( - hidden_size, - num_q_head * head_dim * 2 + num_kv_head * head_dim * calculate_kv_replicas(num_kv_head, rank_info.tp_size) * 2, - quantization == nullptr ? std::make_shared() : quantization, - bias, - dtype, - device, - rank_info.tp_rank, - rank_info.tp_size), + hidden_size, + num_q_head * head_dim * 2 + num_kv_head * head_dim * calculate_kv_replicas(num_kv_head, rank_info.tp_size) * 2, + quantization == nullptr ? std::make_shared() : quantization, + bias, + dtype, + device, + rank_info.tp_rank, + rank_info.tp_size, + -1, + prefix), head_dim_(head_dim), local_num_q_heads_(num_q_head / tp_size_), q_proj_out_size_(num_q_head * head_dim * 2 / tp_size_), @@ -32,12 +35,39 @@ Qwen35FusedQKVLinear::Qwen35FusedQKVLinear(size_t hidden_size, v_out_size_(calculate_kv_replicas(num_kv_head, rank_info.tp_size) * num_kv_head * head_dim / tp_size_), num_kv_head_(num_kv_head), register_fn_(register_fn) { - split_infos_ = { - {q_name, 0, q_proj_out_size_, 0}, - {k_name, q_proj_out_size_, k_out_size_, num_kv_head_}, - {v_name, q_proj_out_size_ + k_out_size_, v_out_size_, num_kv_head_}, - }; - auto params = this->split_params(split_infos_, tp_rank_, tp_size_, num_kv_head_); + if (this->sharded_) { + // GGUF stores three separate checkpoint tensors. The Q projection + // contains the interleaved gate, and there is no fused buffer to + // narrow. Each stem must end in `.` and match the type-table key. + if (prefix.empty()) { + throw std::runtime_error( + "Qwen35FusedQKVLinear requires a layer prefix such as " + "`layers.3.self_attn` for GGUF quantization."); + } + shard_specs_ = { + {q_name, q_proj_out_size_, prefix + "." + q_name + "."}, + {k_name, k_out_size_, prefix + "." + k_name + "."}, + {v_name, v_out_size_, prefix + "." + v_name + "."}, + }; + } else { + split_infos_ = { + {q_name, 0, q_proj_out_size_, 0}, + {k_name, q_proj_out_size_, k_out_size_, num_kv_head_}, + {v_name, q_proj_out_size_ + k_out_size_, v_out_size_, num_kv_head_}, + }; + } + register_fused_params(); +} + +void Qwen35FusedQKVLinear::register_fused_params() { + if (!register_fn_) { + return; + } + // The GGUF path runs one GEMM per shard and concatenates the results into + // `[B, S, Q|K|V]`, so the existing `forward_split()` offsets remain valid. + auto params = this->sharded_ + ? this->init_fused_shards(shard_specs_) + : this->split_params(split_infos_, tp_rank_, tp_size_, num_kv_head_); for (auto &sp : params) { register_fn_(sp.full_name, std::move(sp.param)); } @@ -62,11 +92,11 @@ Qwen35FusedQKVLinear::forward_split(infinicore::Tensor &input) { void Qwen35FusedQKVLinear::process_weights_after_loading() { BaseLinear::process_weights_after_loading(); + // `split_infos_` is empty for sharded layouts such as GGUF because the + // shard parameters are the load targets. Reallocation would discard the + // block bytes that were already loaded. if (register_fn_ && !split_infos_.empty()) { - auto params = this->split_params(split_infos_, tp_rank_, tp_size_, num_kv_head_); - for (auto &sp : params) { - register_fn_(sp.full_name, std::move(sp.param)); - } + register_fused_params(); } } diff --git a/csrc/models/qwen3_5/qwen3_5_fused_qkv_linear.hpp b/csrc/models/qwen3_5/qwen3_5_fused_qkv_linear.hpp index 55d2b762e..d0f756c30 100644 --- a/csrc/models/qwen3_5/qwen3_5_fused_qkv_linear.hpp +++ b/csrc/models/qwen3_5/qwen3_5_fused_qkv_linear.hpp @@ -18,7 +18,8 @@ class Qwen35FusedQKVLinear : public infinilm::layers::linear::ColumnParallelLine bool bias = false, const infinicore::DataType &dtype = infinicore::DataType::F32, const infinicore::Device &device = infinicore::Device(), - engine::distributed::RankInfo rank_info = engine::distributed::RankInfo()); + engine::distributed::RankInfo rank_info = engine::distributed::RankInfo(), + const std::string &prefix = ""); void process_weights_after_loading() override; @@ -45,6 +46,11 @@ class Qwen35FusedQKVLinear : public infinilm::layers::linear::ColumnParallelLine size_t num_kv_head_; infinilm::layers::linear::RegisterParamFn register_fn_; std::vector split_infos_; + // GGUF Q|gate, K, and V shards may use different GGML types, so each + // shard owns a separate buffer instead of using `split_infos_`. + std::vector shard_specs_; + + void register_fused_params(); }; } // namespace infinilm::models::qwen3_5 diff --git a/csrc/models/qwen3_next/qwen3_next_gated_deltanet.cpp b/csrc/models/qwen3_next/qwen3_next_gated_deltanet.cpp index 022454247..eac3ae58f 100644 --- a/csrc/models/qwen3_next/qwen3_next_gated_deltanet.cpp +++ b/csrc/models/qwen3_next/qwen3_next_gated_deltanet.cpp @@ -10,11 +10,38 @@ #include #include +#include #include #include +#include #include namespace infinilm::models::qwen3_next { +namespace { + +bool should_dump_gdn(size_t layer_idx, size_t seq_len) { + const char *target_layer = std::getenv("INFINILM_GDN_DUMP_LAYER"); + const char *target_seq_len = std::getenv("INFINILM_GDN_DUMP_SEQ_LEN"); + return target_layer != nullptr && target_layer[0] != '\0' + && target_seq_len != nullptr && target_seq_len[0] != '\0' + && layer_idx == std::strtoull(target_layer, nullptr, 10) + && seq_len == std::strtoull(target_seq_len, nullptr, 10); +} + +void dump_gdn_tensor(const infinicore::Tensor &tensor, + const std::string &name, + size_t layer_idx, + size_t seq_len) { + const char *dump_dir = std::getenv("INFINILM_LAYER_DUMP_DIR"); + if (dump_dir == nullptr || dump_dir[0] == '\0' || !tensor + || !should_dump_gdn(layer_idx, seq_len)) { + return; + } + tensor->debug(std::string(dump_dir) + "/infini_gdn_" + name + "_" + + std::to_string(layer_idx) + ".bin"); +} + +} // namespace Qwen3NextCausalConv1D::Qwen3NextCausalConv1D(std::shared_ptr model_config, size_t layer_idx, @@ -125,12 +152,15 @@ Qwen3NextGatedDeltaNet::Qwen3NextGatedDeltaNet(std::shared_ptrget_quantization_method(); auto register_fn = [this](const std::string &n, infinicore::nn::Parameter p) { this->register_parameter(n, std::move(p)); }; + // Checkpoint path for this module. This is used only by quantization + // schemes, such as GGUF, that resolve layouts by tensor name. + const std::string prefix = "layers." + std::to_string(layer_idx_) + ".linear_attn"; in_proj_qkv_ = std::make_shared( hidden_size, linear_key_head_dim, linear_key_head_dim, linear_value_head_dim, linear_num_key_heads, linear_num_key_heads, linear_num_value_heads, false, false, false, "in_proj_q", "in_proj_k", "in_proj_v", register_fn, - quantization_method, dtype, device, rank_info); - in_proj_z_ = this->register_module("in_proj_z", hidden_size, value_dim, false, dtype, device, tp_rank, tp_size); + quantization_method, dtype, device, rank_info, prefix); + in_proj_z_ = this->register_module("in_proj_z", hidden_size, value_dim, quantization_method, false, dtype, device, tp_rank, tp_size, -1, prefix + ".in_proj_z."); in_proj_a_ = this->register_module("in_proj_a", hidden_size, linear_num_value_heads, false, dtype, device, tp_rank, tp_size); in_proj_b_ = this->register_module("in_proj_b", hidden_size, linear_num_value_heads, false, dtype, device, tp_rank, tp_size); @@ -140,7 +170,8 @@ Qwen3NextGatedDeltaNet::Qwen3NextGatedDeltaNet(std::shared_ptrregister_module( "out_proj", value_dim, hidden_size, quantization_method, - false, dtype, device, rank_info.tp_rank, rank_info.tp_size, rank_info.comm); + false, dtype, device, rank_info.tp_rank, rank_info.tp_size, rank_info.comm, + prefix + ".out_proj."); } infinicore::Tensor Qwen3NextGatedDeltaNet::forward(const infinicore::Tensor &hidden_states) const { @@ -154,11 +185,16 @@ infinicore::Tensor Qwen3NextGatedDeltaNet::forward(const infinicore::Tensor &hid auto z = in_proj_z_->forward(hidden_states_mutable); auto a = in_proj_a_->forward(hidden_states_mutable); auto b = in_proj_b_->forward(hidden_states_mutable); + dump_gdn_tensor(qkv, "qkv_mixed", layer_idx_, seq_len); + dump_gdn_tensor(z, "z", layer_idx_, seq_len); + dump_gdn_tensor(a, "alpha", layer_idx_, seq_len); + dump_gdn_tensor(b, "beta", layer_idx_, seq_len); auto &forward_context = infinilm::global_state::get_forward_context(); auto &mamba_metadata = forward_context.mamba_metadata; auto conv_qkv = this->conv1d_->forward(qkv); + dump_gdn_tensor(conv_qkv, "conv_output_silu", layer_idx_, seq_len); auto q = conv_qkv->narrow({{2, 0, local_key_dim_}}); auto k = conv_qkv->narrow({{2, local_key_dim_, local_key_dim_}}); @@ -184,6 +220,8 @@ infinicore::Tensor Qwen3NextGatedDeltaNet::forward(const infinicore::Tensor &hid {seq_len, 1, local_num_value_heads_}, {b->stride(1), b->stride(0), 1}); auto [g, beta] = infinicore::op::fused_gated_delta_net_gating(A_log_, a_heads, b_heads, dt_bias_); + dump_gdn_tensor(g, "gate", layer_idx_, seq_len); + dump_gdn_tensor(beta, "beta_sigmoid", layer_idx_, seq_len); delta_out = infinicore::op::recurrent_gated_delta_rule_indexed( q_delta, @@ -217,6 +255,8 @@ infinicore::Tensor Qwen3NextGatedDeltaNet::forward(const infinicore::Tensor &hid {1, seq_len, local_num_value_heads_}, {b->stride(0), b->stride(1), 1}); auto [g, beta] = infinicore::op::fused_gated_delta_net_gating(A_log_, a_heads, b_heads, dt_bias_); + dump_gdn_tensor(g, "gate", layer_idx_, seq_len); + dump_gdn_tensor(beta, "beta_sigmoid", layer_idx_, seq_len); delta_out = infinicore::op::chunk_gated_delta_rule( q_delta, @@ -237,12 +277,18 @@ infinicore::Tensor Qwen3NextGatedDeltaNet::forward(const infinicore::Tensor &hid auto delta_out_2d = delta_out->as_strided( {batch_size * seq_len * local_num_value_heads_, value_head_dim_}, {static_cast(value_head_dim_), 1}); + dump_gdn_tensor(delta_out, "delta_out", layer_idx_, seq_len); auto v_norm_2d = norm_->forward(delta_out_2d); auto v_norm = v_norm_2d->as_strided( {batch_size, seq_len, local_value_dim_}, {static_cast(seq_len * local_value_dim_), static_cast(local_value_dim_), 1}); + dump_gdn_tensor(v_norm, "v_norm", layer_idx_, seq_len); auto gated = infinicore::op::mul(v_norm, infinicore::op::silu(z)); - return out_proj_->forward(gated); + dump_gdn_tensor(gated, "gated", layer_idx_, seq_len); + dump_gdn_tensor(gated, "final_output", layer_idx_, seq_len); + auto output = out_proj_->forward(gated); + dump_gdn_tensor(output, "linear_attn_out", layer_idx_, seq_len); + return output; } } // namespace infinilm::models::qwen3_next diff --git a/csrc/pybind11/engine/engine.hpp b/csrc/pybind11/engine/engine.hpp index c5e85577c..5e8a6caea 100644 --- a/csrc/pybind11/engine/engine.hpp +++ b/csrc/pybind11/engine/engine.hpp @@ -122,18 +122,14 @@ inline void bind_infer_engine(py::module &m) { return state_dict_tp_all; }) .def("process_weights_after_loading", &InferEngine::process_weights_after_loading, "Process the weights after loading on all workers (e.g., for quantization)") - .def( - "forward", [](InferEngine &self, const InferEngine::Input &input) -> InferEngine::Output { + .def("forward", [](InferEngine &self, const InferEngine::Input &input) -> InferEngine::Output { // IMPORTANT: Release the GIL before calling forward() to allow other Python threads // to run concurrently during inference (which may block for a long time). // Do NOT remove this — without it, the GIL is held throughout inference and will // deadlock or stall any other Python thread (e.g., request handling, scheduling). py::gil_scoped_release release; - return self.forward(input); - }, - "Run inference on all ranks with arbitrary arguments") - .def( - "reset_cache", [](InferEngine &self, std::shared_ptr cfg) { self.reset_cache(cfg ? cfg.get() : nullptr); }, py::arg("cache_config") = py::none()) + return self.forward(input); }, "Run inference on all ranks with arbitrary arguments") + .def("reset_cache", [](InferEngine &self, std::shared_ptr cfg) { self.reset_cache(cfg ? cfg.get() : nullptr); }, py::arg("cache_config") = py::none()) .def("get_kv_cache", &InferEngine::get_kv_cache, "Get per-rank kv cache list") .def("get_cache_config", [](const InferEngine &self) -> std::shared_ptr { auto cfg = self.get_cache_config(); @@ -193,6 +189,7 @@ inline void bind_infer_engine(py::module &m) { "temperature", "top_p", "top_k", + "suppressed_token_ids", }; for (auto &item : kwargs) { @@ -209,6 +206,8 @@ inline void bind_infer_engine(py::module &m) { input.top_p = py::cast(item.second); } else if (key == "top_k") { input.top_k = py::cast(item.second); + } else if (key == "suppressed_token_ids") { + input.suppressed_token_ids = py::cast>>(item.second); } } @@ -250,6 +249,7 @@ inline void bind_infer_engine(py::module &m) { .def_readwrite("visual_token_ranges", &InferEngine::Input::visual_token_ranges) .def_readwrite("target_hidden_states", &InferEngine::Input::target_hidden_states) .def_readwrite("sample_all_positions", &InferEngine::Input::sample_all_positions) + .def_readwrite("suppressed_token_ids", &InferEngine::Input::suppressed_token_ids) .def_readwrite("temperature", &InferEngine::Input::temperature) .def_readwrite("top_k", &InferEngine::Input::top_k) .def_readwrite("top_p", &InferEngine::Input::top_p); diff --git a/python/infinilm/infer_engine.py b/python/infinilm/infer_engine.py index 117d82f9a..6db40d8ad 100644 --- a/python/infinilm/infer_engine.py +++ b/python/infinilm/infer_engine.py @@ -124,6 +124,7 @@ class GenerationConfig: eos_token_id: list[int] | None = None stop_on_eos: bool = True + ignore_eos: bool = False def _infer_position_id_axes(hf_config: dict) -> int: @@ -275,6 +276,7 @@ def _build_input( visual_token_ranges=None, target_hidden_states=None, sample_all_positions=False, + suppressed_token_ids=None, temperature=None, top_k=None, top_p=None, @@ -333,6 +335,9 @@ def convert_tensor_list(tensor_list_): visual_token_ranges=visual_token_ranges, target_hidden_states=target_hidden_states, sample_all_positions=sample_all_positions, + suppressed_token_ids=( + [] if suppressed_token_ids is None else suppressed_token_ids + ), temperature=temperature, top_k=top_k, top_p=top_p, @@ -358,6 +363,7 @@ def forward( image_req_ids=None, visual_token_ranges=None, target_hidden_states=None, + suppressed_token_ids=None, temperature=None, top_k=None, top_p=None, @@ -430,6 +436,7 @@ def convert_tensor_list(tensor_list_): image_req_ids=image_req_ids, visual_token_ranges=visual_token_ranges, target_hidden_states=target_hidden_states, + suppressed_token_ids=suppressed_token_ids, temperature=temperature, top_k=top_k, top_p=top_p, @@ -459,6 +466,7 @@ def forward_raw( visual_token_ranges=None, target_hidden_states=None, sample_all_positions=True, + suppressed_token_ids=None, temperature=None, top_k=None, top_p=None, @@ -481,6 +489,7 @@ def forward_raw( visual_token_ranges=visual_token_ranges, target_hidden_states=target_hidden_states, sample_all_positions=sample_all_positions, + suppressed_token_ids=suppressed_token_ids, temperature=temperature, top_k=top_k, top_p=top_p, @@ -509,7 +518,11 @@ def generate( position_id_delta=0, _measure_and_log_time=False, ): - eos_token_id = self.eos_token_id + eos_token_id = generation_config.eos_token_id + if eos_token_id is None: + eos_token_id = self.eos_token_id + elif isinstance(eos_token_id, int): + eos_token_id = [eos_token_id] past_seq_len = 0 output_ids = [] @@ -687,6 +700,11 @@ def generate( tgt_sizes=tgt_sizes if iter == 0 else None, image_grid_thw=image_grid_thw if iter == 0 else None, image_req_ids=image_req_ids if iter == 0 else None, + suppressed_token_ids=( + [list(eos_token_id) for _ in range(batch_size)] + if generation_config.ignore_eos + else [] + ), temperature=generation_config.temperature, top_k=generation_config.top_k, top_p=generation_config.top_p, diff --git a/python/infinilm/llm/llm.py b/python/infinilm/llm/llm.py index 59d5a1eca..ffc90c844 100644 --- a/python/infinilm/llm/llm.py +++ b/python/infinilm/llm/llm.py @@ -67,6 +67,12 @@ def __init__(self, config: EngineConfig): # Initialize KV cache based on cache type if config.cache_type == "static": + if has_mamba_cache: + model_type = hf_config["model_type"] + raise RuntimeError( + "Static KV cache is not supported for Mamba-cache model " + f"{model_type!r} yet. Use --cache-type paged instead." + ) self.scheduler = StaticScheduler( max_cache_len=config.max_cache_len, enable_prefix_caching=config.enable_prefix_caching, diff --git a/python/infinilm/llm/model_runner/model_runner.py b/python/infinilm/llm/model_runner/model_runner.py index a1696f848..a642b71b9 100644 --- a/python/infinilm/llm/model_runner/model_runner.py +++ b/python/infinilm/llm/model_runner/model_runner.py @@ -217,6 +217,12 @@ def _model_forward(self, scheduler_output): self.config.top_p, self.config.top_k, ) + model_input["suppressed_token_ids"] = [ + list(req.eos_token_ids or self.model_engine.eos_token_id or []) + if req.sampling_params.ignore_eos + else [] + for req in scheduler_output.scheduled_requests + ] if self.speculative_runner is not None: return self._model_forward_with_speculative(scheduler_output, model_input) diff --git a/python/infinilm/modeling_utils.py b/python/infinilm/modeling_utils.py index b210b0c56..5e879ac6a 100644 --- a/python/infinilm/modeling_utils.py +++ b/python/infinilm/modeling_utils.py @@ -55,6 +55,15 @@ def parse_dtype(dtype_str: str): "F8_E5M2": torch.float8_e5m2, } +_FP8_DTYPES = tuple( + x + for x in ( + getattr(torch, "float8_e4m3fn", None), + getattr(torch, "float8_e5m2", None), + ) + if x is not None +) + def _is_internal_moe_packed_weight(key: str) -> bool: # InfiniLM registers packed MoE parameters internally. HF checkpoints @@ -129,7 +138,8 @@ def load_state_dict( for k in f.keys(): tensor = f.get_tensor(k) preserve_fp32 = k.endswith(preserve_fp32_suffixes) - if tensor.is_floating_point() and not preserve_fp32: + preserve_fp8 = tensor.dtype in _FP8_DTYPES + if tensor.is_floating_point() and not preserve_fp32 and not preserve_fp8: tensor = tensor.to(device=device, dtype=dtype) else: tensor = tensor.to(device=device) @@ -204,6 +214,10 @@ def load_model_state_dict_by_file( preserve_fp32_suffixes = (".e_score_correction_bias",) if model_type == "kimi_k3": preserve_fp32_suffixes += (".A_log", ".dt_bias") + if model.hf_config.get("lm_head_output_dtype") == "float32": + # The Qwen3.5 GGUF route keeps BF16 head weights but can request an + # FP32 output accumulator; do not downcast a future FP32 head artifact. + preserve_fp32_suffixes += ("lm_head.weight",) torch_device = "cpu" torch_dtype = infinicore.utils.to_torch_dtype(dtype) @@ -252,6 +266,14 @@ def load_model_state_dict_by_file( if remapper is not None: model_param = remapper(model_param, config=model.hf_config) + # Convert FP8 block scales from BF16 to FP32 for CUTLASS GEMM + for key in list(model_param.keys()): + if ( + key.endswith("weight_scale_inv") + and model_param[key].dtype == torch.bfloat16 + ): + model_param[key] = model_param[key].float() + # --------------------------------------------------------- # # Scale embed_tokens on torch side before converting # --------------------------------------------------------- # @@ -324,6 +346,7 @@ def load_model_state_dict_by_file( target_dtype = ( model_params[key].dtype if key.endswith(preserve_fp32_suffixes) + or model_params[key].dtype in _FP8_DTYPES else torch_dtype ) model_param_infini[key] = infinicore.from_torch( @@ -387,7 +410,11 @@ def load_model_state_dict_by_tensor( with safe_open(file_path, "pt", "cpu") as f: for name in f.keys(): - tensor = f.get_tensor(name).to(dtype=torch_dtype) + raw_tensor = f.get_tensor(name) + if raw_tensor.dtype in _FP8_DTYPES: + tensor = raw_tensor.to(device="cpu") + else: + tensor = raw_tensor.to(dtype=torch_dtype) if name == "model.embed_tokens.weight": embed_tokens_torch_unscaled = tensor @@ -407,7 +434,11 @@ def load_model_state_dict_by_tensor( model_params = torch.load(file_path, weights_only=True, map_location="cpu") for key in model_params.keys(): - tensor = model_params[key].to(dtype=torch_dtype) + raw_tensor = model_params[key] + if raw_tensor.dtype in _FP8_DTYPES: + tensor = raw_tensor.to(device="cpu") + else: + tensor = raw_tensor.to(dtype=torch_dtype) if key == "model.embed_tokens.weight": embed_tokens_torch_unscaled = tensor if scale_emb != 1.0: @@ -758,8 +789,21 @@ def _remap_videonsa(state_dict, config=None): def _remap_qwen3_5(state_dict, config): """Apply Qwen3.5-specific load-time weight fixes.""" state_dict = drop_keys(state_dict, ["mtp."]) + + # Filter out visual encoder keys (not used in language-only mode) + state_dict = { + k: v for k, v in state_dict.items() if not k.startswith("model.visual.") + } + llm_config = config["text_config"] key_dim = llm_config["linear_key_head_dim"] * llm_config["linear_num_key_heads"] + block_size = 128 # FP8 block size for scale splitting + + # The llama.cpp converter has already applied `norm.weight + 1` to GGUF + # Route B checkpoints, except for `linear_attn.norm`. The packer preserves + # those values, so applying the remap again would produce `2 + weight`. + # Fused QKV tensors are also split into `in_proj_q/k/v` while packaging. + gguf = (config.get("quantization_config") or {}).get("quant_method", "") == "gguf" norm_weight_suffixes = ( "input_layernorm.weight", @@ -771,9 +815,11 @@ def _remap_qwen3_5(state_dict, config): to_drop = [] to_add = {} for key, tensor in state_dict.items(): - if key == "model.norm.weight" or key.endswith(norm_weight_suffixes): + if not gguf and ( + key == "model.norm.weight" or key.endswith(norm_weight_suffixes) + ): state_dict[key] = tensor + torch.ones_like(tensor) - elif key.endswith("linear_attn.in_proj_qkv.weight"): + elif key.endswith("linear_attn.in_proj_qkv.weight") and not gguf: prefix = key[: -len("in_proj_qkv.weight")] to_add[prefix + "in_proj_q.weight"] = state_dict[key][ :key_dim, : @@ -785,6 +831,22 @@ def _remap_qwen3_5(state_dict, config): key_dim * 2 :, : ].contiguous() to_drop.append(key) + elif key.endswith("linear_attn.in_proj_qkv.weight_scale_inv"): + # Split fused QKV scale into separate q/k/v scales + # Scale shape: [num_out_blocks, num_in_blocks] + # out dim is split: q(key_dim) | k(key_dim) | v(rest) + prefix = key[: -len("in_proj_qkv.weight_scale_inv")] + key_blocks = key_dim // block_size + to_add[prefix + "in_proj_q.weight_scale_inv"] = state_dict[key][ + :key_blocks, : + ].contiguous() + to_add[prefix + "in_proj_k.weight_scale_inv"] = state_dict[key][ + key_blocks : key_blocks * 2, : + ].contiguous() + to_add[prefix + "in_proj_v.weight_scale_inv"] = state_dict[key][ + key_blocks * 2 :, : + ].contiguous() + to_drop.append(key) state_dict = drop_keys(state_dict, to_drop) state_dict.update(to_add) @@ -850,6 +912,7 @@ def _remap_ernie4_5_moe_vl(state_dict, config=None): if ( key.endswith((".mlp.gate.weight", ".mlp.gate.weight_1")) and tensor.is_floating_point() + and tensor.dtype not in _FP8_DTYPES ): remapped[key] = tensor.to(dtype=target_dtype).contiguous() else: @@ -896,13 +959,14 @@ def fuse_expert_group(expert_ids): b1_tensors.append(torch.cat([gate_bias, up_bias], dim=0)) b2_tensors.append(down_bias) + fused_dtype = ( + w1_tensors[0].dtype + if w1_tensors[0].dtype in _FP8_DTYPES + else target_dtype + ) fused = { - "w1": torch.stack(w1_tensors, dim=0) - .to(dtype=target_dtype) - .contiguous(), - "w2": torch.stack(w2_tensors, dim=0) - .to(dtype=target_dtype) - .contiguous(), + "w1": torch.stack(w1_tensors, dim=0).to(dtype=fused_dtype).contiguous(), + "w2": torch.stack(w2_tensors, dim=0).to(dtype=fused_dtype).contiguous(), } if has_all_bias: fused["b1"] = ( diff --git a/scripts/gguf_mapping.py b/scripts/gguf_mapping.py new file mode 100644 index 000000000..53acdfa0f --- /dev/null +++ b/scripts/gguf_mapping.py @@ -0,0 +1,534 @@ +#!/usr/bin/env python3 +"""Single source of truth for GGUF-to-InfiniLM tensor mapping. + +InfiniLM weights use [out, in] orientation, matching GGUF packed row order, so +conversion does not transpose weight data. Transform semantics follow +llama.cpp's Qwen conversion: recover A_log from -exp(A_log), preserve baked +normalization offsets, restore grouped value-head rows, and describe runtime +activation permutation for column-reordered output projections. +""" + +from __future__ import annotations + +import re +from dataclasses import dataclass + +# --------------------------------------------------------------------------- +# Transform semantics +# --------------------------------------------------------------------------- +T_NONE = "" # Preserve packed bytes, or only cast dense values. +T_VROWS = "vrows" # Restore value-head row blocks to grouped order. +T_VELEM = "velem" # One scalar per head; implemented by the same row transform. +T_ALOG = "alog" # Recover A_log = log(-ssm_a), then permute. +T_DENSE = "dense" # Dequantize to BF16 for parameters without a packed path. + +# Value-head permutation scope along dimension 0: +# all = the entire dimension contains value heads +# v_tail = only the trailing value_dim segment contains value heads +VPERM_ALL, VPERM_TAIL = "all", "v_tail" + +# Checkpoint suffix shared with the C++ packed-weight layout. +BLOB_SUFFIX = "weight_bytes" + +# Both forms share one implementation; elements per head are derived from shape. +VPERM_TRANSFORMS = (T_VROWS, T_VELEM) + + +def needs_vperm(e: "Entry") -> bool: + return bool(set(e.transforms) & set(VPERM_TRANSFORMS)) + + +# --------------------------------------------------------------------------- +# GGML type ids from ggml.h. Keep this module independent of gguf-py. +# --------------------------------------------------------------------------- +F32, Q8_0, Q4_K, Q5_K, Q6_K = "F32", "Q8_0", "Q4_K", "Q5_K", "Q6_K" +IQ4_NL, IQ4_XS = "IQ4_NL", "IQ4_XS" + +# Packed block types supported by the runtime kernel. +NATIVE_BLOB_TYPES = (Q8_0, Q4_K, Q5_K, Q6_K) +# I-quants converted to dense BF16 until native kernels are available. +V1_IQUANT_DENSE = (IQ4_NL, IQ4_XS) +DENSE_SRC_TYPES = (F32, Q8_0, Q6_K) + V1_IQUANT_DENSE + + +def apply_v1_exceptions(plan, gguf_types, enabled=True): + """Convert I-quant entries without native kernels to dense BF16 in place. + + ``gguf_types`` maps tensor names to GGML type names collected by the caller. + Set ``enabled=False`` when native IQ4 kernels become available. + """ + n = 0 + if enabled: + for e in plan: + if e.blob and gguf_types.get(e.gguf) in V1_IQUANT_DENSE: + e.blob = False + e.transforms = e.transforms + (T_DENSE,) + e.note = ( + (e.note + "；" if e.note else "") + + "dense fallback for source %s; remove when a native kernel is available" + % gguf_types[e.gguf] + ) + n += 1 + return n + + +@dataclass +class Entry: + """Map one GGUF tensor to one InfiniLM parameter.""" + + infinilm: str # InfiniLM parameter name including model.language_model prefix. + gguf: str # GGUF tensor name. + shape: tuple # Full InfiniLM shape before TP, in [out, in] orientation. + blob: bool # Preserve original blocks as U8 [out, row_bytes]. + transforms: tuple = () + types: tuple = () # Allowed source types; empty accepts any reported type. + slices: tuple = () # [start, end) ranges along output dimension. + vperm: str = VPERM_ALL # Scope for T_VROWS. + # Column permutations cross quantization blocks and would require + # requantization. Record them as runtime activation-permutation rules rather + # than conversion-time transforms. + act_vperm: bool = False + note: str = "" + + +@dataclass +class Dims: + hidden: int + n_q_heads: int + n_kv_heads: int + head_dim: int + ffn: int + lin_k_heads: int + lin_v_heads: int + lin_k_dim: int + lin_v_dim: int + conv_kernel: int + vocab: int + n_layers: int + interval: int + mrope_section: tuple = (11, 11, 10) + rope_theta: float = 1e7 + partial_rotary_factor: float = 0.25 + rms_norm_eps: float = 1e-6 + max_position_embeddings: int = 262144 + architectures: str = "Qwen3_5ForConditionalGeneration" + + # Derived dimensions. + @property + def q_rows(self) -> int: # q_proj rows with interleaved query and gate values. + return self.n_q_heads * self.head_dim * 2 + + @property + def kv_rows(self) -> int: + return self.n_kv_heads * self.head_dim + + @property + def o_in(self) -> int: + return self.n_q_heads * self.head_dim + + @property + def key_dim(self) -> int: + return self.lin_k_heads * self.lin_k_dim + + @property + def value_dim(self) -> int: + return self.lin_v_heads * self.lin_v_dim + + @property + def qkv_rows(self) -> int: # Fused q | k | v rows matching GGUF attn_qkv. + return self.key_dim * 2 + self.value_dim + + @property + def conv_channels(self) -> int: + return self.qkv_rows + + @property + def v_per_k(self) -> int: + return self.lin_v_heads // self.lin_k_heads + + def layer_types(self) -> list: + """Match prepare_qwen3_5_model_config: (i + 1) % interval == 0.""" + return [ + "full_attention" if (i + 1) % self.interval == 0 else "linear_attention" + for i in range(self.n_layers) + ] + + +REAL = Dims( + hidden=5120, + n_q_heads=24, + n_kv_heads=4, + head_dim=256, + ffn=17408, + lin_k_heads=16, + lin_v_heads=48, + lin_k_dim=128, + lin_v_dim=128, + conv_kernel=4, + vocab=248320, + n_layers=64, + interval=4, +) + +MINI = Dims( + hidden=512, + n_q_heads=2, + n_kv_heads=1, + head_dim=256, + ffn=1024, + lin_k_heads=2, + lin_v_heads=6, + lin_k_dim=128, + lin_v_dim=128, + conv_kernel=4, + vocab=1024, + n_layers=8, + interval=4, +) + + +PREFIX = "model.language_model." + + +def layer_entries(d: Dims, i: int, role: str) -> list: + """Return entries for layer ``i`` and its attention role. + + Source types are discovered from the input because the same suffix may use + different quantization types across layers. + """ + L = f"{PREFIX}layers.{i}." + G = f"blk.{i}." + kd, vd = d.key_dim, d.value_dim + out = [ + Entry( + L + "input_layernorm.weight", + G + "attn_norm.weight", + (d.hidden,), + False, + (T_DENSE,), + note="GGUF already contains the baked +1 offset", + ), + Entry( + L + "post_attention_layernorm.weight", + G + "post_attention_norm.weight", + (d.hidden,), + False, + (T_DENSE,), + note="GGUF already contains the baked +1 offset", + ), + Entry( + L + "mlp.gate_proj.weight", G + "ffn_gate.weight", (d.ffn, d.hidden), True + ), + Entry(L + "mlp.up_proj.weight", G + "ffn_up.weight", (d.ffn, d.hidden), True), + Entry( + L + "mlp.down_proj.weight", G + "ffn_down.weight", (d.hidden, d.ffn), True + ), + ] + if role == "full_attention": + out += [ + Entry( + L + "self_attn.q_proj.weight", + G + "attn_q.weight", + (d.q_rows, d.hidden), + True, + (), + note="rows contain interleaved query and gate values per head", + ), + Entry( + L + "self_attn.k_proj.weight", + G + "attn_k.weight", + (d.kv_rows, d.hidden), + True, + ), + Entry( + L + "self_attn.v_proj.weight", + G + "attn_v.weight", + (d.kv_rows, d.hidden), + True, + ), + Entry( + L + "self_attn.o_proj.weight", + G + "attn_output.weight", + (d.hidden, d.o_in), + True, + ), + Entry( + L + "self_attn.q_norm.weight", + G + "attn_q_norm.weight", + (d.head_dim,), + False, + (T_DENSE,), + note="GGUF already contains the baked +1 offset", + ), + Entry( + L + "self_attn.k_norm.weight", + G + "attn_k_norm.weight", + (d.head_dim,), + False, + (T_DENSE,), + note="GGUF already contains the baked +1 offset", + ), + ] + else: + out += [ + Entry( + L + "linear_attn.in_proj_q.weight", + G + "attn_qkv.weight", + (kd, d.hidden), + True, + (), + slices=((0, kd),), + note="attn_qkv rows [0:kd]", + ), + Entry( + L + "linear_attn.in_proj_k.weight", + G + "attn_qkv.weight", + (kd, d.hidden), + True, + (), + slices=((kd, 2 * kd),), + note="attn_qkv rows [kd:2kd]", + ), + Entry( + L + "linear_attn.in_proj_v.weight", + G + "attn_qkv.weight", + (vd, d.hidden), + True, + (T_VROWS,), + slices=((2 * kd, 2 * kd + vd),), + note="attn_qkv rows [2kd:] with tiled-to-grouped value heads", + ), + Entry( + L + "linear_attn.in_proj_z.weight", + G + "attn_gate.weight", + (vd, d.hidden), + True, + (T_VROWS,), + note="row permutation from qwen.py using head_v_dim", + ), + Entry( + L + "linear_attn.in_proj_a.weight", + G + "ssm_alpha.weight", + (d.lin_v_heads, d.hidden), + False, + (T_DENSE, T_VROWS), + note="dense fallback plus qwen.py row permutation with head_dim=1", + ), + Entry( + L + "linear_attn.in_proj_b.weight", + G + "ssm_beta.weight", + (d.lin_v_heads, d.hidden), + False, + (T_DENSE, T_VROWS), + note="dense fallback plus row permutation with head_dim=1", + ), + Entry( + L + "linear_attn.A_log", + G + "ssm_a", + (d.lin_v_heads,), + False, + (T_ALOG, T_VROWS), + note="GGUF stores -exp(A_log); recover it with log(-x)", + ), + Entry( + L + "linear_attn.dt_bias", + G + "ssm_dt.bias", + (d.lin_v_heads,), + False, + (T_VELEM,), + note="per-head qwen.py permutation without changing values", + ), + Entry( + L + "linear_attn.conv1d.weight", + G + "ssm_conv1d.weight", + (d.conv_channels, 1, d.conv_kernel), + False, + (T_DENSE, T_VROWS), + vperm=VPERM_TAIL, + note="restore squeezed [C,K] shape and permute only trailing V channels", + ), + Entry( + L + "linear_attn.norm.weight", + G + "ssm_norm.weight", + (d.lin_v_dim,), + False, + (T_DENSE,), + note="not permuted by qwen.py and no normalization offset", + ), + Entry( + L + "linear_attn.out_proj.weight", + G + "ssm_out.weight", + (d.hidden, vd), + True, + (), + act_vperm=True, + note="column permutation requires grouped-to-tiled runtime activation mapping", + ), + ] + return out + + +def build_plan(d: Dims) -> list: + """Build full-model mapping entries, including root-level tensors.""" + entries = [ + Entry( + PREFIX + "embed_tokens.weight", + "token_embd.weight", + (d.vocab, d.hidden), + False, + (T_DENSE,), + note="dequantize embedding to dense BF16", + ), + ] + for i, role in enumerate(d.layer_types()): + entries += layer_entries(d, i, role) + entries += [ + Entry( + PREFIX + "norm.weight", + "output_norm.weight", + (d.hidden,), + False, + (T_DENSE,), + note="GGUF already contains the baked +1 offset", + ), + Entry( + "lm_head.weight", + "output.weight", + (d.vocab, d.hidden), + False, + (T_DENSE,), + note="dequantize output head to dense BF16", + ), + ] + return entries + + +def activation_vperm_suffix(e: "Entry") -> str: + """Return the layer-independent checkpoint-stem suffix for C++ matching.""" + name = re.sub(r"^" + re.escape(PREFIX) + r"layers\.\d+\.", "", e.infinilm) + if name.endswith(".weight"): + name = name[: -len(".weight")] + return name + "." + + +def activation_vperm_rules(d: "Dims", plan: list) -> list: + """Derive runtime value-head activation permutations for quantization_config. + + llama.cpp exports selected output-projection columns in tiled order, while + the GDN kernel produces grouped activations. Packed columns cannot be moved + across blocks without requantization, so the runtime permutes activations. + """ + n_k, r, hd = d.lin_k_heads, d.v_per_k, d.lin_v_dim + rules, seen = [], set() + for e in plan: + if not e.act_vperm: + continue + in_dim = int(e.shape[1]) + if in_dim != n_k * r * hd: + raise ValueError( + "%s: input dimension %d != num_k_heads*num_v_per_k*head_dim %d; " + "cannot permute complete heads" % (e.infinilm, in_dim, n_k * r * hd) + ) + suffix = activation_vperm_suffix(e) + if suffix in seen: + continue + seen.add(suffix) + rules.append( + {"suffix": suffix, "num_k_heads": n_k, "num_v_per_k": r, "head_dim": hd} + ) + return rules + + +def expected_keys(d: Dims) -> list: + return [e.infinilm for e in build_plan(d)] + + +# GGUF tensor prefixes excluded from the main model, including the MTP block. +DROP_PREFIXES = ("blk.64.",) +MTP_BLOCK = 64 + + +def compress(shape: tuple) -> tuple: + """Remove singleton dimensions when comparing squeezed GGUF tensors.""" + return tuple(int(x) for x in shape if int(x) != 1) + + +# --------------------------------------------------------------------------- +# Derived checkpoint names, type-table keys, packed row sizes, and config data. +# --------------------------------------------------------------------------- +def ckpt_name(e: "Entry") -> str: + """Return the safetensors and framework state-dict parameter name.""" + if e.blob and e.infinilm.endswith(".weight"): + return e.infinilm[: -len(".weight")] + "." + BLOB_SUFFIX + return e.infinilm + + +def type_table_key(name: str) -> str: + """Return the exact checkpoint name used as the ggml_types table key.""" + return name + + +def row_bytes(n_in: int, block_size: int, type_size: int) -> int: + """Return bytes per packed row using caller-provided GGML block metadata.""" + if n_in % block_size: + raise ValueError( + "input size %d is not divisible by block size %d" % (n_in, block_size) + ) + return n_in // block_size * type_size + + +def make_text_config(d: "Dims") -> dict: + """Build the Qwen3.5 text_config consumed by InfiniLM.""" + return { + "model_type": "qwen3_5_text", + "hidden_size": d.hidden, + "num_hidden_layers": d.n_layers, + "num_attention_heads": d.n_q_heads, + "num_key_value_heads": d.n_kv_heads, + "head_dim": d.head_dim, + "intermediate_size": d.ffn, + "rms_norm_eps": d.rms_norm_eps, + "max_position_embeddings": d.max_position_embeddings, + "vocab_size": d.vocab, + "full_attention_interval": d.interval, + "linear_num_key_heads": d.lin_k_heads, + "linear_num_value_heads": d.lin_v_heads, + "linear_key_head_dim": d.lin_k_dim, + "linear_value_head_dim": d.lin_v_dim, + "linear_conv_kernel_dim": d.conv_kernel, + "attention_bias": False, + "rope_parameters": { + "rope_type": "mrope", + "rope_theta": d.rope_theta, + "partial_rotary_factor": d.partial_rotary_factor, + # InfiniLM requires three MRoPE sections. + "mrope_section": list(d.mrope_section), + # Qwen3.5 always uses interleaved MRoPE. + "mrope_interleaved": True, + }, + } + + +def make_root_config(d: "Dims", ggml_types: dict, act_vperm: list = None) -> dict: + """Build root config with top-level quantization metadata. + + ModelConfig reads quantization_config before merging text_config, so placing + it inside text_config would silently select NoneQuantization. + """ + return { + "model_type": "qwen3_5", + "torch_dtype": "bfloat16", + # BF16 logits collapse close top candidates into exact ties. Keep the + # dense BF16 head weights, but accumulate/write its output in FP32. + "lm_head_output_dtype": "float32", + "tie_word_embeddings": False, + "text_config": make_text_config(d), + "quantization_config": { + "quant_method": "gguf", + # Nested C++ modules remove this prefix before type-table lookup. + "key_prefix": PREFIX, + "ggml_types": ggml_types, + # Empty means that no runtime activation permutation is required. + "activation_vperm": act_vperm or [], + }, + } diff --git a/scripts/gguf_to_infinilm.py b/scripts/gguf_to_infinilm.py new file mode 100644 index 000000000..cb0a3d9af --- /dev/null +++ b/scripts/gguf_to_infinilm.py @@ -0,0 +1,886 @@ +#!/usr/bin/env python3 +"""Convert a GGUF model into an InfiniLM packed-weight checkpoint. + +All names, shapes, packed/dense choices, and transformations come from +``gguf_mapping``. Dequantization uses ``gguf.quants.dequantize``. Packed data +may only be moved as complete rows; bytes within quantization blocks remain +unchanged. GGUF tensors already match InfiniLM's [out, in] orientation. + +Example: + python3 scripts/gguf_to_infinilm.py --gguf MODEL.gguf --out OUT_DIR +""" + +from __future__ import annotations + +import argparse +import json +import os +import shutil +import sys +import time +from math import prod + +_HERE = os.path.dirname(os.path.abspath(__file__)) +sys.path.insert(0, _HERE) +if os.environ.get("LLAMA_CPP_DIR"): + sys.path.insert(0, os.path.join(os.environ["LLAMA_CPP_DIR"], "gguf-py")) + +import gguf_mapping as M # noqa: E402 +import gguf_transforms as X # noqa: E402 +import numpy as np # noqa: E402 +from gguf import GGUFReader # noqa: E402 +from gguf.constants import GGML_QUANT_SIZES # noqa: E402 +from gguf.constants import GGMLQuantizationType as Q # noqa: E402 +from gguf.quants import dequantize # noqa: E402 + +TYPE_NAME = {int(v.value): str(v.name) for v in Q} +TYPE_ID = {str(v.name): int(v.value) for v in Q} +UNQUANTIZED = ("F32", "F16", "BF16") + +# Tokenizer vocabulary comes from GGUF; copy auxiliary files when available. +TOKENIZER_FILES = ( + "tokenizer_config.json", + "chat_template.jinja", + "generation_config.json", + "preprocessor_config.json", + "video_preprocessor_config.json", + "special_tokens_map.json", + "merges.txt", + "vocab.json", + "tokenizer.json", +) + +_GiB = 2**30 + + +def log(msg: str) -> None: + print(msg, flush=True) + + +def blk_sizes(type_name: str) -> tuple[int, int]: + bs, ts = GGML_QUANT_SIZES[Q[type_name]] + return int(bs), int(ts) + + +# Normalize safetensors and torch dtype names before comparison. +_DTYPE_ALIAS = {"BF16": "bfloat16", "F16": "float16", "F32": "float32", "U8": "uint8"} + + +def norm_dtype(s) -> str: + s = str(s) + return _DTYPE_ALIAS.get(s.upper() if s.isupper() else s, s.lower()) + + +# --------------------------------------------------------------------------- +# Source-to-target conversion +# --------------------------------------------------------------------------- + + +def dense_float32(src: np.ndarray, type_name: str, chunk_rows: int) -> np.ndarray: + """Convert source rows to float32 [rows, in].""" + if type_name in UNQUANTIZED: + return np.asarray(src, dtype=np.float32) + if src.ndim != 2: + raise ValueError( + "quantized source must have shape [out, row_bytes], got %s" % (src.shape,) + ) + rows = src.shape[0] + if rows == 0: + return np.zeros((0,), dtype=np.float32) + q = Q[type_name] + first = np.asarray(dequantize(src[:chunk_rows], q), dtype=np.float32) + if rows <= chunk_rows: + return first + # Preallocate to bound peak memory for large tensors such as lm_head. + out = np.empty((rows,) + first.shape[1:], dtype=np.float32) + out[:chunk_rows] = first + for i in range(chunk_rows, rows, chunk_rows): + out[i : i + chunk_rows] = np.asarray( + dequantize(src[i : i + chunk_rows], q), dtype=np.float32 + ) + return out + + +def make_blob(e, t, dims, opt): + """Build a U8 [out, row_bytes] tensor, permuting only complete rows.""" + bs, ts = blk_sizes(opt.types[t.name]) + n_out, n_in = int(e.shape[0]), int(e.shape[1]) + rb = M.row_bytes(n_in, bs, ts) + if int(t.data.shape[-1]) != rb: + raise ValueError( + "%s: source row bytes %d != expected %d" + % (e.gguf, int(t.data.shape[-1]), rb) + ) + arr = t.data + if e.slices: + s, ep = e.slices[0] + arr = arr[s:ep] + if int(arr.shape[0]) != n_out: + raise ValueError( + "%s: sliced rows %d != mapping rows %d" % (e.gguf, arr.shape[0], n_out) + ) + if M.needs_vperm(e): + arr = X.apply_vperm(arr, e, dims, opt.vperm) + return torch_from(arr, np.uint8) + + +def entry_float32(e, t, dims, opt) -> np.ndarray: + """Return float32 values shared by dense output and verification.""" + tn = opt.types[t.name] + src = t.data + if e.slices: + s, ep = e.slices[0] + src = src[s:ep] + if tn in UNQUANTIZED: + arr = np.asarray(src, dtype=np.float32) + else: + arr = dense_float32(src, tn, opt.chunk_rows) + for tr in e.transforms: + if tr == M.T_ALOG: + arr = X.alog_from_ssm_a(arr) + elif tr in M.VPERM_TRANSFORMS: + arr = X.apply_vperm(arr, e, dims, opt.vperm) + elif tr in (M.T_DENSE, M.T_NONE): + continue + else: + raise ValueError("%s: unknown transform %r" % (e.infinilm, tr)) + want = tuple(int(x) for x in e.shape) + if tuple(arr.shape) != want: + if arr.size != prod(want): + raise ValueError( + "%s: transformed shape %s != mapped shape %s" + % (e.infinilm, arr.shape, want) + ) + arr = arr.reshape(want) # Restore a squeezed singleton convolution dimension. + return arr + + +def make_dense(e, t, dims, opt): + """Dequantize or cast through float32, then let torch produce BF16.""" + return torch_from(entry_float32(e, t, dims, opt), "bf16") + + +def torch_from(arr: np.ndarray, dtype): + import torch + + t = torch.from_numpy(np.ascontiguousarray(arr)) + return t.to(torch.bfloat16) if dtype == "bf16" else t + + +def _is_baked_plus1_norm(name: str) -> bool: + """Return whether llama.cpp baked a +1 normalization offset into GGUF.""" + return name.endswith("norm.weight") and not name.endswith("linear_attn.norm.weight") + + +def build(e, t, dims, opt, dense_all: bool): + """Build one output tensor and name; dense_all creates the dense reference.""" + e2 = e + if dense_all and e.blob: + e2 = _as_dense(e) + tens = make_blob(e2, t, dims, opt) if e2.blob else make_dense(e2, t, dims, opt) + # Dense-reference output can permute columns directly. Match the packed path, + # which performs the equivalent activation permutation at runtime. + if dense_all and e.act_vperm and opt.vperm != "none": + n_k, r, hd = dims.lin_k_heads, dims.v_per_k, dims.lin_v_dim + out_dim, in_dim = int(e.shape[0]), int(e.shape[1]) + if in_dim != n_k * r * hd: + raise ValueError( + "%s: input dimension %d != num_k_heads*num_v_per_k*head_dim %d; cannot permute complete heads" + % (e.infinilm, in_dim, n_k * r * hd) + ) + # [out, r, n_k, hd] tiled -> [out, n_k, r, hd] grouped -> flatten. + tens = ( + tens.view(out_dim, r, n_k, hd) + .transpose(1, 2) + .contiguous() + .view(out_dim, in_dim) + ) + # The dense reference loads through the non-GGUF remap, which adds +1 to + # selected norm weights. Store w-1 so loading reconstructs the baked GGUF w. + if dense_all and _is_baked_plus1_norm(e.infinilm): + tens = tens - 1 + name = M.ckpt_name(e2) + return name, tens + + +_DENSE_CACHE: dict = {} + + +def _as_dense(e): + """Return a dense-reference view of a packed entry without changing the plan.""" + key = (e.infinilm, e.vperm) + v = _DENSE_CACHE.get(key) + if v is None: + tr = tuple(x for x in e.transforms if x != M.T_NONE) + (M.T_DENSE,) + v = M.Entry( + e.infinilm, + e.gguf, + e.shape, + False, + tr, + e.types, + e.slices, + e.vperm, + "dense-ref " + (e.note or ""), + ) + _DENSE_CACHE[key] = v + return v + + +# --------------------------------------------------------------------------- +# Derive dimensions from GGUF metadata and validate the target profile. +# --------------------------------------------------------------------------- + + +def _dec(x) -> float: + """Render float32 metadata with a stable seven-significant-digit decimal.""" + return float("%.7g" % float(x)) + + +def dims_from_gguf(reader) -> M.Dims: + """Derive dimensions from standard llama.cpp GGUF metadata keys.""" + g = lambda suffix, idx=0: X.gguf_meta(reader, suffix)[idx] # noqa: E731 + n_layers = int(g("block_count")) - int(g("nextn_predict_layers")) + inner = int(g("ssm.inner_size")) + state = int(g("ssm.state_size")) + head_dim = int(g("attention.key_length")) + dim_cnt = int(g("rope.dimension_count")) + sec = [int(x) for x in X.gguf_meta(reader, "rope.dimension_sections")] + vocab = len(X.gguf_meta(reader, "tokenizer.ggml.tokens")) + return M.Dims( + hidden=int(g("embedding_length")), + n_q_heads=int(g("attention.head_count")), + n_kv_heads=int(g("attention.head_count_kv")), + head_dim=head_dim, + ffn=int(g("feed_forward_length")), + lin_k_heads=int(g("ssm.group_count")), + lin_v_heads=inner // state, + lin_k_dim=state, + lin_v_dim=state, + conv_kernel=int(g("ssm.conv_kernel")), + vocab=vocab, + n_layers=n_layers, + interval=int(g("full_attention_interval")), + mrope_section=tuple(sec[:3]), # InfiniLM consumes three MRoPE sections. + rope_theta=_dec(g("rope.freq_base")), + partial_rotary_factor=_dec(dim_cnt / head_dim), + rms_norm_eps=_dec(g("attention.layer_norm_rms_epsilon")), + max_position_embeddings=int(g("context_length")), + ) + + +def check_dims(d: M.Dims) -> None: + """Reject inputs whose metadata does not match the target model profile.""" + diff = [] + for f in _DIM_FIELDS: + got, want = getattr(d, f.name), getattr(M.REAL, f.name) + if isinstance(want, float) or isinstance(got, float): + if not np.isclose(float(got), float(want), rtol=1e-6, atol=1e-12): + diff.append("%s: %r != %r" % (f.name, got, want)) + elif got != want: + diff.append("%s: %r != %r" % (f.name, got, want)) + if diff: + raise SystemExit( + "GGUF metadata does not match gguf_mapping.REAL: %s\n" + "Create and validate a model-specific mapping before conversion." % diff + ) + log( + " rms_norm_eps: GGUF float32 %r -> config decimal %r" + % (float(d.rms_norm_eps), M.REAL.rms_norm_eps) + ) + + +from dataclasses import fields as _dc_fields # noqa: E402 + +_DIM_FIELDS = [f for f in _dc_fields(M.Dims) if f.name != "architectures"] + + +# --------------------------------------------------------------------------- +# Sharded output +# --------------------------------------------------------------------------- + + +class ShardWriter: + def __init__(self, out_dir: str, max_bytes: int): + self.dir, self.max = out_dir, max_bytes + self.buf: dict[str, object] = {} + self.buf_bytes = 0 + self.shards: list[str] = [] + self.weight_map: dict[str, str] = {} + self.total = 0 + + def add(self, name: str, tens) -> None: + nbytes = int(tens.numel()) * int(tens.element_size()) + if self.buf and self.buf_bytes + nbytes > self.max: + self.flush() + self.buf[name] = tens + self.buf_bytes += nbytes + self.total += nbytes + + def flush(self) -> None: + if not self.buf: + return + self.shards.append("__pending__") + idx = len(self.shards) + fname = "model-%05d.safetensors" % idx + from safetensors.torch import save_file + + save_file(self.buf, os.path.join(self.dir, fname), metadata={"format": "pt"}) + for k in self.buf: + self.weight_map[k] = fname + log( + " wrote %s (%.2f GiB, %d tensors)" + % (fname, self.buf_bytes / _GiB, len(self.buf)) + ) + self.shards[-1] = fname + self.buf, self.buf_bytes = {}, 0 + + def finish(self) -> None: + self.flush() + n = len(self.shards) + renamed = {} + for i, f in enumerate(self.shards, 1): + new = "model-%05d-of-%05d.safetensors" % (i, n) + if f != new: + os.rename(os.path.join(self.dir, f), os.path.join(self.dir, new)) + renamed[f] = new + self.weight_map = {k: renamed[v] for k, v in self.weight_map.items()} + with open(os.path.join(self.dir, "model.safetensors.index.json"), "w") as fp: + json.dump( + {"metadata": {"total_size": self.total}, "weight_map": self.weight_map}, + fp, + indent=1, + sort_keys=True, + ) + log(" %d shards, %.3f GiB total" % (n, self.total / _GiB)) + + +# --------------------------------------------------------------------------- +# Verification +# --------------------------------------------------------------------------- + + +def rows_hash(a) -> str: + """Hash the multiset of rows in a [rows, cols] byte array. + + A row permutation cannot be compared positionally, but it must preserve the + complete multiset of packed rows. + """ + import hashlib + + a = np.ascontiguousarray(a) + v = a.view(np.void(a.shape[1] * a.dtype.itemsize)).ravel() + h = hashlib.sha256() + for x in np.sort(v): + h.update(x.tobytes()) + return h.hexdigest()[:16] + + +def dense_bits_check(e, t, dims, opt, prod_t) -> bool: + """Verify dense BF16 entries bitwise in bounded row chunks.""" + import torch + + if M.needs_vperm(e): + return bool( + np.array_equal( + prod_t.view(torch.uint16).numpy(), + X.bf16_bits(entry_float32(e, t, dims, opt)), + ) + ) + src = np.asarray(t.data) + if e.slices: + s, ep = e.slices[0] + src = src[s:ep] + tn = opt.types[t.name] + tail = tuple(int(x) for x in e.shape[1:]) + per_row = prod(tail) if tail else 1 + rows = max(1, int(_BIG_ELEMS // per_row)) + n = int(e.shape[0]) + if src.shape[0] != n: + return False + for i in range(0, n, rows): + blk = src[i : i + rows] + arr = ( + np.asarray(blk, dtype=np.float32) + if tn in UNQUANTIZED + else dense_float32(blk, tn, opt.chunk_rows) + ) + exp = X.bf16_bits(arr.reshape((blk.shape[0],) + tail)) + got = prod_t[i : i + rows].view(torch.uint16).numpy() + if not np.array_equal(got, exp): + return False + return True + + +_BIG_ELEMS = 64 * 1024 * 1024 # At most 64M float32 elements per verification chunk. + + +def verify(out_dir: str, plan, tensors, dims, opt, sample) -> int: + """Reload output, validate metadata, and sample bytes. Return failure count.""" + import torch + from safetensors import safe_open + + log("\n== Verification: reload output ==") + bs_files = sorted( + f + for f in os.listdir(out_dir) + if f.endswith(".safetensors") and not f.startswith(".") + ) + with open(os.path.join(out_dir, "model.safetensors.index.json")) as fp: + index = json.load(fp) + got: dict[str, tuple] = {} + handles = {} + for f in bs_files: + h = safe_open(os.path.join(out_dir, f), framework="pt") + handles[f] = h + for k in h.keys(): + t = h.get_slice(k) + got[k] = (tuple(int(x) for x in t.get_shape()), norm_dtype(t.get_dtype())) + fails = 0 + want = {} + for e in plan: + name = M.ckpt_name(e) + if e.blob: + bs, ts = blk_sizes(opt.types[e.gguf]) + shape, dt = (int(e.shape[0]), M.row_bytes(int(e.shape[1]), bs, ts)), "uint8" + else: + shape, dt = tuple(int(x) for x in e.shape), "bfloat16" + want[name] = (shape, dt, e) + missing = sorted(set(want) - set(got)) + extra = sorted(set(got) - set(want)) + for label, keys in (("missing keys", missing), ("extra keys", extra)): + if keys: + fails += 1 + log(" FAIL %s (%d): %s" % (label, len(keys), keys[:6])) + else: + log(" PASS no %s" % label) + bad = [k for k in set(want) & set(got) if want[k][:2] != got[k]] + if bad: + fails += 1 + log( + " FAIL shape/dtype mismatch (%d): %s" + % (len(bad), [(k, want[k][:2], got[k]) for k in sorted(bad)[:4]]) + ) + else: + log(" PASS shape and dtype for all %d keys" % len(want)) + + # Type-table keys must exactly match output tensor names in both directions. + with open(os.path.join(out_dir, "config.json")) as fp: + cfg = json.load(fp) + qcfg = cfg.get("quantization_config") or {} + table = qcfg.get("ggml_types") or {} + if qcfg.get("quant_method") != "gguf": + fails += 1 + log(" FAIL top-level quantization_config.quant_method is not 'gguf'") + elif qcfg.get("key_prefix") != M.PREFIX: + fails += 1 + log(" FAIL config.json is missing key_prefix=%r" % M.PREFIX) + else: + log(" PASS top-level quantization_config with key_prefix=%r" % M.PREFIX) + for label, keys in ( + ("type-table missing keys", sorted(set(got) - set(table))), + ("type-table extra keys", sorted(set(table) - set(got))), + ): + if keys: + fails += 1 + log(" FAIL %s (%d): %s" % (label, len(keys), keys[:6])) + else: + log(" PASS no %s (%d exact tensor-name matches)" % (label, len(table))) + + # Deterministically sample every conversion category rather than relying on + # random samples that may miss permutations and slices. + def sel(pred): + return sorted(k for k, (_, _, e) in want.items() if pred(e)) + + cats = [ + ( + "packed unchanged", + sel(lambda e: e.blob and not e.slices and not M.needs_vperm(e)), + ), + ( + "packed V permutation", + sel(lambda e: e.blob and not e.slices and M.needs_vperm(e)), + ), + ("packed fused slice", sel(lambda e: e.blob and e.slices)), + ( + "BF16 dequantization", + sel(lambda e: not e.blob and not e.slices and not M.needs_vperm(e)), + ), + ("BF16 permutation and A_log", sel(lambda e: not e.blob and M.needs_vperm(e))), + ("BF16 fused slice", sel(lambda e: not e.blob and e.slices)), + ] + picks = [c[1][0] for c in cats if c[1]] + if sample == "all": + picks = sorted(k for k, (_, _, e) in want.items() if e.blob) + log( + " sampled %d entries: %s" + % ( + len(picks), + "all packed entries" + if sample == "all" + else " ".join("%s=%s" % (c, len(v)) for c, v in cats), + ) + ) + for k in picks: + shape, dt, e = want[k] + prod_t = handles[index["weight_map"][k]].get_tensor(k) + src = np.asarray(tensors[e.gguf].data) + if e.slices: + s, ep = e.slices[0] + src = src[s:ep] + ref = build(e, tensors[e.gguf], dims, opt, False)[1] + checks = [] + if tuple(int(x) for x in prod_t.shape) != tuple(int(x) for x in ref.shape): + checks.append(("shape", False)) + elif e.blob: + p = prod_t.numpy() + checks.append(("matches reconstruction", bool(torch.equal(prod_t, ref)))) + if M.needs_vperm(e): + checks.append( + ("same source-row multiset", rows_hash(p) == rows_hash(src)) + ) + else: + checks.append(("byte-identical to GGUF source", np.array_equal(p, src))) + else: + # Independently compare BF16 bits against NumPy RNE conversion. + checks.append( + ( + "BF16 bits match NumPy RNE", + dense_bits_check(e, tensors[e.gguf], dims, opt, prod_t), + ) + ) + ok = all(v for _, v in checks) + fails += 0 if ok else 1 + log( + " %s %-52s %-16s %s" + % ( + "PASS" if ok else "FAIL", + k, + str(tuple(int(x) for x in prod_t.shape)), + ", ".join("%s=%s" % (n, "Y" if v else "N") for n, v in checks), + ) + ) + return fails + + +# --------------------------------------------------------------------------- +# main +# --------------------------------------------------------------------------- + + +def main() -> int: + ap = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + ap.add_argument("--gguf", required=True, help="source GGUF file") + ap.add_argument("--out", required=True, help="output checkpoint directory") + ap.add_argument( + "--tokenizer-dir", + default="", + help="optional directory with auxiliary tokenizer files", + ) + ap.add_argument( + "--dense-iq", + action=argparse.BooleanOptionalAction, + default=True, + help="convert IQ4_NL/IQ4_XS tensors to dense BF16", + ) + ap.add_argument( + "--dense-embed", + action=argparse.BooleanOptionalAction, + default=True, + help="convert embedding and output head to dense BF16", + ) + ap.add_argument( + "--vperm", + choices=("inv", "fwd", "none"), + default="inv", + help="value-head permutation direction", + ) + ap.add_argument( + "--emit-dense-ref", + metavar="PATH", + default=None, + help="also emit a fully dequantized BF16 reference checkpoint", + ) + ap.add_argument("--max-shard-gib", type=float, default=4.0) + ap.add_argument( + "--chunk-rows", type=int, default=8192, help="rows per dequantization chunk" + ) + ap.add_argument( + "--layers", + type=int, + default=None, + help="convert only the first N layers and update num_hidden_layers", + ) + ap.add_argument("--verify", choices=("off", "sample", "all"), default="sample") + ap.add_argument( + "--skip-pack", + action="store_true", + help="reuse existing weights while refreshing config, tokenizer, and verification", + ) + ap.add_argument( + "--dry-run", + action="store_true", + help="validate orientation, shapes, and packed byte sizes without writing output", + ) + a = ap.parse_args() + + if not a.dense_embed: + raise SystemExit( + "--no-dense-embed requires native packed embedding and lm_head kernels" + ) + + t0 = time.time() + log("Reading GGUF metadata: %s" % a.gguf) + reader = GGUFReader(a.gguf) + tensors = {t.name: t for t in reader.tensors} + dims = dims_from_gguf(reader) + check_dims(dims) + log( + " dimensions match mapping profile: %d layers, hidden=%d, vocab=%d" + % (dims.n_layers, dims.hidden, dims.vocab) + ) + + # Validate the full model profile before applying the optional layer limit, + # then keep the emitted config consistent with the truncated checkpoint. + if a.layers is not None: + if not 0 < a.layers < dims.n_layers: + raise SystemExit( + "--layers must be in (0, %d), got %d" % (dims.n_layers, a.layers) + ) + log( + " --layers %d: num_hidden_layers %d -> %d" + % (a.layers, dims.n_layers, a.layers) + ) + dims.n_layers = a.layers + + opt = type("Opt", (), {})() + opt.vperm, opt.chunk_rows = a.vperm, a.chunk_rows + opt.types = {n: TYPE_NAME[int(t.tensor_type)] for n, t in tensors.items()} + plan = M.build_plan(dims) + n_exc = M.apply_v1_exceptions(plan, opt.types, enabled=a.dense_iq) + log(" mapping entries %d, dense IQ4 fallbacks %d" % (len(plan), n_exc)) + blob = [e for e in plan if e.blob] + log( + " packed %d / dense %d / excluded MTP prefixes %s" + % (len(blob), len(plan) - len(blob), M.DROP_PREFIXES) + ) + + if a.dry_run: + log("\n== Dry run: validate orientation and byte sizes ==") + blob_bytes = dense_bytes = 0 + for e in plan: + t = tensors[e.gguf] + n_out = (e.slices[0][1] - e.slices[0][0]) if e.slices else int(e.shape[0]) + if e.blob: + bs, ts = blk_sizes(opt.types[e.gguf]) + rb = M.row_bytes(int(e.shape[1]), bs, ts) + if int(t.data.shape[-1]) != rb: + raise ValueError( + "%s: source row bytes %d != expected %d" + % (e.gguf, int(t.data.shape[-1]), rb) + ) + if int(t.data.shape[0]) < n_out: + raise ValueError( + "%s: source has %d rows, entry requires %d" + % (e.gguf, t.data.shape[0], n_out) + ) + blob_bytes += n_out * rb + else: + n = prod(tuple(int(x) for x in e.shape)) + if not e.slices and prod(int(x) for x in t.shape) != n: + raise ValueError( + "%s: source element count %s != entry shape %s" + % (e.gguf, t.shape, tuple(e.shape)) + ) + dense_bytes += n * 2 + log( + " PASS %d entries: packed %.3f GiB + dense BF16 %.3f GiB" + " = expected output %.3f GiB" + % ( + len(plan), + blob_bytes / _GiB, + dense_bytes / _GiB, + (blob_bytes + dense_bytes) / _GiB, + ) + ) + return 0 + + os.makedirs(a.out, exist_ok=True) + w = ShardWriter(a.out, int(a.max_shard_gib * _GiB)) + ggml_types = {} + for e in plan: + ggml_types[M.type_table_key(M.ckpt_name(e))] = ( + TYPE_ID[opt.types[e.gguf]] if e.blob else "dense_bf16" + ) + if a.skip_pack: + with open(os.path.join(a.out, "model.safetensors.index.json")) as fp: + w.total = json.load(fp)["metadata"]["total_size"] + log("\n== --skip-pack: refresh config, tokenizer, and verification ==") + else: + log("\n== Writing %s ==" % a.out) + for i, e in enumerate(plan): + t = tensors[e.gguf] + name, tens = build(e, t, dims, opt, False) + w.add(name, tens) + if (i + 1) % 100 == 0: + log( + " ... %d/%d entries (%.1f s)" + % (i + 1, len(plan), time.time() - t0) + ) + w.finish() + + # Refresh semantic metadata even with --skip-pack. Derive rules from the + # mapping so Python and C++ do not maintain duplicate definitions. + rules = M.activation_vperm_rules(dims, plan) + if a.vperm == "none": + # Keep conversion, runtime, and dense-reference paths aligned. + rules = [] + cfg = M.make_root_config(dims, ggml_types, rules) + with open(os.path.join(a.out, "config.json"), "w") as fp: + json.dump(cfg, fp, indent=1, sort_keys=True) + log( + " config.json: %d ggml_types keys and %d activation V-head rules: %s" + % ( + len(ggml_types), + len(rules), + " ".join( + "%s=%dx%dx%d" + % (r["suffix"], r["num_k_heads"], r["num_v_per_k"], r["head_dim"]) + for r in rules + ) + or "none", + ) + ) + + fails = export_tokenizer(reader, a.out, a.tokenizer_dir, dims) + + if not a.skip_pack: + with open(os.path.join(a.out, "pack_report.json"), "w") as fp: + json.dump( + { + "gguf": os.path.abspath(a.gguf), + # Tensor payload excludes file metadata and alignment padding. + "gguf_file_bytes": os.path.getsize(os.path.abspath(a.gguf)), + "gguf_tensor_data_bytes": sum( + int(t.n_bytes) for t in reader.tensors + ), + "n_gguf_tensors": len(tensors), + "v1_dense_iq": bool(a.dense_iq), + "vperm": a.vperm, + "n_entries": len(plan), + "n_blob": len(blob), + "n_v1_exceptions": n_exc, + "blob_type_ids": sorted({TYPE_ID[opt.types[e.gguf]] for e in blob}), + "out_bytes": w.total, + "shards": w.shards, + "seconds": round(time.time() - t0, 1), + }, + fp, + indent=1, + sort_keys=True, + ) + + if a.verify != "off": + fails += verify(a.out, plan, tensors, dims, opt, a.verify) + + if a.emit_dense_ref: + log("\n== Writing dense reference %s ==" % a.emit_dense_ref) + os.makedirs(a.emit_dense_ref, exist_ok=True) + wr = ShardWriter(a.emit_dense_ref, int(a.max_shard_gib * _GiB)) + for e in plan: + name, tens = build(e, tensors[e.gguf], dims, opt, True) + wr.add(name, tens) + wr.finish() + ref_cfg = M.make_root_config( + dims, {M.type_table_key(k): "dense_bf16" for k in ggml_types}, rules + ) + # The dense reference omits quantization_config and stores grouped + # columns directly, matching the packed path's runtime semantics. + del ref_cfg["quantization_config"] + with open(os.path.join(a.emit_dense_ref, "config.json"), "w") as fp: + json.dump(ref_cfg, fp, indent=1, sort_keys=True) + export_tokenizer(reader, a.emit_dense_ref, a.tokenizer_dir, dims) + + log( + "\n===== Complete: %.1f s, output %.3f GiB, verification failures %d =====" + % (time.time() - t0, w.total / _GiB, fails) + ) + return 1 if fails else 0 + + +def export_tokenizer(reader, out_dir: str, tokenizer_dir: str, dims) -> int: + """Export GGUF BPE vocabulary and copy optional tokenizer configuration.""" + tokens = [str(t) for t in X.gguf_meta(reader, "tokenizer.ggml.tokens")] + merges = [str(m) for m in X.gguf_meta(reader, "tokenizer.ggml.merges")] + model = str(X.gguf_meta(reader, "tokenizer.ggml.model")[0]) + if len(tokens) != dims.vocab: + raise SystemExit( + "GGUF vocabulary size %d != config vocab_size %d" + % (len(tokens), dims.vocab) + ) + if model != "gpt2": + log( + " WARNING tokenizer.ggml.model=%r is not gpt2; verify vocab/merges export" + % model + ) + with open(os.path.join(out_dir, "vocab.json"), "w", encoding="utf-8") as fp: + json.dump({t: i for i, t in enumerate(tokens)}, fp, ensure_ascii=False) + with open(os.path.join(out_dir, "merges.txt"), "w", encoding="utf-8") as fp: + fp.write("#version: 0.2\n" + "\n".join(merges) + "\n") + log( + " tokenizer from GGUF: vocab %d / merges %d (model=%s)" + % (len(tokens), len(merges), model) + ) + + have = ( + os.listdir(tokenizer_dir) + if tokenizer_dir and os.path.isdir(tokenizer_dir) + else [] + ) + if not have: + log(" WARNING tokenizer config directory not found: %s" % tokenizer_dir) + copied = [] + for f in TOKENIZER_FILES: + dst = os.path.join(out_dir, f) + if f in have and not os.path.exists(dst): + shutil.copy2(os.path.join(tokenizer_dir, f), dst) + copied.append(f) + log( + " copied %d auxiliary tokenizer files: %s" + % (len(copied), " ".join(sorted(copied))) + ) + if "tokenizer_config.json" not in copied + have: + raise SystemExit( + "output is missing tokenizer_config.json; it was not available in %s" + % tokenizer_dir + ) + return check_tokenizer(out_dir, dims) + + +def check_tokenizer(out_dir: str, dims) -> int: + """Load AutoTokenizer and verify an encode/decode round trip.""" + try: + from transformers import AutoTokenizer + except ImportError: + log(" SKIP tokenizer verification: transformers is unavailable") + return 0 + try: + tk = AutoTokenizer.from_pretrained(out_dir) + n, cls = len(tk), type(tk).__name__ + s = "Hello, world 27B" + ids = tk.encode(s) + ok = n == dims.vocab and tk.decode(ids) == s + log( + " %s tokenizer %s vocab=%d round_trip=%s" + % ("PASS" if ok else "FAIL", cls, n, tk.decode(ids) == s) + ) + return 0 if ok else 1 + except Exception as exc: # noqa: BLE001 + log(" FAIL tokenizer load: %s: %s" % (type(exc).__name__, str(exc)[:200])) + return 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/gguf_transforms.py b/scripts/gguf_transforms.py new file mode 100644 index 000000000..7be63e141 --- /dev/null +++ b/scripts/gguf_transforms.py @@ -0,0 +1,123 @@ +#!/usr/bin/env python3 +"""Pure NumPy transforms shared by the GGUF-to-InfiniLM converter. + +Value-head ordering follows llama.cpp's Qwen conversion: + HF / InfiniLM = grouped [key][value] + GGUF = tiled [value][key] + +Therefore ``reorder_v`` converts grouped to tiled order and +``reorder_v_inverse`` converts tiled to grouped order. +""" + +from __future__ import annotations + +import numpy as np + +# --------------------------------------------------------------------------- +# Value-head permutation +# --------------------------------------------------------------------------- + + +def reorder_v(t: np.ndarray, n_k: int, n_v_per_k: int, hd: int) -> np.ndarray: + """Convert grouped to tiled order along dimension 0. + + Trailing dimensions are preserved, including 1-D scalars per head, + 2-D weight rows, and 3-D convolution weights. + """ + rest = t.shape[1:] + return ( + t.reshape((n_k, n_v_per_k, hd) + rest) + .transpose((1, 0, 2) + tuple(range(3, 3 + len(rest)))) + .reshape((n_k * n_v_per_k * hd,) + rest) + ) + + +def reorder_v_inverse(t: np.ndarray, n_k: int, n_v_per_k: int, hd: int) -> np.ndarray: + """Convert tiled to grouped order along dimension 0.""" + rest = t.shape[1:] + return ( + t.reshape((n_v_per_k, n_k, hd) + rest) + .transpose((1, 0, 2) + tuple(range(3, 3 + len(rest)))) + .reshape((n_k * n_v_per_k * hd,) + rest) + ) + + +_VPERM = {"inv": reorder_v_inverse, "fwd": reorder_v, "none": None} + + +def vperm_head_dim(e, dims) -> int: + """Return the number of elements per value head for a mapping entry. + + Derive the value from shape rather than tensor names so the converter does + not need a second name table. + """ + n_heads = dims.lin_v_heads + rows = int(e.shape[0]) if e.vperm == "all" else dims.value_dim + if rows % n_heads: + raise ValueError( + "%s: scope rows %d are not divisible by %d value heads" + % (e.infinilm, rows, n_heads) + ) + return rows // n_heads + + +def apply_vperm(arr: np.ndarray, e, dims, direction: str = "inv") -> np.ndarray: + """Permute value heads along dimension 0 within an all or v_tail scope.""" + fn = _VPERM[direction] + if fn is None: + return arr + n_k, hd = dims.lin_k_heads, vperm_head_dim(e, dims) + v_per_k = dims.lin_v_heads // n_k + if int(dims.lin_v_heads) % n_k: + raise ValueError( + "lin_v_heads %d is not divisible by lin_k_heads %d" + % (dims.lin_v_heads, n_k) + ) + if e.vperm == "v_tail": + n_v = n_k * v_per_k * hd + if arr.shape[0] < n_v: + raise ValueError( + "%s: dimension 0 size %d is smaller than value segment %d" + % (e.infinilm, arr.shape[0], n_v) + ) + out = np.asarray(arr, dtype=arr.dtype) + return np.concatenate([out[:-n_v], fn(out[-n_v:], n_k, v_per_k, hd)], axis=0) + return fn(np.asarray(arr, dtype=arr.dtype), n_k, v_per_k, hd) + + +# --------------------------------------------------------------------------- +# Other transforms +# --------------------------------------------------------------------------- + + +def alog_from_ssm_a(a: np.ndarray) -> np.ndarray: + """Recover the HF ``A_log`` convention from GGUF ``-exp(A_log)`` values.""" + a = np.asarray(a, dtype=np.float32) + if not np.all(a < 0): + raise ValueError( + "ssm_a contains a non-negative value (min=%g); cannot compute log(-x). " + "Check the A_log convention in conversion/qwen.py." % float(a.min()) + ) + return np.log(-a) + + +def gguf_meta(reader, suffix: str): + """Read metadata with architecture/general prefixes and return a list.""" + for key in ("qwen35.%s" % suffix, "general.%s" % suffix, suffix): + if key in reader.fields: + v = reader.fields[key].contents() + return v if isinstance(v, (list, tuple, np.ndarray)) else [v] + raise KeyError( + "missing GGUF metadata %s (no qwen35/general/unprefixed match)" % suffix + ) + + +def bf16_bits(x: np.ndarray) -> np.ndarray: + """Return round-to-nearest-even bfloat16 bit patterns as uint16. + + Keep arithmetic in uint32 to avoid doubling memory for large tensors. An + overflow beyond bit 31 cannot affect the retained bfloat16 bits. + """ + u = np.ascontiguousarray(x, dtype=np.float32).view(np.uint32) + bias = ((u >> np.uint32(16)) & np.uint32(1)) + np.uint32(0x7FFF) + return ((u + bias) >> np.uint32(16)).astype(np.uint16) diff --git a/test/bench/backends/infinilm.py b/test/bench/backends/infinilm.py index fbd99379d..7c10291ca 100644 --- a/test/bench/backends/infinilm.py +++ b/test/bench/backends/infinilm.py @@ -50,6 +50,9 @@ def __init__( print(f"Graph compilation: {'enabled' if enable_graph else 'disabled'}") print(f"Attention backend: {attn_backend}") + model_type = self.config_dict.get("model_type") + enable_prefix_caching = model_type not in {"qwen3_5", "qwen3_5_moe"} + self.model = LLM( model_path=model_dir_path, device=device_name, @@ -60,6 +63,7 @@ def __init__( block_size=256, enable_graph=enable_graph, attn_backend=attn_backend, + enable_prefix_caching=enable_prefix_caching, ) self.processor = self.model.engine.processor self.tokenizer = self.processor.get_tokenizer() diff --git a/test/scripts/test_gguf_routeb.py b/test/scripts/test_gguf_routeb.py new file mode 100644 index 000000000..fa0add8f3 --- /dev/null +++ b/test/scripts/test_gguf_routeb.py @@ -0,0 +1,75 @@ +import sys +import unittest +from pathlib import Path +from types import SimpleNamespace + +import numpy as np + +ROOT = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(ROOT / "scripts")) + +import gguf_mapping as mapping # noqa: E402 +import gguf_transforms as transforms # noqa: E402 + + +class GGUFTransformsTest(unittest.TestCase): + def test_value_head_permutation_round_trip(self): + source = np.arange(2 * 3 * 4 * 5, dtype=np.uint8).reshape(24, 5) + tiled = transforms.reorder_v(source, n_k=2, n_v_per_k=3, hd=4) + restored = transforms.reorder_v_inverse(tiled, n_k=2, n_v_per_k=3, hd=4) + + np.testing.assert_array_equal(restored, source) + self.assertEqual(sorted(map(bytes, tiled)), sorted(map(bytes, source))) + + def test_tail_permutation_preserves_prefix_and_rows(self): + dims = SimpleNamespace(lin_k_heads=2, lin_v_heads=6, value_dim=12) + entry = SimpleNamespace(shape=(20, 3), vperm="v_tail", infinilm="conv") + source = np.arange(60, dtype=np.uint8).reshape(20, 3) + + tiled = transforms.apply_vperm(source, entry, dims, direction="fwd") + restored = transforms.apply_vperm(tiled, entry, dims, direction="inv") + + np.testing.assert_array_equal(tiled[:8], source[:8]) + np.testing.assert_array_equal(restored, source) + self.assertEqual(sorted(map(bytes, tiled[8:])), sorted(map(bytes, source[8:]))) + + def test_bf16_bits_for_exact_values(self): + values = np.array([0.0, 1.0, -2.0, np.inf], dtype=np.float32) + expected = np.array([0x0000, 0x3F80, 0xC000, 0x7F80], dtype=np.uint16) + np.testing.assert_array_equal(transforms.bf16_bits(values), expected) + + +class GGUFMappingTest(unittest.TestCase): + def test_generated_config_keeps_quantization_at_root(self): + table = { + "model.language_model.layers.0.mlp.down_proj.weight_bytes": mapping.Q6_K + } + rules = mapping.activation_vperm_rules( + mapping.REAL, mapping.build_plan(mapping.REAL) + ) + config = mapping.make_root_config(mapping.REAL, table, rules) + + self.assertNotIn("quantization_config", config["text_config"]) + quant = config["quantization_config"] + self.assertEqual(quant["quant_method"], "gguf") + self.assertEqual(quant["key_prefix"], mapping.PREFIX) + self.assertEqual(quant["ggml_types"], table) + self.assertEqual(quant["activation_vperm"], rules) + self.assertTrue(rules) + self.assertEqual(len({rule["suffix"] for rule in rules}), len(rules)) + + def test_packed_checkpoint_name_and_row_size(self): + entry = SimpleNamespace( + blob=True, infinilm="model.language_model.layers.0.mlp.down_proj.weight" + ) + self.assertEqual( + mapping.ckpt_name(entry), + "model.language_model.layers.0.mlp.down_proj.weight_bytes", + ) + self.assertEqual(mapping.row_bytes(5120, block_size=256, type_size=210), 4200) + with self.assertRaises(ValueError): + mapping.row_bytes(5119, block_size=256, type_size=210) + + +if __name__ == "__main__": + unittest.main()