model: Add Mimo v2.5 model support (#22493)
* add mimo-v2.5 support * mimo-v2.5: fix modify_tensors row split * mimi-v2.5: forgot `add_attn_value_scale` plumbing * mimi-v2.5: fix tp dequant to detect tp rows * mimo-v2.5: fix TP iteration to be descending * mimo-v2.5: fix comment * mimo-v2.5: retain fused qkv * mimo-v2.5: missed the attn_value scale during merge * mimo-v2.5: fused QKV needs contiguous for scaling attention value * mimo-v2.5: move `speech_embeddings.` to TextModel filter_tensors * Update src/llama-hparams.h Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@scala.com> * Update src/models/mimo2.cpp Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@scala.com> * Update src/models/mimo2.cpp Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@scala.com> * Update convert_hf_to_gguf.py Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@scala.com> * Update convert_hf_to_gguf.py Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@scala.com> * Update src/models/mimo2.cpp Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@scala.com> * mimo-v2.5: include MTP weights in gguf --------- Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@scala.com>
This commit is contained in:
AesSedai
GitHub
+141
-8
@@ -1063,7 +1063,7 @@ class TextModel(ModelBase):
|
||||
name, gen = item
|
||||
|
||||
# Skip multimodal tensors
|
||||
if name.startswith(("mlp", "vit.", "vpm.", "siglip2.", "conformer.", "merger.", "resampler.", "sound_encoder.", "sound_projection.")) \
|
||||
if name.startswith(("mlp", "vit.", "vpm.", "siglip2.", "conformer.", "merger.", "resampler.", "sound_encoder.", "sound_projection.", "speech_embeddings.")) \
|
||||
or "visual." in name or "vision." in name or "audio." in name or "talker." in name \
|
||||
or "vision_" in name or "audio_" in name or "sam_model" in name \
|
||||
or "token2wav." in name or "code2wav." in name \
|
||||
@@ -9493,10 +9493,126 @@ class MiniMaxM2Model(TextModel):
|
||||
yield from super().modify_tensors(data_torch, name, bid)
|
||||
|
||||
|
||||
@ModelBase.register("MiMoV2FlashForCausalLM")
|
||||
@ModelBase.register("MiMoV2FlashForCausalLM", "MiMoV2ForCausalLM")
|
||||
class MimoV2Model(TextModel):
|
||||
model_arch = gguf.MODEL_ARCH.MIMO2
|
||||
|
||||
# MiMo V2-Flash, V2.5 and V2.5-Pro all ship 3 trained MTP layers under model.mtp.layers.{0,1,2}.
|
||||
# The HF config does not expose the count, so it's hardcoded to match the count found in the safetensors.
|
||||
_n_nextn = 3
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
self.block_count = self.hparams["num_hidden_layers"] + self._n_nextn
|
||||
self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)
|
||||
|
||||
@staticmethod
|
||||
def _tp_aware_qkv_dequant(weight: Tensor, scale_inv: Tensor,
|
||||
n_q: int, n_kv: int, hd: int, vhd: int,
|
||||
bs: int = 128) -> Tensor:
|
||||
# MiMo-V2.5 (TP=4) and V2.5-Pro (TP=8) ship qkv_proj sharded across TP
|
||||
# ranks; per rank, rows are stacked as [Q_per | K_per | V_per].
|
||||
# weight_scale_inv has ceil(rows_per_rank/bs) block-rows per rank (last
|
||||
# may extend past rows_per_rank with phantom rows not in the weight).
|
||||
# Naive repeat_interleave aligns rank 0 only and mis-applies scales to
|
||||
# later ranks once rows_per_rank isn't a multiple of bs.
|
||||
# Re-group the per-rank [Q_per|K_per|V_per] rows into a single fused
|
||||
# [Q | K | V] tensor matching the un-sharded original layout.
|
||||
q_size = n_q * hd
|
||||
k_size = n_kv * hd
|
||||
v_size = n_kv * vhd
|
||||
total_rows = q_size + k_size + v_size
|
||||
if weight.shape[0] != total_rows:
|
||||
raise ValueError(f"qkv_proj weight rows {weight.shape[0]} != q+k+v {total_rows}")
|
||||
|
||||
# detect TP from scale_inv block count, descending order so larger matches first
|
||||
tp = None
|
||||
for cand in (8, 4):
|
||||
if total_rows % cand != 0:
|
||||
continue
|
||||
rpr = total_rows // cand
|
||||
bpr = (rpr + bs - 1) // bs
|
||||
if scale_inv.shape[0] == cand * bpr:
|
||||
tp = cand
|
||||
break
|
||||
if tp is None:
|
||||
raise ValueError(
|
||||
f"qkv_proj: cannot detect TP - scale_inv rows {scale_inv.shape[0]}, "
|
||||
f"q+k+v {total_rows}")
|
||||
|
||||
q_per = q_size // tp
|
||||
k_per = k_size // tp
|
||||
v_per = v_size // tp
|
||||
rows_per_rank = q_per + k_per + v_per
|
||||
blocks_per_rank = (rows_per_rank + bs - 1) // bs
|
||||
|
||||
scale_inv = scale_inv.float()
|
||||
# per-row scale-row index: rank * blocks_per_rank + (rr_in_rank // bs)
|
||||
row_idx = torch.arange(total_rows)
|
||||
rr = row_idx % rows_per_rank
|
||||
rank = row_idx // rows_per_rank
|
||||
scale_row_idx = rank * blocks_per_rank + (rr // bs)
|
||||
# gather: (total_rows, n_col_blocks)
|
||||
scale_per_row_block = scale_inv[scale_row_idx]
|
||||
# expand col-blocks -> cols: each block-col covers `bs` weight cols
|
||||
scale_full = scale_per_row_block.repeat_interleave(bs, dim=1)
|
||||
# crop to weight col count (in case last col-block isn't full)
|
||||
scale_full = scale_full[:, : weight.shape[1]]
|
||||
dequant = weight.float() * scale_full
|
||||
|
||||
if tp == 1:
|
||||
return dequant
|
||||
|
||||
# Re-group per-rank [Q_per|K_per|V_per] rows into unified [Q | K | V]
|
||||
qs, ks, vs = [], [], []
|
||||
for r in range(tp):
|
||||
base = r * rows_per_rank
|
||||
qs.append(dequant[base : base + q_per])
|
||||
ks.append(dequant[base + q_per : base + q_per + k_per])
|
||||
vs.append(dequant[base + q_per + k_per : base + rows_per_rank])
|
||||
return torch.cat(qs + ks + vs, dim=0)
|
||||
|
||||
def dequant_model(self):
|
||||
# Capture raw FP8 (weight, scale_inv) lambdas for qkv_proj BEFORE super
|
||||
# rewrites them with the existing dequant. Replace super's lambda after
|
||||
# it runs so scale_inv removal still happens via the standard path.
|
||||
qkv_overrides: dict[str, tuple[Callable, Callable, int]] = {}
|
||||
qc = self.hparams.get("quantization_config")
|
||||
if isinstance(qc, dict) and qc.get("quant_method") == "fp8":
|
||||
pat = re.compile(r"^model\.layers\.(\d+)\.self_attn\.qkv_proj\.weight_scale_inv$")
|
||||
for name in list(self.model_tensors.keys()):
|
||||
m = pat.match(name)
|
||||
if not m:
|
||||
continue
|
||||
weight_name = name.removesuffix("_scale_inv")
|
||||
if weight_name not in self.model_tensors:
|
||||
continue
|
||||
qkv_overrides[weight_name] = (
|
||||
self.model_tensors[weight_name],
|
||||
self.model_tensors[name],
|
||||
int(m.group(1)),
|
||||
)
|
||||
|
||||
super().dequant_model()
|
||||
|
||||
if not qkv_overrides:
|
||||
return
|
||||
|
||||
n_q = self.hparams["num_attention_heads"]
|
||||
hd = self.hparams["head_dim"]
|
||||
vhd = self.hparams["v_head_dim"]
|
||||
hybrid = self.hparams["hybrid_layer_pattern"]
|
||||
n_layer_text = self.hparams["num_hidden_layers"]
|
||||
for weight_name, (w_fn, s_fn, bid) in qkv_overrides.items():
|
||||
# MTP layers (bid >= n_layer_text) use SWA-style attention dims
|
||||
is_swa = True if bid >= n_layer_text else hybrid[bid] == 1
|
||||
n_kv = self.hparams["swa_num_key_value_heads" if is_swa else "num_key_value_heads"]
|
||||
self.model_tensors[weight_name] = (
|
||||
lambda w_fn=w_fn, s_fn=s_fn, n_q=n_q, n_kv=n_kv, hd=hd, vhd=vhd:
|
||||
MimoV2Model._tp_aware_qkv_dequant(w_fn(), s_fn(), n_q, n_kv, hd, vhd)
|
||||
)
|
||||
|
||||
def set_gguf_parameters(self):
|
||||
super().set_gguf_parameters()
|
||||
|
||||
@@ -9507,11 +9623,14 @@ class MimoV2Model(TextModel):
|
||||
|
||||
n_head_kv = self.hparams["num_key_value_heads"]
|
||||
n_head_kv_swa = self.hparams["swa_num_key_value_heads"]
|
||||
n_head_kv_arr = [n_head_kv_swa if use_swa == 1 else n_head_kv for use_swa in self.hparams["hybrid_layer_pattern"]]
|
||||
# Extend the per-layer pattern with SWA entries for the MTP blocks so the
|
||||
# runtime arrays (sized to extended block_count) are fully populated.
|
||||
hybrid = list(self.hparams["hybrid_layer_pattern"]) + [1] * self._n_nextn
|
||||
n_head_kv_arr = [n_head_kv_swa if use_swa == 1 else n_head_kv for use_swa in hybrid]
|
||||
self.gguf_writer.add_head_count_kv(n_head_kv_arr)
|
||||
|
||||
self.gguf_writer.add_sliding_window(self.hparams["sliding_window"])
|
||||
self.gguf_writer.add_sliding_window_pattern(self.hparams["hybrid_layer_pattern"])
|
||||
self.gguf_writer.add_sliding_window_pattern(hybrid)
|
||||
self.gguf_writer.add_value_length(self.hparams["v_head_dim"])
|
||||
self.gguf_writer.add_expert_count(self.hparams["n_routed_experts"])
|
||||
self.gguf_writer.add_expert_feed_forward_length(self.hparams["moe_intermediate_size"])
|
||||
@@ -9521,6 +9640,12 @@ class MimoV2Model(TextModel):
|
||||
|
||||
self.gguf_writer.add_layer_norm_rms_eps(self.hparams.get("layernorm_epsilon", 1e-5))
|
||||
|
||||
v_scale = self.hparams.get("attention_value_scale")
|
||||
if v_scale is not None:
|
||||
self.gguf_writer.add_attn_value_scale(float(v_scale))
|
||||
|
||||
self.gguf_writer.add_nextn_predict_layers(self._n_nextn)
|
||||
|
||||
_experts: list[dict[str, Tensor]] | None = None
|
||||
|
||||
@classmethod
|
||||
@@ -9530,13 +9655,21 @@ class MimoV2Model(TextModel):
|
||||
if "attention_sink" in name and not name.endswith(".weight"):
|
||||
name += ".weight"
|
||||
|
||||
# TODO: mimo v2 does not indicate the number of next-token-prediction layers, therefore we cannot do the same way as GLM4_MOE
|
||||
if "model.mtp." in name:
|
||||
return None
|
||||
|
||||
return super().filter_tensors((name, gen))
|
||||
|
||||
def modify_tensors(self, data_torch, name, bid):
|
||||
# Remap MTP/NextN tensors to additional layer slots so the standard tensor map handles them.
|
||||
# HF: model.mtp.layers.{i}.foo -> model.layers.{n_layer_text + i}.foo
|
||||
m = re.match(r"^model\.mtp\.layers\.(\d+)\.(.*)$", name)
|
||||
if m is not None:
|
||||
mtp_idx = int(m.group(1))
|
||||
assert mtp_idx < self._n_nextn, f"MTP layer index {mtp_idx} >= _n_nextn ({self._n_nextn})"
|
||||
rest = m.group(2)
|
||||
n_layer_text = self.hparams["num_hidden_layers"]
|
||||
new_bid = n_layer_text + mtp_idx
|
||||
name = f"model.layers.{new_bid}.{rest}"
|
||||
bid = new_bid
|
||||
|
||||
# process the experts separately
|
||||
if name.find("mlp.experts") != -1:
|
||||
n_experts = self.hparams["n_routed_experts"]
|
||||
|
||||
@@ -175,6 +175,7 @@ class Keys:
|
||||
SLIDING_WINDOW = "{arch}.attention.sliding_window"
|
||||
SCALE = "{arch}.attention.scale"
|
||||
OUTPUT_SCALE = "{arch}.attention.output_scale"
|
||||
VALUE_SCALE = "{arch}.attention.value_scale"
|
||||
TEMPERATURE_LENGTH = "{arch}.attention.temperature_length"
|
||||
KEY_LENGTH_MLA = "{arch}.attention.key_length_mla"
|
||||
VALUE_LENGTH_MLA = "{arch}.attention.value_length_mla"
|
||||
@@ -3868,6 +3869,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
|
||||
MODEL_TENSOR.OUTPUT_NORM,
|
||||
MODEL_TENSOR.OUTPUT,
|
||||
MODEL_TENSOR.ATTN_NORM,
|
||||
MODEL_TENSOR.ATTN_QKV,
|
||||
MODEL_TENSOR.ATTN_Q,
|
||||
MODEL_TENSOR.ATTN_K,
|
||||
MODEL_TENSOR.ATTN_V,
|
||||
@@ -3882,6 +3884,10 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
|
||||
MODEL_TENSOR.FFN_DOWN_EXP,
|
||||
MODEL_TENSOR.FFN_UP_EXP,
|
||||
MODEL_TENSOR.FFN_EXP_PROBS_B,
|
||||
MODEL_TENSOR.LAYER_OUT_NORM,
|
||||
MODEL_TENSOR.NEXTN_EH_PROJ,
|
||||
MODEL_TENSOR.NEXTN_ENORM,
|
||||
MODEL_TENSOR.NEXTN_HNORM,
|
||||
],
|
||||
MODEL_ARCH.STEP35: [
|
||||
MODEL_TENSOR.TOKEN_EMBD,
|
||||
|
||||
@@ -943,6 +943,9 @@ class GGUFWriter:
|
||||
def add_attn_output_scale(self, value: float) -> None:
|
||||
self.add_float32(Keys.Attention.OUTPUT_SCALE.format(arch=self.arch), value)
|
||||
|
||||
def add_attn_value_scale(self, value: float) -> None:
|
||||
self.add_float32(Keys.Attention.VALUE_SCALE.format(arch=self.arch), value)
|
||||
|
||||
def add_attn_temperature_length(self, value: int) -> None:
|
||||
self.add_uint32(Keys.Attention.TEMPERATURE_LENGTH.format(arch=self.arch), value)
|
||||
|
||||
|
||||
@@ -232,6 +232,7 @@ static const std::map<llm_kv, const char *> LLM_KV_NAMES = {
|
||||
{ LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, "%s.attention.sliding_window_pattern" },
|
||||
{ LLM_KV_ATTENTION_SCALE, "%s.attention.scale" },
|
||||
{ LLM_KV_ATTENTION_OUTPUT_SCALE, "%s.attention.output_scale" },
|
||||
{ LLM_KV_ATTENTION_VALUE_SCALE, "%s.attention.value_scale" },
|
||||
{ LLM_KV_ATTENTION_TEMPERATURE_LENGTH, "%s.attention.temperature_length" },
|
||||
{ LLM_KV_ATTENTION_TEMPERATURE_SCALE, "%s.attention.temperature_scale" },
|
||||
{ LLM_KV_ATTENTION_KEY_LENGTH_MLA, "%s.attention.key_length_mla" },
|
||||
|
||||
@@ -236,6 +236,7 @@ enum llm_kv {
|
||||
LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN,
|
||||
LLM_KV_ATTENTION_SCALE,
|
||||
LLM_KV_ATTENTION_OUTPUT_SCALE,
|
||||
LLM_KV_ATTENTION_VALUE_SCALE,
|
||||
LLM_KV_ATTENTION_TEMPERATURE_LENGTH,
|
||||
LLM_KV_ATTENTION_TEMPERATURE_SCALE,
|
||||
LLM_KV_ATTENTION_KEY_LENGTH_MLA,
|
||||
|
||||
@@ -166,6 +166,8 @@ struct llama_hparams {
|
||||
float f_attn_out_scale = 0.0f;
|
||||
uint32_t attn_temp_length = 0;
|
||||
|
||||
float f_attn_value_scale = 0.0f;
|
||||
|
||||
bool causal_attn = true;
|
||||
bool use_alibi = false;
|
||||
bool attn_soft_cap = false;
|
||||
|
||||
@@ -268,6 +268,7 @@ void llama_model_saver::add_kv_from_model() {
|
||||
// add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, ???);
|
||||
add_kv(LLM_KV_ATTENTION_SCALE, hparams.f_attention_scale);
|
||||
add_kv(LLM_KV_ATTENTION_OUTPUT_SCALE, hparams.f_attn_out_scale);
|
||||
add_kv(LLM_KV_ATTENTION_VALUE_SCALE, hparams.f_attn_value_scale);
|
||||
add_kv(LLM_KV_ATTENTION_TEMPERATURE_LENGTH, hparams.attn_temp_length);
|
||||
add_kv(LLM_KV_ATTENTION_TEMPERATURE_SCALE, hparams.f_attn_temp_scale);
|
||||
add_kv(LLM_KV_ATTENTION_KEY_LENGTH_MLA, hparams.n_embd_head_k_mla_impl);
|
||||
|
||||
@@ -1671,6 +1671,7 @@ void llama_model::print_info() const {
|
||||
LLAMA_LOG_INFO("%s: f_max_alibi_bias = %.1e\n", __func__, hparams.f_max_alibi_bias);
|
||||
LLAMA_LOG_INFO("%s: f_logit_scale = %.1e\n", __func__, hparams.f_logit_scale);
|
||||
LLAMA_LOG_INFO("%s: f_attn_scale = %.1e\n", __func__, hparams.f_attention_scale);
|
||||
LLAMA_LOG_INFO("%s: f_attn_value_scale = %.4f\n", __func__, hparams.f_attn_value_scale);
|
||||
LLAMA_LOG_INFO("%s: n_ff = %s\n", __func__, print_f([&](uint32_t il) { return hparams.n_ff(il); }, hparams.n_layer).c_str());
|
||||
LLAMA_LOG_INFO("%s: n_expert = %u\n", __func__, hparams.n_expert);
|
||||
LLAMA_LOG_INFO("%s: n_expert_used = %u\n", __func__, hparams.n_expert_used);
|
||||
|
||||
+79
-26
@@ -10,7 +10,16 @@ void llama_model_mimo2::load_arch_hparams(llama_model_loader & ml) {
|
||||
ml.get_key(LLM_KV_ROPE_FREQ_BASE_SWA, hparams.rope_freq_base_train_swa, false);
|
||||
ml.get_key_or_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.swa_layers, hparams.n_layer);
|
||||
|
||||
switch (hparams.n_layer) {
|
||||
float value_scale = 0.0f;
|
||||
if (ml.get_key(LLM_KV_ATTENTION_VALUE_SCALE, value_scale, false) && value_scale != 1.0f) {
|
||||
hparams.f_attn_value_scale = value_scale;
|
||||
}
|
||||
|
||||
ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.nextn_predict_layers, false);
|
||||
GGML_ASSERT(hparams.nextn_predict_layers < hparams.n_layer && "nextn_predict_layers must be < n_layer");
|
||||
hparams.n_layer_kv_from_start = hparams.n_layer - hparams.nextn_predict_layers;
|
||||
|
||||
switch (hparams.n_layer - hparams.nextn_predict_layers) {
|
||||
case 48: type = LLM_TYPE_310B_A15B; break;
|
||||
default: type = LLM_TYPE_UNKNOWN;
|
||||
}
|
||||
@@ -25,32 +34,45 @@ void llama_model_mimo2::load_arch_tensors(llama_model_loader &) {
|
||||
output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
|
||||
output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, 0);
|
||||
|
||||
const uint32_t n_nextn = hparams.nextn_predict_layers;
|
||||
|
||||
for (int i = 0; i < n_layer; ++i) {
|
||||
auto & layer = layers[i];
|
||||
uint32_t n_embd_k_gqa = hparams.n_embd_k_gqa(i);
|
||||
uint32_t n_embd_v_gqa = hparams.n_embd_v_gqa(i);
|
||||
uint32_t n_head = hparams.n_head(i);
|
||||
|
||||
create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_k_gqa, n_embd_v_gqa, 0);
|
||||
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), { n_embd_head_v * n_head, n_embd }, 0);
|
||||
// NextN/MTP layers (the last n_nextn blocks) are preserved but disabled pending support
|
||||
const bool is_nextn = (n_nextn > 0) && (static_cast<uint32_t>(i) >= n_layer - n_nextn);
|
||||
const int skip = is_nextn ? TENSOR_SKIP : 0;
|
||||
|
||||
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
|
||||
layer.attn_sinks = create_tensor(tn(LLM_TENSOR_ATTN_SINKS, "weight", i), {n_head}, TENSOR_NOT_REQUIRED);
|
||||
create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_k_gqa, n_embd_v_gqa, skip);
|
||||
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), { n_embd_head_v * n_head, n_embd }, skip);
|
||||
|
||||
layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
|
||||
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, skip);
|
||||
layer.attn_sinks = create_tensor(tn(LLM_TENSOR_ATTN_SINKS, "weight", i), {n_head}, TENSOR_NOT_REQUIRED | skip);
|
||||
|
||||
layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, skip);
|
||||
|
||||
// non-MoE branch
|
||||
layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, TENSOR_NOT_REQUIRED);
|
||||
layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd}, TENSOR_NOT_REQUIRED);
|
||||
layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, TENSOR_NOT_REQUIRED);
|
||||
layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, TENSOR_NOT_REQUIRED | skip);
|
||||
layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd}, TENSOR_NOT_REQUIRED | skip);
|
||||
layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, TENSOR_NOT_REQUIRED | skip);
|
||||
|
||||
// MoE branch
|
||||
int64_t n_ff_exp = hparams.n_ff_exp;
|
||||
layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, TENSOR_NOT_REQUIRED);
|
||||
layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, TENSOR_NOT_REQUIRED);
|
||||
layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, TENSOR_NOT_REQUIRED);
|
||||
layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, TENSOR_NOT_REQUIRED);
|
||||
layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, TENSOR_NOT_REQUIRED);
|
||||
layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, TENSOR_NOT_REQUIRED | skip);
|
||||
layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, TENSOR_NOT_REQUIRED | skip);
|
||||
layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, TENSOR_NOT_REQUIRED | skip);
|
||||
layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, TENSOR_NOT_REQUIRED | skip);
|
||||
layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, TENSOR_NOT_REQUIRED | skip);
|
||||
|
||||
if (is_nextn) {
|
||||
layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i), {2 * n_embd, n_embd}, skip);
|
||||
layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), {n_embd}, skip);
|
||||
layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", i), {n_embd}, skip);
|
||||
layer.layer_out_norm = create_tensor(tn(LLM_TENSOR_LAYER_OUT_NORM, "weight", i), {n_embd}, skip);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -68,7 +90,12 @@ llama_model_mimo2::graph::graph(const llama_model & model, const llm_graph_param
|
||||
auto * inp_attn = build_attn_inp_kv_iswa();
|
||||
ggml_tensor * inp_out_ids = build_inp_out_ids();
|
||||
|
||||
for (int il = 0; il < n_layer; ++il) {
|
||||
const float v_scale = hparams.f_attn_value_scale;
|
||||
|
||||
// The last hparams.nextn_predict_layers blocks are MTP heads, currently inactive
|
||||
const int n_transformer_layers = n_layer - hparams.nextn_predict_layers;
|
||||
|
||||
for (int il = 0; il < n_transformer_layers; ++il) {
|
||||
ggml_tensor * inpSA = inpL;
|
||||
|
||||
uint32_t n_head_l = hparams.n_head(il);
|
||||
@@ -83,19 +110,39 @@ llama_model_mimo2::graph::graph(const llama_model & model, const llm_graph_param
|
||||
cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il);
|
||||
cb(cur, "attn_norm", il);
|
||||
|
||||
// compute Q and K and RoPE them
|
||||
ggml_tensor * Qcur = build_lora_mm(model.layers[il].wq, cur);
|
||||
cb(Qcur, "Qcur", il);
|
||||
ggml_tensor * Qcur;
|
||||
ggml_tensor * Kcur;
|
||||
ggml_tensor * Vcur;
|
||||
|
||||
ggml_tensor * Kcur = build_lora_mm(model.layers[il].wk, cur);
|
||||
cb(Kcur, "Kcur", il);
|
||||
if (model.layers[il].wqkv) {
|
||||
// Fused qkv_proj - Q/K share head_dim_k, V uses head_dim_v
|
||||
ggml_tensor * qkv = build_lora_mm(model.layers[il].wqkv, cur);
|
||||
cb(qkv, "wqkv", il);
|
||||
|
||||
ggml_tensor * Vcur = build_lora_mm(model.layers[il].wv, cur);
|
||||
cb(Vcur, "Vcur", il);
|
||||
const size_t row_k = ggml_row_size(qkv->type, n_embd_head_k);
|
||||
const size_t row_v = ggml_row_size(qkv->type, n_embd_head_v);
|
||||
const size_t row_full = qkv->nb[1];
|
||||
const size_t k_off = row_k * n_head_l;
|
||||
const size_t v_off = k_off + row_k * n_head_kv_l;
|
||||
|
||||
Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head_k, n_head_l, n_tokens);
|
||||
Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head_k, n_head_kv_l, n_tokens);
|
||||
Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head_v, n_head_kv_l, n_tokens);
|
||||
Qcur = ggml_view_3d(ctx0, qkv, n_embd_head_k, n_head_l, n_tokens, row_k, row_full, 0);
|
||||
Kcur = ggml_view_3d(ctx0, qkv, n_embd_head_k, n_head_kv_l, n_tokens, row_k, row_full, k_off);
|
||||
Vcur = ggml_view_3d(ctx0, qkv, n_embd_head_v, n_head_kv_l, n_tokens, row_v, row_full, v_off);
|
||||
} else {
|
||||
// Split path
|
||||
Qcur = build_lora_mm(model.layers[il].wq, cur);
|
||||
cb(Qcur, "Qcur", il);
|
||||
|
||||
Kcur = build_lora_mm(model.layers[il].wk, cur);
|
||||
cb(Kcur, "Kcur", il);
|
||||
|
||||
Vcur = build_lora_mm(model.layers[il].wv, cur);
|
||||
cb(Vcur, "Vcur", il);
|
||||
|
||||
Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head_k, n_head_l, n_tokens);
|
||||
Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head_k, n_head_kv_l, n_tokens);
|
||||
Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head_v, n_head_kv_l, n_tokens);
|
||||
}
|
||||
|
||||
Qcur = ggml_rope_ext(
|
||||
ctx0, Qcur, inp_pos, nullptr,
|
||||
@@ -118,9 +165,15 @@ llama_model_mimo2::graph::graph(const llama_model & model, const llm_graph_param
|
||||
cur = build_attn(inp_attn,
|
||||
model.layers[il].wo, NULL, model.layers[il].wo_s,
|
||||
Qcur, Kcur, Vcur, nullptr, sinks, nullptr, 1.0f/sqrtf(float(n_embd_head_k)), il);
|
||||
cb(cur, "attn_out", il);
|
||||
|
||||
if (v_scale) {
|
||||
cur = ggml_scale(ctx0, cur, v_scale);
|
||||
cb(cur, "attn_out_scaled", il);
|
||||
}
|
||||
}
|
||||
|
||||
if (il == n_layer - 1 && inp_out_ids) {
|
||||
if (il == n_transformer_layers - 1 && inp_out_ids) {
|
||||
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
|
||||
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user