mirror of
https://github.com/ggml-org/llama.cpp.git
synced 2026-08-25 22:21:03 +02:00
model : Granite-Switch Architecture (#25107)
* granite-switch: add llama.cpp backend (POC, CPU)
New "granite-switch" architecture: a dense, all-attention Granite-4.1
model with N embedded LoRA adapters selected per-token by control tokens.
- gguf-py schema (arch, KV keys, stacked LoRA tensor names) + writer helpers
- conversion/granite.py: GraniteSwitchModel converter (stacks N adapters +
zero base slot into per-projection A/B tensors; emits switch metadata)
- C++ arch registration (llama-arch.{h,cpp}, llama-model.{h,cpp})
- src/models/granite_switch.cpp: load + per-token switched-LoRA graph via
ggml_mul_mat_id over stacked tensors; sticky per-token index + control-token
substitution in llm_graph_input_switch::set_input
- llm_graph_input_switch in src/models/models.h
Runs end-to-end on CPU: convert 3b checkpoint (842 tensors, stacked dim 13)
and generate on both base and control-token paths. Sticky switch state is
single-sequence (POC); full multi-sequence machinery is a follow-up.
* granite-switch: add Mac (Metal) build + mid-sequence switch demo script
Self-contained script to build llama.cpp on Apple Silicon (Metal),
convert the composed 3b checkpoint, and run the crisp mid-sequence
adapter-switch demos verified on Vela:
- answerability: <|answerability|> mid-seq -> "unanswerable"
- query_rewrite: <|query_rewrite|> mid-seq -> {"rewritten_question": ...}
Each demo runs the same prompt twice, differing only by a control token
placed before the assistant turn, so the per-token switch is visible.
* granite-switch mac demo: add -no-cnv so each run is one-shot
The composed model ships a chat template, so llama-completion auto-enables
interactive conversation mode and halts at a `>` prompt after generating,
stalling the script. -no-cnv disables conversation mode: generate once from
the raw prompt and exit (also prints special tokens, making the switch visible).
* granite-switch: replace global sticky index with in-graph router attention
The POC computed the per-token adapter index on the CPU and carried it
across ubatches in ONE global `mutable int32_t poc_sticky_index`, reset
only when a ubatch contained sequence position 0. That global had two
problems:
1. Concurrency: with multiple sequences in a batch it was last-writer-
wins — one sequence's adapter leaked into the others.
2. Multi-turn: an interactive `ollama run` chat continues one KV cache,
so turn 2 never saw position 0 and the index never reset — the
adapter stayed stuck on across turns.
Port the vLLM/HF backend mechanism faithfully: a single-head causal
"router" attention recovers the adapter index in-graph. Per token, only
dim 0 carries signal — Q[0]=1, K[0]=+gain for a control token / -gain
otherwise, V[0]=adapter slot / 0 — and the causal softmax over the single
visible control token recovers that adapter's slot (readback =
clamp(round(V[0]), 0, n_adapters)). gain=15 matches config.py and is
F16-safe (no F32 cache).
The router's K/V live in the model KV cache at an extra layer
R == hparams.router_layer (== n_layer). We bump n_layer_all to n_real+1
so the cache allocator gives the router its own per-sequence slot, and
set n_layer_nextn=1 so n_layer() stays n_real — the decoder loop and
tensor loading are untouched and never reference layer R. The router K is
exempted from the k-shift RoPE loop (its dim-0 value is a literal
magnitude, not a rotation).
Because the selection now lives in the per-sequence KV cache, CONCURRENT
requests are isolated for free (problem 1 fixed; verified by
scratch/concurrent_switch_test.cpp). set_input becomes stateless pure
per-token maps; the global is gone.
Single-switch contract / known limitation, identical to vLLM & HF: the
gain is flat (no recency), so within one sequence there is no mechanism to
revert to base mid-sequence — once an adapter fires it stays on until that
sequence ends (problem 2 is therefore NOT fixed by a faithful copy; vLLM/HF
avoid it only because each served request is a fresh sequence). A client
continuing one KV cache across turns must start a fresh sequence per turn,
or opt into a recency-biased router (a deliberate divergence, not done
here). Documented in granite_switch.cpp and asserted by
scratch/multiturn_leak_test.cpp.
Verified (CPU): both demos unchanged (answerability -> "unanswerable",
query_rewrite -> rewritten query); concurrent two-sequence isolation
passes; multi-turn carry-over matches the vLLM/HF contract.
* granite-switch: drop scratch tests and mac demo for upstream PR
Remove the local-only development artifacts that should not ship in the
upstream PR:
- granite-switch-mac-demo.sh (local Metal build + demo driver)
- scratch/concurrent_switch_test.cpp
- scratch/multiturn_leak_test.cpp
Also drop the now-dangling reference to the scratch tests from the
granite_switch.cpp header comment. Leaves only the core architecture
support (conversion, gguf constants, llama-arch/model/kv-cache, and the
granite_switch graph).
* granite-switch: trim comments to match native llama.cpp style
* granite-switch: trim conversion comments to match native style
* granite-switch: drop unused adapter_ranks metadata
* granite-switch: rename arch to graniteswitch and drop obid alias
* granite-switch: fix non-ASCII comments and document router gain assumption
* granite-switch: drop section comments from constants.py to match native style
* granite-switch: add functional tensor block comments matching Granite4 Vision style
* granite-switch: clarify n_expert_used comment
State the actual constraint: mul_mat_id needs n_expert_used == 1, and
since the GGUF carries expert_count = 0 the generic loader's
n_expert == 0 => n_expert_used == 0 assertion has already passed by the
time load_arch_hparams runs, so it is forced to 1 here.
* granite-switch: note n_layer_nextn reuse has no MTP
The router carving reuses n_layer_nextn, normally the MTP/next-token
count. Clarify in the comment that it is borrowed here purely as the
trailing-layers lever and that there is no MTP head, to spare readers
the double-take.
* granite-switch: rename source file and apply review nits
* granite-switch: don't force LoRA tensors to F16, follow --outtype instead
* granite-switch: drop redundant _permute_qk wrapper, call LlamaModel.permute directly
* granite-switch: read router gain from GGUF (control_token_gain) instead of hardcoding 15.0
* granite-switch: derive n_slots()
* granite-switch: move llm_graph_input_switch into granite-switch.cpp
* granite-switch: cut AI-style narration comments
* granite-switch: collapse multi-line comments
* granite-switch: rename control_token_* maps to adapter_token_*
* granite-switch: cut noise comments
* granite-switch: rename embedded LoRA tensors to <base>.lora_a/lora_b
* granite-switch: GGML_ASSERT token input to avoid UB on embeddings
* granite-switch: TODO for raw embedding input support
* granite-switch: collapse LoRA tensor constants to .lora_a/.lora_b suffix
* granite-switch: drop n_expert_used hack, guard mul_mat_id buft probe
* granite-switch: stop forcing dense expert counts, read from config
* granite-switch: renamed control_token_gain metadata key to router_gain
* granite-switch: trim header comments to match native style
* granite-switch: collapse LoRA tensors to base name + suffix
* granite-switch: inline suffix checks in tensor op resolution
* granite-switch: drop switch-lora struct comment
* granite-switch: guard router layer index and inline n_slots
* granite-switch: group adapter metadata under {arch}.adapters.* namespace
* granite-switch: add hparams.has_rope(il) for KV-shift rope skipping
* granite-switch: skip arch in test-llama-archs (adapter fixture missing, TODO)
* granite-switch: Keys.Adapters namespace + simplify n_slots
* granite-switch: validate substitute token ids against n_vocab
* granite-switch: bound adapter count and lora rank from GGUF
* granite-switch: reject MTP context type when router_layer is set
* granite-switch: throw on bad adapter metadata instead of GGML_ASSERT
* granite-switch: use ASCII +/- in router K signal comment
* granite-switch: document n_layer_nextn repurpose and its leak points
* granite-switch: gate lora_a/lora_b op mapping on router_layer
* granite-switch: label all three preview model sizes
This commit is contained in:
@@ -103,6 +103,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
|
||||
"GraniteMoeForCausalLM": "granite",
|
||||
"GraniteMoeHybridForCausalLM": "granite",
|
||||
"GraniteMoeSharedForCausalLM": "granite",
|
||||
"GraniteSwitchForCausalLM": "granite",
|
||||
"GraniteSpeechForConditionalGeneration": "granite",
|
||||
"GraniteSpeechPlusForConditionalGeneration": "granite",
|
||||
"Grok1ForCausalLM": "grok",
|
||||
|
||||
@@ -123,6 +123,166 @@ class GraniteMoeModel(GraniteModel):
|
||||
yield from super().modify_tensors(data_torch, name, bid)
|
||||
|
||||
|
||||
@ModelBase.register("GraniteSwitchForCausalLM")
|
||||
class GraniteSwitchModel(GraniteMoeModel):
|
||||
"""Dense, all-attention Granite with N per-token embedded LoRA adapters, stacked
|
||||
over the adapter dim with a zero adapter at slot 0 (N = num_adapters + 1)."""
|
||||
model_arch = gguf.MODEL_ARCH.GRANITE_SWITCH
|
||||
|
||||
# permute q/k per-slice below (NORM-rope layout), not via the parent's auto-permute
|
||||
undo_permute = False
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
# the weightless switch reserves one cache slot: one fewer block than num_hidden_layers
|
||||
self.block_count = self.block_count - 1
|
||||
self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)
|
||||
|
||||
self._n_adapters = int(self.hparams["num_adapters"])
|
||||
self._max_lora_rank = int(self.hparams["max_lora_rank"])
|
||||
self._n_slots = self._n_adapters + 1 # +1 for the zero slot at index 0
|
||||
|
||||
n_head = int(self.hparams["num_attention_heads"])
|
||||
n_kv_head = int(self.hparams["num_key_value_heads"])
|
||||
head_dim = (
|
||||
self.hparams.get("projection_head_dim")
|
||||
or self.hparams.get("head_dim")
|
||||
or (self.hparams["hidden_size"] // n_head)
|
||||
)
|
||||
self._n_head = n_head
|
||||
self._n_kv_head = n_kv_head
|
||||
self._head_dim = int(head_dim)
|
||||
self._q_size = n_head * self._head_dim
|
||||
self._kv_size = n_kv_head * self._head_dim
|
||||
|
||||
def set_gguf_parameters(self):
|
||||
super().set_gguf_parameters()
|
||||
|
||||
# dense: pin expert_used_count to 0 (config carries a leftover num_experts_per_tok)
|
||||
if not self.hparams.get("num_local_experts"):
|
||||
self.gguf_writer.add_expert_used_count(0)
|
||||
|
||||
self.gguf_writer.add_adapter_count(self._n_adapters)
|
||||
self.gguf_writer.add_adapter_lora_rank(self._max_lora_rank)
|
||||
self.gguf_writer.add_adapter_token_ids_activate(self.hparams["adapter_token_ids"])
|
||||
self.gguf_writer.add_adapter_token_ids_substitute(self.hparams["adapter_substitute_token_ids"])
|
||||
router_gain = float(self.hparams.get("control_token_gain", 15.0))
|
||||
self.gguf_writer.add_adapter_router_gain(router_gain)
|
||||
logger.info("gguf: (graniteswitch) num_adapters=%s max_lora_rank=%s n_slots=%s router_gain=%s", self._n_adapters, self._max_lora_rank, self._n_slots, router_gain)
|
||||
|
||||
def _lora_a(self, data: Tensor) -> Tensor:
|
||||
# on-disk A: [n_adapters, 1, max_rank, in] -> [n_adapters+1, max_rank, in]
|
||||
a = data.squeeze(1)
|
||||
zero = torch.zeros_like(a[:1])
|
||||
return torch.cat([zero, a], dim=0).contiguous()
|
||||
|
||||
def _lora_b(self, data: Tensor, permute_n_head: int | None = None) -> Tensor:
|
||||
# on-disk B: [n_adapters, 1, out, max_rank] -> [n_adapters+1, out, max_rank]
|
||||
b = data.squeeze(1)
|
||||
if permute_n_head is not None:
|
||||
# permute each adapter's B output rows to match the permuted q/k base
|
||||
b = torch.stack([self.permute(b[i], permute_n_head, permute_n_head) for i in range(b.shape[0])], dim=0)
|
||||
zero = torch.zeros_like(b[:1])
|
||||
return torch.cat([zero, b], dim=0).contiguous()
|
||||
|
||||
def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
|
||||
T = gguf.MODEL_TENSOR
|
||||
|
||||
# skip the weightless switch + control-token buffers (rebuilt at load time)
|
||||
bare = name.split(".")[-1]
|
||||
if (
|
||||
name.startswith("model.switch.") or name.startswith("switch.")
|
||||
or bare in ("adapter_token_ids", "control_to_substitute_lut")
|
||||
):
|
||||
return
|
||||
|
||||
if "self_attn.qkv_proj" in name:
|
||||
if name.endswith("base_layer.weight"):
|
||||
# fused [q|k|v] rows: permute q/k row-blocks for ggml's NORM-rope layout
|
||||
q, k, v = data_torch.split([self._q_size, self._kv_size, self._kv_size], dim=0)
|
||||
q = self.permute(q, self._n_head, self._n_head)
|
||||
k = self.permute(k, self._n_kv_head, self._n_kv_head)
|
||||
fused = torch.cat([q, k, v], dim=0)
|
||||
yield (self.format_tensor_name(T.ATTN_QKV, bid), fused)
|
||||
return
|
||||
if "lora_A_slices." in name:
|
||||
slot = int(name.rsplit(".", 1)[1])
|
||||
key = {0: T.ATTN_Q, 1: T.ATTN_K, 2: T.ATTN_V}[slot]
|
||||
yield (self.format_tensor_name(key, bid, suffix=".lora_a"), self._lora_a(data_torch))
|
||||
return
|
||||
if "lora_B_slices." in name:
|
||||
slot = int(name.rsplit(".", 1)[1])
|
||||
key, ph = {
|
||||
0: (T.ATTN_Q, self._n_head),
|
||||
1: (T.ATTN_K, self._n_kv_head),
|
||||
2: (T.ATTN_V, None),
|
||||
}[slot]
|
||||
yield (self.format_tensor_name(key, bid, suffix=".lora_b"), self._lora_b(data_torch, ph))
|
||||
return
|
||||
raise ValueError(f"Unexpected qkv_proj tensor: {name}")
|
||||
|
||||
if "self_attn.o_proj" in name:
|
||||
if name.endswith("base_layer.weight"):
|
||||
yield (self.format_tensor_name(T.ATTN_OUT, bid), data_torch)
|
||||
return
|
||||
if name.endswith("lora_A"):
|
||||
yield (self.format_tensor_name(T.ATTN_OUT, bid, suffix=".lora_a"), self._lora_a(data_torch))
|
||||
return
|
||||
if name.endswith("lora_B"):
|
||||
yield (self.format_tensor_name(T.ATTN_OUT, bid, suffix=".lora_b"), self._lora_b(data_torch))
|
||||
return
|
||||
raise ValueError(f"Unexpected o_proj tensor: {name}")
|
||||
|
||||
if "shared_mlp.input_linear" in name:
|
||||
ffn = self.hparams["shared_intermediate_size"]
|
||||
if name.endswith("base_layer.weight"):
|
||||
gate, up = data_torch.split([ffn, ffn], dim=0)
|
||||
yield (self.format_tensor_name(T.FFN_GATE, bid), gate)
|
||||
yield (self.format_tensor_name(T.FFN_UP, bid), up)
|
||||
return
|
||||
if "lora_A_slices." in name:
|
||||
slot = int(name.rsplit(".", 1)[1])
|
||||
key = {0: T.FFN_GATE, 1: T.FFN_UP}[slot]
|
||||
yield (self.format_tensor_name(key, bid, suffix=".lora_a"), self._lora_a(data_torch))
|
||||
return
|
||||
if "lora_B_slices." in name:
|
||||
slot = int(name.rsplit(".", 1)[1])
|
||||
key = {0: T.FFN_GATE, 1: T.FFN_UP}[slot]
|
||||
yield (self.format_tensor_name(key, bid, suffix=".lora_b"), self._lora_b(data_torch))
|
||||
return
|
||||
raise ValueError(f"Unexpected shared_mlp.input_linear tensor: {name}")
|
||||
|
||||
if "shared_mlp.output_linear" in name:
|
||||
if name.endswith("base_layer.weight"):
|
||||
yield (self.format_tensor_name(T.FFN_DOWN, bid), data_torch)
|
||||
return
|
||||
if name.endswith("lora_A"):
|
||||
yield (self.format_tensor_name(T.FFN_DOWN, bid, suffix=".lora_a"), self._lora_a(data_torch))
|
||||
return
|
||||
if name.endswith("lora_B"):
|
||||
yield (self.format_tensor_name(T.FFN_DOWN, bid, suffix=".lora_b"), self._lora_b(data_torch))
|
||||
return
|
||||
raise ValueError(f"Unexpected shared_mlp.output_linear tensor: {name}")
|
||||
|
||||
if bid is not None and ".layers." in name and (
|
||||
"input_layernorm" in name or "post_attention_layernorm" in name
|
||||
):
|
||||
key = T.ATTN_NORM if "input_layernorm" in name else T.FFN_NORM
|
||||
yield (self.format_tensor_name(key, bid), data_torch)
|
||||
return
|
||||
|
||||
if name in ("model.embed_tokens.weight", "embed_tokens.weight"):
|
||||
yield (self.format_tensor_name(T.TOKEN_EMBD), data_torch)
|
||||
return
|
||||
if name in ("model.norm.weight", "norm.weight"):
|
||||
yield (self.format_tensor_name(T.OUTPUT_NORM), data_torch)
|
||||
return
|
||||
if name == "lm_head.weight":
|
||||
return # tied to token_embd
|
||||
|
||||
raise ValueError(f"graniteswitch: unhandled tensor {name!r} (bid={bid})")
|
||||
|
||||
|
||||
@ModelBase.register("GraniteMoeHybridForCausalLM", "BambaForCausalLM")
|
||||
class GraniteHybridModel(Mamba2Model, GraniteMoeModel):
|
||||
"""GraniteHybrid is a hybrid SSM + Attention model that uses Mamba2 SSM
|
||||
|
||||
@@ -164,6 +164,13 @@ class Keys:
|
||||
NORM_BEFORE_RESIDUAL = "{arch}.norm_before_residual"
|
||||
NORM_BEFORE_FC = "{arch}.norm_before_fc"
|
||||
|
||||
class Adapters:
|
||||
COUNT = "{arch}.adapters.count"
|
||||
TOKEN_IDS_ACTIVATE = "{arch}.adapters.token_ids_activate"
|
||||
TOKEN_IDS_SUBSTITUTE = "{arch}.adapters.token_ids_substitute"
|
||||
LORA_RANK = "{arch}.adapters.lora_rank"
|
||||
ROUTER_GAIN = "{arch}.adapters.router_gain"
|
||||
|
||||
class Attention:
|
||||
HEAD_COUNT = "{arch}.attention.head_count"
|
||||
HEAD_COUNT_KV = "{arch}.attention.head_count_kv"
|
||||
@@ -527,6 +534,7 @@ class MODEL_ARCH(IntEnum):
|
||||
GRANITE = auto()
|
||||
GRANITE_MOE = auto()
|
||||
GRANITE_HYBRID = auto()
|
||||
GRANITE_SWITCH = auto()
|
||||
CHAMELEON = auto()
|
||||
WAVTOKENIZER_DEC = auto()
|
||||
PLM = auto()
|
||||
@@ -1198,6 +1206,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = {
|
||||
MODEL_ARCH.GRANITE: "granite",
|
||||
MODEL_ARCH.GRANITE_MOE: "granitemoe",
|
||||
MODEL_ARCH.GRANITE_HYBRID: "granitehybrid",
|
||||
MODEL_ARCH.GRANITE_SWITCH: "graniteswitch",
|
||||
MODEL_ARCH.CHAMELEON: "chameleon",
|
||||
MODEL_ARCH.WAVTOKENIZER_DEC: "wavtokenizer-dec",
|
||||
MODEL_ARCH.PLM: "plm",
|
||||
@@ -3972,6 +3981,21 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
|
||||
MODEL_TENSOR.FFN_DOWN,
|
||||
MODEL_TENSOR.FFN_UP,
|
||||
],
|
||||
MODEL_ARCH.GRANITE_SWITCH: [
|
||||
MODEL_TENSOR.TOKEN_EMBD,
|
||||
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,
|
||||
MODEL_TENSOR.ATTN_OUT,
|
||||
MODEL_TENSOR.FFN_NORM,
|
||||
MODEL_TENSOR.FFN_GATE,
|
||||
MODEL_TENSOR.FFN_DOWN,
|
||||
MODEL_TENSOR.FFN_UP,
|
||||
],
|
||||
MODEL_ARCH.CHAMELEON: [
|
||||
MODEL_TENSOR.TOKEN_EMBD,
|
||||
MODEL_TENSOR.OUTPUT_NORM,
|
||||
|
||||
@@ -906,6 +906,21 @@ class GGUFWriter:
|
||||
def add_embedding_scale(self, value: float) -> None:
|
||||
self.add_float32(Keys.LLM.EMBEDDING_SCALE.format(arch=self.arch), value)
|
||||
|
||||
def add_adapter_count(self, count: int) -> None:
|
||||
self.add_uint32(Keys.Adapters.COUNT.format(arch=self.arch), count)
|
||||
|
||||
def add_adapter_token_ids_activate(self, ids: Sequence[int]) -> None:
|
||||
self.add_array(Keys.Adapters.TOKEN_IDS_ACTIVATE.format(arch=self.arch), ids)
|
||||
|
||||
def add_adapter_token_ids_substitute(self, ids: Sequence[int]) -> None:
|
||||
self.add_array(Keys.Adapters.TOKEN_IDS_SUBSTITUTE.format(arch=self.arch), ids)
|
||||
|
||||
def add_adapter_lora_rank(self, rank: int) -> None:
|
||||
self.add_uint32(Keys.Adapters.LORA_RANK.format(arch=self.arch), rank)
|
||||
|
||||
def add_adapter_router_gain(self, gain: float) -> None:
|
||||
self.add_float32(Keys.Adapters.ROUTER_GAIN.format(arch=self.arch), gain)
|
||||
|
||||
def add_wkv_head_size(self, size: int) -> None:
|
||||
self.add_uint32(Keys.WKV.HEAD_SIZE.format(arch=self.arch), size)
|
||||
|
||||
|
||||
@@ -100,6 +100,7 @@ static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = {
|
||||
{ LLM_ARCH_GRANITE, "granite" },
|
||||
{ LLM_ARCH_GRANITE_MOE, "granitemoe" },
|
||||
{ LLM_ARCH_GRANITE_HYBRID, "granitehybrid" },
|
||||
{ LLM_ARCH_GRANITE_SWITCH, "graniteswitch" },
|
||||
{ LLM_ARCH_CHAMELEON, "chameleon" },
|
||||
{ LLM_ARCH_WAVTOKENIZER_DEC, "wavtokenizer-dec" },
|
||||
{ LLM_ARCH_PLM, "plm" },
|
||||
@@ -220,6 +221,11 @@ static const std::map<llm_kv, const char *> LLM_KV_NAMES = {
|
||||
{ LLM_KV_TIME_DECAY_EXTRA_DIM, "%s.time_decay_extra_dim" },
|
||||
{ LLM_KV_RESIDUAL_SCALE, "%s.residual_scale" },
|
||||
{ LLM_KV_EMBEDDING_SCALE, "%s.embedding_scale" },
|
||||
{ LLM_KV_ADAPTER_COUNT, "%s.adapters.count" },
|
||||
{ LLM_KV_ADAPTER_TOKEN_IDS_ACTIVATE, "%s.adapters.token_ids_activate" },
|
||||
{ LLM_KV_ADAPTER_TOKEN_IDS_SUBSTITUTE, "%s.adapters.token_ids_substitute" },
|
||||
{ LLM_KV_ADAPTER_LORA_RANK, "%s.adapters.lora_rank" },
|
||||
{ LLM_KV_ADAPTER_ROUTER_GAIN, "%s.adapters.router_gain" },
|
||||
{ LLM_KV_TOKEN_SHIFT_COUNT, "%s.token_shift_count" },
|
||||
{ LLM_KV_INTERLEAVE_MOE_LAYER_STEP, "%s.interleave_moe_layer_step" },
|
||||
{ LLM_KV_FULL_ATTENTION_INTERVAL, "%s.full_attention_interval" },
|
||||
|
||||
@@ -105,6 +105,7 @@ enum llm_arch {
|
||||
LLM_ARCH_GRANITE,
|
||||
LLM_ARCH_GRANITE_MOE,
|
||||
LLM_ARCH_GRANITE_HYBRID,
|
||||
LLM_ARCH_GRANITE_SWITCH,
|
||||
LLM_ARCH_CHAMELEON,
|
||||
LLM_ARCH_WAVTOKENIZER_DEC,
|
||||
LLM_ARCH_PLM,
|
||||
@@ -225,6 +226,11 @@ enum llm_kv {
|
||||
LLM_KV_TIME_DECAY_EXTRA_DIM,
|
||||
LLM_KV_RESIDUAL_SCALE,
|
||||
LLM_KV_EMBEDDING_SCALE,
|
||||
LLM_KV_ADAPTER_COUNT,
|
||||
LLM_KV_ADAPTER_TOKEN_IDS_ACTIVATE,
|
||||
LLM_KV_ADAPTER_TOKEN_IDS_SUBSTITUTE,
|
||||
LLM_KV_ADAPTER_LORA_RANK,
|
||||
LLM_KV_ADAPTER_ROUTER_GAIN,
|
||||
LLM_KV_TOKEN_SHIFT_COUNT,
|
||||
LLM_KV_INTERLEAVE_MOE_LAYER_STEP,
|
||||
LLM_KV_FULL_ATTENTION_INTERVAL,
|
||||
|
||||
@@ -3602,8 +3602,9 @@ llama_context * llama_init_from_model(
|
||||
model->hparams.pooling_type, params.pooling_type);
|
||||
}
|
||||
|
||||
// router_layer >= 0 means n_layer_nextn is repurposed for a router layer, not real MTP
|
||||
if (params.ctx_type == LLAMA_CONTEXT_TYPE_MTP &&
|
||||
model->hparams.n_layer_nextn == 0) {
|
||||
(model->hparams.n_layer_nextn == 0 || model->hparams.router_layer >= 0)) {
|
||||
LLAMA_LOG_WARN("%s: context type MTP requested but model doesn't contain MTP layers\n", __func__);
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
@@ -277,6 +277,16 @@ bool llama_hparams::has_kv(uint32_t il) const {
|
||||
return true;
|
||||
}
|
||||
|
||||
bool llama_hparams::has_rope(uint32_t il) const {
|
||||
// the router layer stores adapter routing signal, not positional info,
|
||||
// so it must not be RoPE-shifted
|
||||
if (router_layer >= 0 && (int32_t) il == router_layer) {
|
||||
return false;
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
uint32_t llama_hparams::n_layer() const {
|
||||
return n_layer_all - n_layer_nextn;
|
||||
}
|
||||
|
||||
@@ -53,6 +53,10 @@ struct llama_hparams {
|
||||
uint32_t n_embd;
|
||||
uint32_t n_layer_all;
|
||||
uint32_t n_layer_nextn = 0;
|
||||
|
||||
// granite-switch: index of the single-head "router" KV layer that encodes
|
||||
// per-token adapter selection. -1 when the model has no such layer.
|
||||
int32_t router_layer = -1;
|
||||
uint32_t n_expert = 0;
|
||||
uint32_t n_expert_used = 0;
|
||||
uint32_t n_rel_attn_bkts = 0;
|
||||
@@ -371,6 +375,8 @@ struct llama_hparams {
|
||||
|
||||
bool has_kv(uint32_t il) const;
|
||||
|
||||
bool has_rope(uint32_t il) const;
|
||||
|
||||
// number of effective layers (excludes nextn layers)
|
||||
uint32_t n_layer() const;
|
||||
|
||||
|
||||
@@ -1931,6 +1931,10 @@ ggml_cgraph * llama_kv_cache::build_graph_shift(llm_graph_result * res, llama_co
|
||||
for (const auto & layer : layers) {
|
||||
const uint32_t il = layer.il;
|
||||
|
||||
if (!hparams.has_rope(il)) {
|
||||
continue;
|
||||
}
|
||||
|
||||
const int64_t n_head_kv = hparams.n_head_kv(il);
|
||||
const int64_t n_embd_k_gqa = hparams.n_embd_k_gqa(il);
|
||||
|
||||
|
||||
+12
-12
@@ -937,10 +937,11 @@ static bool weight_buft_supported(const llama_hparams & hparams, ggml_tensor * w
|
||||
} break;
|
||||
case GGML_OP_MUL_MAT_ID:
|
||||
{
|
||||
const int n_expert_used = hparams.n_expert_used;
|
||||
GGML_ASSERT(n_expert_used > 0);
|
||||
ggml_tensor * b = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, w->ne[0], n_expert_used, 512);
|
||||
ggml_tensor * ids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_expert_used, 512);
|
||||
// Used for either MoE expert routing or embedded adapter routing
|
||||
const int n_ids_used = hparams.router_layer >= 0 ? 1 : hparams.n_expert_used;
|
||||
GGML_ASSERT(n_ids_used > 0);
|
||||
ggml_tensor * b = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, w->ne[0], n_ids_used, 512);
|
||||
ggml_tensor * ids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_ids_used, 512);
|
||||
op_tensor = ggml_mul_mat_id(ctx, w, b, ids);
|
||||
} break;
|
||||
case GGML_OP_ADD:
|
||||
@@ -1123,15 +1124,14 @@ struct ggml_tensor * llama_model_loader::create_tensor(
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
// tensors with "bias" suffix are always used with GGML_OP_ADD or GGML_OP_ADD_ID
|
||||
// tensors with "bias" suffix are always used with GGML_OP_ADD or GGML_OP_ADD_ID;
|
||||
// embedded-adapter ".lora_a"/".lora_b" tensors are always used with GGML_OP_MUL_MAT_ID
|
||||
ggml_op op;
|
||||
bool bias = tn.suffix != nullptr && strcmp(tn.suffix, "bias") == 0;
|
||||
if (bias) {
|
||||
if (info.op == GGML_OP_MUL_MAT_ID) {
|
||||
op = GGML_OP_ADD_ID;
|
||||
} else {
|
||||
op = GGML_OP_ADD;
|
||||
}
|
||||
if (tn.suffix != nullptr && strcmp(tn.suffix, "bias") == 0) {
|
||||
op = info.op == GGML_OP_MUL_MAT_ID ? GGML_OP_ADD_ID : GGML_OP_ADD;
|
||||
} else if (hparams.router_layer >= 0 && tn.suffix != nullptr &&
|
||||
(strcmp(tn.suffix, "lora_a") == 0 || strcmp(tn.suffix, "lora_b") == 0)) {
|
||||
op = GGML_OP_MUL_MAT_ID;
|
||||
} else {
|
||||
op = info.op;
|
||||
}
|
||||
|
||||
@@ -234,6 +234,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params
|
||||
return new llama_model_granite(params);
|
||||
case LLM_ARCH_GRANITE_MOE:
|
||||
return new llama_model_granite_moe(params);
|
||||
case LLM_ARCH_GRANITE_SWITCH:
|
||||
return new llama_model_granite_switch(params);
|
||||
case LLM_ARCH_MINICPM:
|
||||
return new llama_model_minicpm(params);
|
||||
case LLM_ARCH_GRANITE_HYBRID:
|
||||
@@ -1912,6 +1914,7 @@ void llama_model::print_info() const {
|
||||
arch == LLM_ARCH_GRANITE ||
|
||||
arch == LLM_ARCH_GRANITE_MOE ||
|
||||
arch == LLM_ARCH_GRANITE_HYBRID ||
|
||||
arch == LLM_ARCH_GRANITE_SWITCH ||
|
||||
arch == LLM_ARCH_NEMOTRON_H_MOE) {
|
||||
LLAMA_LOG_INFO("%s: f_embedding_scale = %f\n", __func__, hparams.f_embedding_scale);
|
||||
LLAMA_LOG_INFO("%s: f_residual_scale = %f\n", __func__, hparams.f_residual_scale);
|
||||
@@ -2596,6 +2599,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) {
|
||||
case LLM_ARCH_GRANITE:
|
||||
case LLM_ARCH_GRANITE_MOE:
|
||||
case LLM_ARCH_GRANITE_HYBRID:
|
||||
case LLM_ARCH_GRANITE_SWITCH:
|
||||
case LLM_ARCH_CHAMELEON:
|
||||
case LLM_ARCH_BAILINGMOE:
|
||||
case LLM_ARCH_NEO_BERT:
|
||||
|
||||
@@ -223,6 +223,24 @@ struct llama_layer_nextn {
|
||||
struct ggml_tensor * shared_head_norm = nullptr;
|
||||
};
|
||||
|
||||
struct llama_layer_switch_lora {
|
||||
struct ggml_tensor * a_q = nullptr;
|
||||
struct ggml_tensor * b_q = nullptr;
|
||||
struct ggml_tensor * a_k = nullptr;
|
||||
struct ggml_tensor * b_k = nullptr;
|
||||
struct ggml_tensor * a_v = nullptr;
|
||||
struct ggml_tensor * b_v = nullptr;
|
||||
struct ggml_tensor * a_o = nullptr;
|
||||
struct ggml_tensor * b_o = nullptr;
|
||||
|
||||
struct ggml_tensor * a_gate = nullptr;
|
||||
struct ggml_tensor * b_gate = nullptr;
|
||||
struct ggml_tensor * a_up = nullptr;
|
||||
struct ggml_tensor * b_up = nullptr;
|
||||
struct ggml_tensor * a_down = nullptr;
|
||||
struct ggml_tensor * b_down = nullptr;
|
||||
};
|
||||
|
||||
struct llama_layer {
|
||||
// normalization
|
||||
struct ggml_tensor * attn_norm = nullptr;
|
||||
@@ -533,6 +551,8 @@ struct llama_layer {
|
||||
struct llama_layer_shortconv shortconv;
|
||||
|
||||
struct llama_layer_nextn nextn;
|
||||
|
||||
struct llama_layer_switch_lora switch_lora;
|
||||
};
|
||||
|
||||
struct llama_device {
|
||||
|
||||
@@ -0,0 +1,426 @@
|
||||
#include "models.h"
|
||||
|
||||
#include <cmath>
|
||||
|
||||
void llama_model_granite_switch::load_arch_hparams(llama_model_loader & ml) {
|
||||
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
|
||||
ml.get_key(LLM_KV_LOGIT_SCALE, hparams.f_logit_scale);
|
||||
ml.get_key(LLM_KV_RESIDUAL_SCALE, hparams.f_residual_scale, false);
|
||||
ml.get_key(LLM_KV_EMBEDDING_SCALE, hparams.f_embedding_scale, false);
|
||||
ml.get_key(LLM_KV_ATTENTION_SCALE, hparams.f_attention_scale, false);
|
||||
|
||||
bool rope_finetuned = true;
|
||||
ml.get_key(LLM_KV_ROPE_SCALING_FINETUNED, rope_finetuned, false);
|
||||
hparams.rope_finetuned = rope_finetuned;
|
||||
|
||||
switch (hparams.n_layer()) {
|
||||
case 40: type = hparams.n_embd == 4096 ? LLM_TYPE_8B : LLM_TYPE_3B; break;
|
||||
case 64: type = LLM_TYPE_30B; break;
|
||||
default: type = LLM_TYPE_UNKNOWN;
|
||||
}
|
||||
|
||||
ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, /* required */ false);
|
||||
|
||||
ml.get_key(LLM_KV_ADAPTER_COUNT, n_adapters);
|
||||
ml.get_key(LLM_KV_ADAPTER_LORA_RANK, max_lora_rank);
|
||||
ml.get_key(LLM_KV_ADAPTER_ROUTER_GAIN, router_gain, /* required */ false);
|
||||
|
||||
// bound counts that size tensors
|
||||
if (n_adapters > 4096) {
|
||||
throw std::runtime_error(format("graniteswitch: invalid adapter count %u", n_adapters));
|
||||
}
|
||||
if (max_lora_rank > 4096) {
|
||||
throw std::runtime_error(format("graniteswitch: invalid lora rank %u", max_lora_rank));
|
||||
}
|
||||
|
||||
std::vector<llama_token> token_ids;
|
||||
std::vector<llama_token> substitute_ids;
|
||||
ml.get_arr(LLM_KV_ADAPTER_TOKEN_IDS_ACTIVATE, token_ids);
|
||||
ml.get_arr(LLM_KV_ADAPTER_TOKEN_IDS_SUBSTITUTE, substitute_ids);
|
||||
|
||||
if (token_ids.size() != n_adapters || substitute_ids.size() != n_adapters) {
|
||||
throw std::runtime_error(format(
|
||||
"graniteswitch: adapter token id arrays (%zu activate, %zu substitute) do not match adapter count %u",
|
||||
token_ids.size(), substitute_ids.size(), n_adapters));
|
||||
}
|
||||
|
||||
adapter_token_to_slot.clear();
|
||||
adapter_token_to_substitute.clear();
|
||||
for (uint32_t i = 0; i < n_adapters; ++i) {
|
||||
// adapter i -> stacked slot i+1 (slot 0 is the base/zero delta)
|
||||
adapter_token_to_slot[token_ids[i]] = (int32_t) (i + 1);
|
||||
adapter_token_to_substitute[token_ids[i]] = substitute_ids[i];
|
||||
}
|
||||
|
||||
// extra single-head attention layer at the END (index n_real) holds the router
|
||||
// K/V. reusing n_layer_nextn keeps n_layer() == n_real, so the regular layers
|
||||
// keep their indices and the KV cache shift/defrag skips the router layer.
|
||||
// n_layer_nextn is repurposed here (no MTP): it leaks as 1 into the
|
||||
// llama_model_n_layer_nextn() getter and a re-saved nextn_predict_layers
|
||||
const uint32_t n_real = hparams.n_layer();
|
||||
if (n_real >= LLAMA_MAX_LAYERS) {
|
||||
throw std::runtime_error(format("graniteswitch: block count %u exceeds LLAMA_MAX_LAYERS", n_real));
|
||||
}
|
||||
hparams.router_layer = (int32_t) n_real;
|
||||
hparams.n_layer_all = n_real + 1;
|
||||
hparams.n_layer_nextn = 1;
|
||||
|
||||
hparams.n_head_arr[n_real] = 1;
|
||||
hparams.n_head_kv_arr[n_real] = 1;
|
||||
hparams.n_ff_arr[n_real] = 0;
|
||||
}
|
||||
|
||||
void llama_model_granite_switch::load_arch_tensors(llama_model_loader &) {
|
||||
LLAMA_LOAD_LOCALS;
|
||||
|
||||
const int64_t n_slots = (int64_t) n_adapters + 1; // slot 0 = base/zero delta
|
||||
const int64_t n_rank = (int64_t) max_lora_rank;
|
||||
const int64_t n_embd_q = n_embd_head_k * n_head;
|
||||
const int64_t n_embd_kv = n_embd_k_gqa;
|
||||
|
||||
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
|
||||
|
||||
// substitute ids index tok_embd rows directly; range-check against n_vocab
|
||||
for (const auto & kv : adapter_token_to_substitute) {
|
||||
const llama_token sub = kv.second;
|
||||
if (sub < 0 || (int64_t) sub >= n_vocab) {
|
||||
throw std::runtime_error(format(
|
||||
"graniteswitch: substitute token id %d out of range [0, %d)", sub, (int) n_vocab));
|
||||
}
|
||||
}
|
||||
|
||||
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}, TENSOR_NOT_REQUIRED);
|
||||
if (output == NULL) {
|
||||
output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED);
|
||||
}
|
||||
|
||||
for (int i = 0; i < n_layer; ++i) {
|
||||
auto & layer = layers[i];
|
||||
|
||||
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
|
||||
|
||||
layer.wqkv = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", i), {n_embd, n_embd_q + 2*n_embd_kv}, 0);
|
||||
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_q, n_embd}, 0);
|
||||
|
||||
layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
|
||||
|
||||
layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0);
|
||||
layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd}, 0);
|
||||
layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0);
|
||||
|
||||
auto & sl = layer.switch_lora;
|
||||
|
||||
sl.a_q = create_tensor(tn(LLM_TENSOR_ATTN_Q, "lora_a", i), {n_embd, n_rank, n_slots}, 0);
|
||||
sl.b_q = create_tensor(tn(LLM_TENSOR_ATTN_Q, "lora_b", i), {n_rank, n_embd_q, n_slots}, 0);
|
||||
sl.a_k = create_tensor(tn(LLM_TENSOR_ATTN_K, "lora_a", i), {n_embd, n_rank, n_slots}, 0);
|
||||
sl.b_k = create_tensor(tn(LLM_TENSOR_ATTN_K, "lora_b", i), {n_rank, n_embd_kv, n_slots}, 0);
|
||||
sl.a_v = create_tensor(tn(LLM_TENSOR_ATTN_V, "lora_a", i), {n_embd, n_rank, n_slots}, 0);
|
||||
sl.b_v = create_tensor(tn(LLM_TENSOR_ATTN_V, "lora_b", i), {n_rank, n_embd_kv, n_slots}, 0);
|
||||
|
||||
sl.a_o = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "lora_a", i), {n_embd_q, n_rank, n_slots}, 0);
|
||||
sl.b_o = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "lora_b", i), {n_rank, n_embd, n_slots}, 0);
|
||||
|
||||
sl.a_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "lora_a", i), {n_embd, n_rank, n_slots}, 0);
|
||||
sl.b_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "lora_b", i), {n_rank, n_ff, n_slots}, 0);
|
||||
sl.a_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "lora_a", i), {n_embd, n_rank, n_slots}, 0);
|
||||
sl.b_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "lora_b", i), {n_rank, n_ff, n_slots}, 0);
|
||||
sl.a_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "lora_a", i), { n_ff, n_rank, n_slots}, 0);
|
||||
sl.b_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "lora_b", i), {n_rank, n_embd, n_slots}, 0);
|
||||
}
|
||||
}
|
||||
|
||||
class llm_graph_input_switch : public llm_graph_input_i {
|
||||
public:
|
||||
llm_graph_input_switch(const llama_model_granite_switch & smodel) : smodel(smodel) {}
|
||||
virtual ~llm_graph_input_switch() = default;
|
||||
|
||||
void set_input(const llama_ubatch * ubatch) override;
|
||||
|
||||
ggml_tensor * sub_tokens = nullptr; // I32 [n_tokens] adapter-substituted token ids
|
||||
ggml_tensor * router_ksig = nullptr; // F32 [n_tokens] router K signal (+/-gain)
|
||||
ggml_tensor * router_vval = nullptr; // F32 [n_tokens] router V value (adapter slot / 0)
|
||||
ggml_tensor * router_q = nullptr; // F32 [n_tokens] router Q value (constant 1.0)
|
||||
|
||||
const llama_model_granite_switch & smodel;
|
||||
};
|
||||
|
||||
// K dim-0 is +gain for an adapter token, -gain otherwise; the causal softmax then
|
||||
// lets a single visible adapter token dominate so the readback recovers its slot.
|
||||
void llm_graph_input_switch::set_input(const llama_ubatch * ubatch) {
|
||||
if (!ubatch->token) {
|
||||
return;
|
||||
}
|
||||
|
||||
const int64_t n_tokens = ubatch->n_tokens;
|
||||
|
||||
std::vector<int32_t> sub (n_tokens);
|
||||
std::vector<float> ksig(n_tokens);
|
||||
std::vector<float> vval(n_tokens);
|
||||
std::vector<float> q (n_tokens, 1.0f);
|
||||
|
||||
for (int64_t i = 0; i < n_tokens; ++i) {
|
||||
const llama_token tok = ubatch->token[i];
|
||||
|
||||
const auto it = smodel.adapter_token_to_slot.find(tok);
|
||||
if (it != smodel.adapter_token_to_slot.end()) {
|
||||
ksig[i] = +smodel.router_gain;
|
||||
vval[i] = (float) it->second;
|
||||
} else {
|
||||
ksig[i] = -smodel.router_gain;
|
||||
vval[i] = 0.0f;
|
||||
}
|
||||
|
||||
const auto sit = smodel.adapter_token_to_substitute.find(tok);
|
||||
sub[i] = (sit != smodel.adapter_token_to_substitute.end())
|
||||
? (int32_t) sit->second
|
||||
: (int32_t) tok;
|
||||
}
|
||||
|
||||
ggml_backend_tensor_set(sub_tokens, sub.data(), 0, n_tokens*ggml_element_size(sub_tokens));
|
||||
ggml_backend_tensor_set(router_ksig, ksig.data(), 0, n_tokens*ggml_element_size(router_ksig));
|
||||
ggml_backend_tensor_set(router_vval, vval.data(), 0, n_tokens*ggml_element_size(router_vval));
|
||||
ggml_backend_tensor_set(router_q, q.data(), 0, n_tokens*ggml_element_size(router_q));
|
||||
}
|
||||
|
||||
std::unique_ptr<llm_graph_context> llama_model_granite_switch::build_arch_graph(const llm_graph_params & params) const {
|
||||
return std::make_unique<graph>(*this, params);
|
||||
}
|
||||
|
||||
// per-token switched LoRA delta: B_a*(A_a*x), adapter selected per token via ids.
|
||||
// cur: {n_in, n_tokens}, ids: {n_tokens} -> {n_out, n_tokens}
|
||||
ggml_tensor * llama_model_granite_switch::graph::build_switched_lora_delta(
|
||||
ggml_tensor * lora_a,
|
||||
ggml_tensor * lora_b,
|
||||
ggml_tensor * cur,
|
||||
ggml_tensor * ids) {
|
||||
const int64_t n_in = cur->ne[0];
|
||||
const int64_t n_tokens = cur->ne[1];
|
||||
|
||||
ggml_tensor * x = ggml_reshape_3d(ctx0, cur, n_in, 1, n_tokens);
|
||||
ggml_tensor * ids2 = ggml_reshape_2d(ctx0, ids, 1, n_tokens);
|
||||
|
||||
ggml_tensor * a = ggml_mul_mat_id(ctx0, lora_a, x, ids2); // {max_rank, 1, n_tokens}
|
||||
ggml_tensor * d = ggml_mul_mat_id(ctx0, lora_b, a, ids2); // {n_out, 1, n_tokens}
|
||||
|
||||
return ggml_reshape_2d(ctx0, d, d->ne[0], n_tokens);
|
||||
}
|
||||
|
||||
ggml_tensor * llama_model_granite_switch::graph::build_switched_lora_mm(
|
||||
ggml_tensor * w,
|
||||
ggml_tensor * lora_a,
|
||||
ggml_tensor * lora_b,
|
||||
ggml_tensor * cur,
|
||||
ggml_tensor * ids) {
|
||||
ggml_tensor * base = ggml_mul_mat(ctx0, w, cur);
|
||||
ggml_tensor * delta = build_switched_lora_delta(lora_a, lora_b, cur, ids);
|
||||
return ggml_add(ctx0, base, delta);
|
||||
}
|
||||
|
||||
llama_model_granite_switch::graph::graph(
|
||||
const llama_model & model,
|
||||
const llm_graph_params & params)
|
||||
: llm_graph_context(params) {
|
||||
|
||||
const auto & smodel = static_cast<const llama_model_granite_switch &>(model);
|
||||
|
||||
// TODO: support raw embedding input (multimodal / pre-embedded tokens) when needed
|
||||
GGML_ASSERT(ubatch.token && "granite-switch requires token input");
|
||||
|
||||
const int64_t n_embd_head = hparams.n_embd_head_v();
|
||||
GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());
|
||||
GGML_ASSERT(n_embd_head == n_rot);
|
||||
|
||||
auto inp_switch = std::make_unique<llm_graph_input_switch>(smodel);
|
||||
inp_switch->sub_tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
|
||||
inp_switch->router_ksig = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, n_tokens);
|
||||
inp_switch->router_vval = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, n_tokens);
|
||||
inp_switch->router_q = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, n_tokens);
|
||||
ggml_set_input(inp_switch->sub_tokens);
|
||||
ggml_set_input(inp_switch->router_ksig);
|
||||
ggml_set_input(inp_switch->router_vval);
|
||||
ggml_set_input(inp_switch->router_q);
|
||||
ggml_tensor * sub_tokens = inp_switch->sub_tokens;
|
||||
ggml_tensor * router_ksig = inp_switch->router_ksig;
|
||||
ggml_tensor * router_vval = inp_switch->router_vval;
|
||||
ggml_tensor * router_q = inp_switch->router_q;
|
||||
res->add_input(std::move(inp_switch));
|
||||
|
||||
// embed the substituted ids directly; build_inp_embd would embed the raw tokens
|
||||
ggml_tensor * inpL = ggml_get_rows(ctx0, model.tok_embd, sub_tokens);
|
||||
if (hparams.f_embedding_scale != 0.0f) {
|
||||
inpL = ggml_scale(ctx0, inpL, hparams.f_embedding_scale);
|
||||
}
|
||||
cb(inpL, "inp_embd", -1);
|
||||
|
||||
ggml_tensor * inp_pos = nullptr;
|
||||
if (hparams.rope_finetuned) {
|
||||
inp_pos = build_inp_pos();
|
||||
}
|
||||
auto * inp_attn = build_attn_inp_kv();
|
||||
|
||||
// single causal head at layer R recovers the adapter index in-graph: only dim 0
|
||||
// carries signal (Q[0]=1, K[0]=+/-gain, V[0]=slot/0), the rest is zero-padded.
|
||||
const int R = hparams.router_layer;
|
||||
GGML_ASSERT(R >= 0);
|
||||
auto router_lane = [&](ggml_tensor * sig1d) {
|
||||
ggml_tensor * t = ggml_reshape_3d(ctx0, sig1d, 1, 1, n_tokens);
|
||||
return ggml_pad(ctx0, t, (int) n_embd_head - 1, 0, 0, 0);
|
||||
};
|
||||
ggml_tensor * Qr = router_lane(router_q);
|
||||
ggml_tensor * Kr = router_lane(router_ksig);
|
||||
ggml_tensor * Vr = router_lane(router_vval);
|
||||
|
||||
ggml_tensor * router_out = build_attn(inp_attn,
|
||||
nullptr, nullptr, nullptr,
|
||||
Qr, Kr, Vr, nullptr, nullptr, nullptr, /*kq_scale=*/1.0f, /*il=*/R);
|
||||
cb(router_out, "router_out", R);
|
||||
|
||||
// row 0 of router_out is the attended slot; clamp+round to an I32 index
|
||||
ggml_tensor * slot_f = ggml_cont(ctx0,
|
||||
ggml_view_2d(ctx0, router_out, 1, n_tokens, router_out->nb[1], 0));
|
||||
slot_f = ggml_reshape_1d(ctx0, slot_f, n_tokens);
|
||||
slot_f = ggml_clamp(ctx0, slot_f, 0.0f, (float) smodel.n_adapters);
|
||||
slot_f = ggml_round(ctx0, slot_f);
|
||||
ggml_tensor * adapter_ids = ggml_cast(ctx0, slot_f, GGML_TYPE_I32);
|
||||
cb(adapter_ids, "adapter_ids", -1);
|
||||
|
||||
ggml_tensor * inp_out_ids = build_inp_out_ids();
|
||||
|
||||
ggml_tensor * cur;
|
||||
|
||||
for (int il = 0; il < n_layer; ++il) {
|
||||
ggml_tensor * inpSA = inpL;
|
||||
|
||||
cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il);
|
||||
cb(cur, "attn_norm", il);
|
||||
|
||||
cur = build_attention_layer(cur, inp_pos, adapter_ids, inp_attn, model, n_embd_head, il);
|
||||
|
||||
if (il == n_layer - 1 && inp_out_ids) {
|
||||
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
|
||||
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
|
||||
// keep adapter_ids aligned to the kept rows (2D round-trip for get_rows)
|
||||
const int64_t n_out = inp_out_ids->ne[0];
|
||||
adapter_ids = ggml_get_rows(ctx0,
|
||||
ggml_reshape_2d(ctx0, adapter_ids, 1, adapter_ids->ne[0]), inp_out_ids);
|
||||
adapter_ids = ggml_reshape_1d(ctx0, adapter_ids, n_out);
|
||||
}
|
||||
|
||||
cur = build_layer_ffn(cur, inpSA, adapter_ids, model, il);
|
||||
|
||||
inpL = cur;
|
||||
}
|
||||
|
||||
cur = inpL;
|
||||
|
||||
cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1);
|
||||
cb(cur, "result_norm", -1);
|
||||
res->t_embd = cur;
|
||||
|
||||
cur = build_lora_mm(model.output, cur, model.output_s);
|
||||
|
||||
cur = ggml_scale(ctx0, cur, 1.0f / hparams.f_logit_scale);
|
||||
cb(cur, "result_output", -1);
|
||||
res->t_logits = cur;
|
||||
|
||||
ggml_build_forward_expand(gf, cur);
|
||||
}
|
||||
|
||||
ggml_tensor * llama_model_granite_switch::graph::build_attention_layer(
|
||||
ggml_tensor * cur,
|
||||
ggml_tensor * inp_pos,
|
||||
ggml_tensor * adapter_ids,
|
||||
llm_graph_input_attn_kv * inp_attn,
|
||||
const llama_model & model,
|
||||
const int64_t n_embd_head,
|
||||
const int il) {
|
||||
|
||||
const auto & layer = model.layers[il];
|
||||
const auto & sl = layer.switch_lora;
|
||||
|
||||
const int64_t n_head = hparams.n_head(il);
|
||||
const int64_t n_head_kv = hparams.n_head_kv(il);
|
||||
|
||||
ggml_tensor * qkv = ggml_mul_mat(ctx0, layer.wqkv, cur);
|
||||
cb(qkv, "wqkv", il);
|
||||
|
||||
const int64_t n_embd_q = n_embd_head * n_head;
|
||||
const int64_t n_embd_kv = n_embd_head * n_head_kv;
|
||||
|
||||
// slice fused qkv into Q/K/V, made contiguous so LoRA deltas can be added
|
||||
ggml_tensor * Qcur = ggml_cont(ctx0, ggml_view_2d(ctx0, qkv, n_embd_q, qkv->ne[1], qkv->nb[1], 0));
|
||||
ggml_tensor * Kcur = ggml_cont(ctx0, ggml_view_2d(ctx0, qkv, n_embd_kv, qkv->ne[1], qkv->nb[1], n_embd_q*ggml_element_size(qkv)));
|
||||
ggml_tensor * Vcur = ggml_cont(ctx0, ggml_view_2d(ctx0, qkv, n_embd_kv, qkv->ne[1], qkv->nb[1], (n_embd_q + n_embd_kv)*ggml_element_size(qkv)));
|
||||
|
||||
Qcur = ggml_add(ctx0, Qcur, build_switched_lora_delta(sl.a_q, sl.b_q, cur, adapter_ids));
|
||||
Kcur = ggml_add(ctx0, Kcur, build_switched_lora_delta(sl.a_k, sl.b_k, cur, adapter_ids));
|
||||
Vcur = ggml_add(ctx0, Vcur, build_switched_lora_delta(sl.a_v, sl.b_v, cur, adapter_ids));
|
||||
|
||||
Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens);
|
||||
Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens);
|
||||
Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens);
|
||||
|
||||
if (hparams.rope_finetuned) {
|
||||
ggml_tensor * rope_factors = model.get_rope_factors(cparams, il);
|
||||
Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, rope_factors,
|
||||
n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
|
||||
ext_factor, attn_factor, beta_fast, beta_slow);
|
||||
Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, rope_factors,
|
||||
n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
|
||||
ext_factor, attn_factor, beta_fast, beta_slow);
|
||||
}
|
||||
cb(Qcur, "Qcur", il);
|
||||
cb(Kcur, "Kcur", il);
|
||||
cb(Vcur, "Vcur", il);
|
||||
|
||||
const float kq_scale = hparams.f_attention_scale == 0.0f
|
||||
? 1.0f/sqrtf(float(n_embd_head)) : hparams.f_attention_scale;
|
||||
|
||||
// wo = nullptr so build_attn returns concatenated heads; o-proj is switched below
|
||||
ggml_tensor * attn = build_attn(inp_attn,
|
||||
nullptr, nullptr, nullptr,
|
||||
Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);
|
||||
cb(attn, "attn_pre_o", il);
|
||||
|
||||
cur = build_switched_lora_mm(layer.wo, sl.a_o, sl.b_o, attn, adapter_ids);
|
||||
cb(cur, "attn_out", il);
|
||||
return cur;
|
||||
}
|
||||
|
||||
ggml_tensor * llama_model_granite_switch::graph::build_layer_ffn(
|
||||
ggml_tensor * cur,
|
||||
ggml_tensor * inpSA,
|
||||
ggml_tensor * adapter_ids,
|
||||
const llama_model & model,
|
||||
const int il) {
|
||||
|
||||
const auto & layer = model.layers[il];
|
||||
const auto & sl = layer.switch_lora;
|
||||
|
||||
if (hparams.f_residual_scale) {
|
||||
cur = ggml_scale(ctx0, cur, hparams.f_residual_scale);
|
||||
}
|
||||
ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);
|
||||
cb(ffn_inp, "ffn_inp", il);
|
||||
|
||||
cur = build_norm(ffn_inp, layer.ffn_norm, NULL, LLM_NORM_RMS, il);
|
||||
cb(cur, "ffn_norm", il);
|
||||
|
||||
ggml_tensor * g = build_switched_lora_mm(layer.ffn_gate, sl.a_gate, sl.b_gate, cur, adapter_ids);
|
||||
ggml_tensor * u = build_switched_lora_mm(layer.ffn_up, sl.a_up, sl.b_up, cur, adapter_ids);
|
||||
g = ggml_silu(ctx0, g);
|
||||
ggml_tensor * gu = ggml_mul(ctx0, g, u);
|
||||
cur = build_switched_lora_mm(layer.ffn_down, sl.a_down, sl.b_down, gu, adapter_ids);
|
||||
cb(cur, "ffn_out", il);
|
||||
|
||||
if (hparams.f_residual_scale) {
|
||||
cur = ggml_scale(ctx0, cur, hparams.f_residual_scale);
|
||||
}
|
||||
cur = ggml_add(ctx0, cur, ffn_inp);
|
||||
|
||||
cur = build_cvec(cur, il);
|
||||
cb(cur, "l_out", il);
|
||||
|
||||
return cur;
|
||||
}
|
||||
@@ -1596,6 +1596,56 @@ struct llama_model_granite_moe : public llama_model_base {
|
||||
};
|
||||
|
||||
|
||||
struct llama_model_granite_switch : public llama_model_base {
|
||||
llama_model_granite_switch(const struct llama_model_params & params) : llama_model_base(params) {}
|
||||
void load_arch_hparams(llama_model_loader & ml) override;
|
||||
void load_arch_tensors(llama_model_loader & ml) override;
|
||||
|
||||
uint32_t n_adapters = 0;
|
||||
uint32_t max_lora_rank = 0;
|
||||
float router_gain = 15.0f;
|
||||
|
||||
std::unordered_map<llama_token, int32_t> adapter_token_to_slot;
|
||||
std::unordered_map<llama_token, llama_token> adapter_token_to_substitute;
|
||||
|
||||
struct graph : public llm_graph_context {
|
||||
graph(const llama_model & model, const llm_graph_params & params);
|
||||
|
||||
private:
|
||||
ggml_tensor * build_switched_lora_delta(
|
||||
ggml_tensor * lora_a,
|
||||
ggml_tensor * lora_b,
|
||||
ggml_tensor * cur,
|
||||
ggml_tensor * ids);
|
||||
|
||||
ggml_tensor * build_switched_lora_mm(
|
||||
ggml_tensor * w,
|
||||
ggml_tensor * lora_a,
|
||||
ggml_tensor * lora_b,
|
||||
ggml_tensor * cur,
|
||||
ggml_tensor * ids);
|
||||
|
||||
ggml_tensor * build_attention_layer(
|
||||
ggml_tensor * cur,
|
||||
ggml_tensor * inp_pos,
|
||||
ggml_tensor * adapter_ids,
|
||||
llm_graph_input_attn_kv * inp_attn,
|
||||
const llama_model & model,
|
||||
const int64_t n_embd_head,
|
||||
const int il);
|
||||
|
||||
ggml_tensor * build_layer_ffn(
|
||||
ggml_tensor * cur,
|
||||
ggml_tensor * inpSA,
|
||||
ggml_tensor * adapter_ids,
|
||||
const llama_model & model,
|
||||
const int il);
|
||||
};
|
||||
|
||||
std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
|
||||
};
|
||||
|
||||
|
||||
struct llama_model_minicpm : public llama_model_base {
|
||||
llama_model_minicpm(const struct llama_model_params & params) : llama_model_base(params) {}
|
||||
void load_arch_hparams(llama_model_loader & ml) override;
|
||||
|
||||
@@ -411,6 +411,9 @@ static bool arch_supported(const llm_arch arch) {
|
||||
if (arch == LLM_ARCH_GEMMA4 || arch == LLM_ARCH_GEMMA4_ASSISTANT) {
|
||||
return false; // FIXME @ngxson
|
||||
}
|
||||
if (arch == LLM_ARCH_GRANITE_SWITCH) {
|
||||
return false; // FIXME adapter fixture
|
||||
}
|
||||
if (arch == LLM_ARCH_LLAMA_EMBED || arch == LLM_ARCH_GEMMA_EMBEDDING || arch == LLM_ARCH_T5ENCODER) {
|
||||
return false; // FIXME Embedding (?) models produce inconsistent results.
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user