diff --git a/conversion/__init__.py b/conversion/__init__.py index 7936f11..9b2107b 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -138,6 +138,7 @@ TEXT_MODEL_MAP: dict[str, str] = { "Llama4ForConditionalGeneration": "llama", "LlamaBidirectionalModel": "llama", "LlamaForCausalLM": "llama", + "NanbeigeForCausalLM": "llama", "LlamaModel": "llama", "Eagle3DraftModel": "llama", "Eagle3Speculator": "llama", diff --git a/conversion/llama.py b/conversion/llama.py index 315a619..7702755 100644 --- a/conversion/llama.py +++ b/conversion/llama.py @@ -344,6 +344,24 @@ class LlamaModel(TextModel): raise ValueError(f"Unprocessed experts: {experts}") +@ModelBase.register("NanbeigeForCausalLM") +class NanbeigeModel(LlamaModel): + model_arch = gguf.MODEL_ARCH.NANBEIGE + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.physical_block_count = self.block_count + self.loop_count = int(self.hparams.get("num_loops", 1)) + if self.loop_count < 1: + raise ValueError("num_loops must be at least 1") + self.block_count *= self.loop_count + self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count) + + def set_gguf_parameters(self): + super().set_gguf_parameters() + self.gguf_writer.add_uint32("nanbeige.loop_count", self.loop_count) + + @ModelBase.register("ArceeForCausalLM") class ArceeModel(LlamaModel): model_arch = gguf.MODEL_ARCH.ARCEE diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 4391bd1..5fa2a33 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -407,6 +407,7 @@ class GGUFType: class MODEL_ARCH(IntEnum): MMPROJ = auto() # dummy arch for clip.cpp LLAMA = auto() + NANBEIGE = auto() LLAMA4 = auto() DECI = auto() FALCON = auto() @@ -990,6 +991,7 @@ class MODEL_TENSOR(IntEnum): MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = { MODEL_ARCH.MMPROJ: "clip", # dummy arch for clip.cpp MODEL_ARCH.LLAMA: "llama", + MODEL_ARCH.NANBEIGE: "nanbeige", MODEL_ARCH.LLAMA4: "llama4", MODEL_ARCH.DECI: "deci", MODEL_ARCH.FALCON: "falcon", @@ -1780,6 +1782,26 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.FFN_DOWN_EXP, MODEL_TENSOR.FFN_UP_EXP, ], + MODEL_ARCH.NANBEIGE: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_ROT_EMBD, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + ], MODEL_ARCH.LLAMA4: [ MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT_NORM, @@ -4469,6 +4491,10 @@ MODEL_TENSOR_SKIP: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_ROT_EMBD, ], + MODEL_ARCH.NANBEIGE: [ + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_ROT_EMBD, + ], MODEL_ARCH.DECI: [ MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_ROT_EMBD, diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index e5a9a29..4f345cc 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -8,6 +8,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_CLIP, "clip" }, // dummy, only used by llama-quantize { LLM_ARCH_LLAMA, "llama" }, + { LLM_ARCH_NANBEIGE, "nanbeige" }, { LLM_ARCH_LLAMA4, "llama4" }, { LLM_ARCH_DECI, "deci" }, { LLM_ARCH_FALCON, "falcon" }, @@ -179,6 +180,7 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_EMBEDDING_LENGTH_PER_LAYER, "%s.embedding_length_per_layer_input" }, { LLM_KV_FEATURES_LENGTH, "%s.features_length" }, { LLM_KV_BLOCK_COUNT, "%s.block_count" }, + { LLM_KV_LOOP_COUNT, "%s.loop_count" }, { LLM_KV_LEADING_DENSE_BLOCK_COUNT, "%s.leading_dense_block_count" }, { LLM_KV_FEED_FORWARD_LENGTH, "%s.feed_forward_length" }, { LLM_KV_EXPERT_FEED_FORWARD_LENGTH, "%s.expert_feed_forward_length" }, diff --git a/src/llama-arch.h b/src/llama-arch.h index 0ef95b0..e1b39b3 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -13,6 +13,7 @@ enum llm_arch { LLM_ARCH_CLIP, LLM_ARCH_LLAMA, + LLM_ARCH_NANBEIGE, LLM_ARCH_LLAMA4, LLM_ARCH_DECI, LLM_ARCH_FALCON, @@ -184,6 +185,7 @@ enum llm_kv { LLM_KV_EMBEDDING_LENGTH_PER_LAYER, LLM_KV_FEATURES_LENGTH, LLM_KV_BLOCK_COUNT, + LLM_KV_LOOP_COUNT, LLM_KV_LEADING_DENSE_BLOCK_COUNT, LLM_KV_FEED_FORWARD_LENGTH, LLM_KV_EXPERT_FEED_FORWARD_LENGTH, diff --git a/src/llama-hparams.h b/src/llama-hparams.h index 8be5f28..1ce3e43 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -51,6 +51,7 @@ struct llama_hparams { uint32_t n_ctx_train; // context size the model was trained on uint32_t n_embd; uint32_t n_layer_all; + uint32_t n_loop = 1; uint32_t n_layer_nextn = 0; uint32_t n_expert = 0; uint32_t n_expert_used = 0; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 7e5bab2..b622931 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -40,6 +40,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params switch (arch) { case LLM_ARCH_LLAMA: return new llama_model_llama(params); + case LLM_ARCH_NANBEIGE: + return new llama_model_nanbeige(params); case LLM_ARCH_LLAMA4: return new llama_model_llama4(params); case LLM_ARCH_LLAMA_EMBED: @@ -1191,7 +1193,7 @@ void llama_model_base::load_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_ROPE_DIMENSION_COUNT, hparams.n_rot_full, false); - if (arch == LLM_ARCH_LLAMA || arch == LLM_ARCH_DECI || arch == LLM_ARCH_FALCON || arch == LLM_ARCH_LLAMA_EMBED) { + if (arch == LLM_ARCH_LLAMA || arch == LLM_ARCH_NANBEIGE || arch == LLM_ARCH_DECI || arch == LLM_ARCH_FALCON || arch == LLM_ARCH_LLAMA_EMBED) { if (hparams.n_rot_full != hparams.n_embd_head_k_full) { throw std::runtime_error(format("invalid n_rot: %u, expected %u", hparams.n_rot_full, hparams.n_embd_head_k_full)); } @@ -2450,6 +2452,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { // use what we call a normal RoPE, operating on pairs of consecutive head values case LLM_ARCH_LLAMA: + case LLM_ARCH_NANBEIGE: case LLM_ARCH_LLADA: case LLM_ARCH_LLAMA4: case LLM_ARCH_DECI: diff --git a/src/models/llama.cpp b/src/models/llama.cpp index 4bfebc8..3c31ac0 100644 --- a/src/models/llama.cpp +++ b/src/models/llama.cpp @@ -31,6 +31,15 @@ void llama_model_llama::load_arch_hparams(llama_model_loader & ml) { } } +void llama_model_nanbeige::load_arch_hparams(llama_model_loader & ml) { + llama_model_llama::load_arch_hparams(ml); + ml.get_key(LLM_KV_LOOP_COUNT, hparams.n_loop); + + if (hparams.n_loop < 1 || hparams.n_layer() % hparams.n_loop != 0) { + throw std::runtime_error("invalid Nanbeige loop count"); + } +} + void llama_model_llama::load_arch_tensors(llama_model_loader &) { LLAMA_LOAD_LOCALS; @@ -45,7 +54,9 @@ void llama_model_llama::load_arch_tensors(llama_model_loader &) { output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED); } - for (int i = 0; i < n_layer; ++i) { + const int n_layer_physical = n_layer / hparams.n_loop; + + for (int i = 0; i < n_layer_physical; ++i) { auto & layer = layers[i]; layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0); @@ -89,6 +100,10 @@ void llama_model_llama::load_arch_tensors(llama_model_loader &) { } } } + + for (int i = n_layer_physical; i < n_layer; ++i) { + layers[i] = layers[i % n_layer_physical]; + } } std::unique_ptr llama_model_llama::build_arch_graph(const llm_graph_params & params) const { @@ -223,6 +238,14 @@ llama_model_llama::graph::graph(const llama_model & model, const llm_grap cur = build_cvec(cur, il); cb(cur, "l_out", il); + const int n_layer_physical = n_layer / hparams.n_loop; + if (hparams.n_loop > 1 && (il + 1) % n_layer_physical == 0 && il + 1 < n_layer) { + cur = build_norm(cur, + model.output_norm, NULL, + LLM_NORM_RMS, il); + cb(cur, "loop_norm", il); + } + // input for next layer inpL = cur; } diff --git a/src/models/models.h b/src/models/models.h index 916b839..9c4b6a3 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -148,6 +148,12 @@ struct llama_model_llama : public llama_model_base { }; +struct llama_model_nanbeige : public llama_model_llama { + llama_model_nanbeige(const struct llama_model_params & params) : llama_model_llama(params) {} + void load_arch_hparams(llama_model_loader & ml) override; +}; + + struct llama_model_llama4 : public llama_model_base { llama_model_llama4(const struct llama_model_params & params) : llama_model_base(params) {} void load_arch_hparams(llama_model_loader & ml) override;