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:
Bar Haim
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user