Change GraniteSwitch inheritance to MoEShared

README.md

@@ -47,15 +47,16 @@ Other install options depending on your use case:
 ```bash
 pip install "granite-switch[compose]"   # Compose modular models
 pip install "granite-switch[hf]"        # HuggingFace inference
-pip install "granite-switch[vllm20]"    # vLLM 0.20+ (requires CUDA 13+)
+pip install "granite-switch[vllm20]"    # newer vLLM line (0.27.x)
 pip install "granite-switch[dev]"       # Everything
 ```
 
-Requires Python 3.10+ and PyTorch 2.0+.
+Requires Python 3.11+ and PyTorch 2.11+.
 
-> **vLLM version note:** This project currently defaults to vLLM 0.19.1 due to vLLM 0.20's
-> dependency on CUDA 13.0+ (via PyTorch 2.11), which is incompatible with many existing
-> environments running CUDA 12.x drivers. Use `.[vllm20]` if your environment supports CUDA 13+.
+> **vLLM version note:** This project requires `transformers>=5.16`, so it pins vLLM to the
+> `0.26.x` line (the earliest vLLM whose model code handles the transformers-5.13+
+> `full_attention` layer-type rename) via the default `[vllm]` extra. `[vllm20]` selects the
+> newer `0.27.x` line. Both require PyTorch 2.11+ (CUDA 13+).
 
 ### Compose a Model
 

docs/AUDIO.md

@@ -22,8 +22,8 @@ non-16 kHz input:
 uv sync --extra vllm --extra audio     # or --extra vllm20 --extra audio
 
 # Development / running the test suite (the dev groups include audio already)
-uv sync --group dev                    # vLLM 0.19.x
-uv sync --group dev-vllm20             # vLLM 0.20.x
+uv sync --group dev                    # vLLM 0.26.x
+uv sync --group dev-vllm20             # vLLM 0.27.x
 ```
 
 ## Building an audio-enabled checkpoint

pyproject.toml

@@ -10,8 +10,8 @@ readme = "README.md"
 license = "Apache-2.0"
 requires-python = ">=3.11,<3.14"
 dependencies = [
-    "torch>=2.10.0",
-    "transformers>=5.5.1,<5.10.0",
+    "torch>=2.11.0",
+    "transformers>=5.16,<=5.17.0",
 ]
 
 [project.urls]
@@ -20,8 +20,8 @@ Documentation = "https://github.com/generative-computing/granite-switch/tree/mai
 
 [project.optional-dependencies]
 hf = ["accelerate>=0.20.0"]
-vllm = ["vllm>=0.19.1,<0.20.0"]
-vllm20 = ["vllm>=0.20.0,<0.21.0"]
+vllm = ["vllm>=0.26.0,<0.27.0"]
+vllm20 = ["vllm>=0.27.0,<0.28.0"]
 compose = ["huggingface_hub", "pyyaml", "tqdm", "safetensors"]
 build = ["huggingface_hub", "pyyaml", "tqdm", "safetensors"]  # Backward compatibility alias for compose
 # Audio (ASR) decode + resample. Reuse vLLM's own audio deps (unversioned, so it
@@ -56,8 +56,8 @@ markers = [
 ]
 
 [dependency-groups]
-vllm19     = ["vllm>=0.19.1,<0.20.0"]
-vllm20     = ["vllm>=0.20.0,<0.21.0"]
+vllm19     = ["vllm>=0.26.0,<0.27.0"]
+vllm20     = ["vllm>=0.27.0,<0.28.0"]
 # `audio` is included so the audio tests can actually run: the ASR path needs
 # vLLM's audio deps (av/soundfile/resampy) at runtime, and no group pulled them
 # in before (integration tests failed with ModuleNotFoundError on a synced pod).

src/granite_switch/composer/__init__.py

@@ -11,8 +11,8 @@
     granite_dense_arch,
     granite_dense_sr_arch,
     granite_moe_arch,
-    granite_moe_hybrid_arch,
-    granite_moe_hybrid_sr_arch,
+    granite_moe_shared_arch,
+    granite_moe_shared_sr_arch,
     granite_moe_sr_arch,
     resolve_arch,
 )
@@ -28,8 +28,8 @@
     "granite_dense_arch",
     "granite_dense_sr_arch",
     "granite_moe_arch",
-    "granite_moe_hybrid_arch",
-    "granite_moe_hybrid_sr_arch",
+    "granite_moe_shared_arch",
+    "granite_moe_shared_sr_arch",
     "granite_moe_sr_arch",
     "resolve_arch",
 ]

src/granite_switch/composer/arch.py

@@ -220,7 +220,7 @@ def _dense_mlp_to_shared_groups() -> list[ModuleDescriptor]:
     Used for dense Granite models whose base uses ``mlp.gate_proj`` /
     ``mlp.up_proj`` / ``mlp.down_proj`` but whose switch model uses
     ``shared_mlp.input_linear`` / ``shared_mlp.output_linear``
-    (the ``GraniteMoeHybridMLP`` layout).
+    (the ``GraniteMoeSharedMLP`` layout).
     """
     return [
         ModuleDescriptor(
@@ -315,7 +315,7 @@ def _cross_stream_groups() -> list[ModuleDescriptor]:
     "embedding_multiplier": 1.0,
     "logits_scaling": 1.0,
     "attention_multiplier": 1.0,
-    # Granite/GraniteMoeHybrid vLLM classes use separate add-then-norm.
+    # Granite/GraniteMoeShared vLLM classes use separate add-then-norm.
     "fused_add_norm": False,
 }
 
@@ -326,8 +326,12 @@ def _cross_stream_groups() -> list[ModuleDescriptor]:
     "shared_intermediate_size": None,
 }
 
-# Layer type fields (propagated for hybrid models)
-_HYBRID_OPTIONAL_FIELDS: dict[str, Any] = {
+# Layer-type / positional fields. The switch model is attention-only, but it
+# still reads ``layer_types`` (lora_target_modules auto-detection) and
+# ``position_embedding_type`` (RoPE gating), so these are propagated from the
+# base config when present. ``GraniteSwitchConfig`` owns them as its own
+# attributes (the GraniteMoeShared parent does not declare them).
+_LAYER_TYPE_OPTIONAL_FIELDS: dict[str, Any] = {
     "layer_types": None,
     "position_embedding_type": "rope",
 }
@@ -338,16 +342,19 @@ def _cross_stream_groups() -> list[ModuleDescriptor]:
 # ---------------------------------------------------------------------------
 
 
-def granite_moe_hybrid_arch(base_config=None) -> ArchDescriptor:
-    """GraniteMoeHybrid architecture (model_type ``granitemoehybrid``).
+def granite_moe_shared_arch(base_config=None) -> ArchDescriptor:
+    """GraniteMoeShared architecture (model_type ``granitemoeshared`` /
+    ``granitemoehybrid``).
 
-    GraniteMoeHybrid models use ``shared_mlp`` module naming
+    Granite 4 MoE-with-shared-expert models use ``shared_mlp`` module naming
     (``shared_input_linear``, ``shared_output_linear``), even dense layers
-    with ``num_local_experts=0``.
+    with ``num_local_experts=0``. Real Granite 4.x dense checkpoints are still
+    typed ``granitemoehybrid`` upstream (they carry no mamba layers), so this
+    same descriptor serves both model_type strings.
     """
     optional_fields = dict(_GRANITE_OPTIONAL_FIELDS)
     optional_fields.update(_MOE_OPTIONAL_FIELDS)
-    optional_fields.update(_HYBRID_OPTIONAL_FIELDS)
+    optional_fields.update(_LAYER_TYPE_OPTIONAL_FIELDS)
 
     return ArchDescriptor(
         groups=list(_common_attn_groups()) + list(_moe_shared_mlp_groups()),
@@ -361,11 +368,12 @@ def granite_moe_arch(base_config=None) -> ArchDescriptor:
 
     Pure sparse MoE: every layer has an expert bank and **no** dense
     ``shared_mlp``.  The descriptor is therefore ``_common_attn_groups()`` and
-    nothing else — a strict subset of :func:`granite_moe_hybrid_arch`.
+    nothing else — a strict subset of :func:`granite_moe_shared_arch`.
 
-    The frozen expert tensors (``block_sparse_moe.input_linear`` /
-    ``output_linear`` / ``router.layer``) are named identically in the switch
-    model, so with no shared-MLP group to shadow them they transfer by identity.
+    The frozen expert tensors (``block_sparse_moe.experts.gate_up_proj`` /
+    ``experts.down_proj`` / ``router.weight`` in the transformers-5.16 layout)
+    are named identically in the switch model, so with no shared-MLP group to
+    shadow them they transfer by identity.
 
     ``shared_intermediate_size`` is pinned to ``0``, which is upstream's own
     encoding for "no shared MLP"
@@ -409,13 +417,13 @@ def granite_dense_arch(base_config=None) -> ArchDescriptor:
 ]
 
 
-def granite_moe_hybrid_sr_arch(base_config=None) -> ArchDescriptor:
-    """GraniteMoeHybrid Shadow Residual architecture.
+def granite_moe_shared_sr_arch(base_config=None) -> ArchDescriptor:
+    """GraniteMoeShared Shadow Residual architecture.
 
-    Identical to :func:`granite_moe_hybrid_arch` (fused projections) plus the
+    Identical to :func:`granite_moe_shared_arch` (fused projections) plus the
     layer-level ``cross_stream`` injection site.
     """
-    arch = granite_moe_hybrid_arch(base_config=base_config)
+    arch = granite_moe_shared_arch(base_config=base_config)
     arch.groups = arch.groups + _cross_stream_groups()
     arch.buffer_keywords = list(_SR_BUFFER_KEYWORDS)
     return arch
@@ -444,7 +452,12 @@ def granite_dense_sr_arch(base_config=None) -> ArchDescriptor:
 _ARCH_REGISTRY = {
     "granite": granite_dense_arch,
     "granitemoe": granite_moe_arch,
-    "granitemoehybrid": granite_moe_hybrid_arch,
+    # granitemoeshared is the de-hybridized base family. The granitemoehybrid
+    # key is retained because real Granite 4.x dense checkpoints are still typed
+    # granitemoehybrid upstream (they carry no mamba layers); both resolve to the
+    # same shared-expert descriptor.
+    "granitemoeshared": granite_moe_shared_arch,
+    "granitemoehybrid": granite_moe_shared_arch,
 }
 
 # Must stay key-for-key in step with _ARCH_REGISTRY: a model_type registered in
@@ -453,7 +466,8 @@ def granite_dense_sr_arch(base_config=None) -> ArchDescriptor:
 _SR_ARCH_REGISTRY = {
     "granite": granite_dense_sr_arch,
     "granitemoe": granite_moe_sr_arch,
-    "granitemoehybrid": granite_moe_hybrid_sr_arch,
+    "granitemoeshared": granite_moe_shared_sr_arch,
+    "granitemoehybrid": granite_moe_shared_sr_arch,
 }
 
 

src/granite_switch/composer/compose_utils.py

@@ -7,7 +7,7 @@
 
 import torch
 
-from ..config import SWITCH_CACHE_LAYERS
+from ..config import ATTENTION_LAYER_TYPE, SWITCH_CACHE_LAYERS
 from .adapter_loader import (
     _extract_modules_from_weights,
     detect_lora_config,
@@ -188,20 +188,22 @@ def from_base_and_adapters(
         for field_name, default in arch.optional_config_fields.items():
             config_kwargs[field_name] = getattr(base_config, field_name, default)
 
-        # For Granite 3.x whose arch descriptor doesn't include
-        # shared_intermediate_size, default it to intermediate_size.
-        # GraniteMoeHybridConfig defaults it to 1024 (not None), so
-        # GraniteSwitchConfig's fallback logic doesn't trigger.
+        # For a dense Granite base whose arch descriptor doesn't include
+        # shared_intermediate_size, supply it explicitly from intermediate_size:
+        # a dense Granite layer always has a shared MLP of that width. This makes
+        # the composer the source of truth and does not rely on any parent-class
+        # default (GraniteMoeShared defaults it to 0, which is the "no shared MLP"
+        # sentinel — GraniteSwitchConfig would keep that 0 verbatim if passed).
         if "shared_intermediate_size" not in config_kwargs:
             config_kwargs["shared_intermediate_size"] = config_kwargs[
                 "intermediate_size"
             ]
 
-        # Normalize layer_types: map everything to "attention" (only attention
-        # layers are supported).
+        # Normalize layer_types: map everything to the attention layer type
+        # (only attention layers are supported).
         lt = config_kwargs.get("layer_types")
         if lt is not None:
-            config_kwargs["layer_types"] = ["attention" for _ in lt]
+            config_kwargs["layer_types"] = [ATTENTION_LAYER_TYPE for _ in lt]
 
         # When adapters are present, reserve the switch's cache slots at the
         # front: MultiSwitch (coded) owns SWITCH_CACHE_LAYERS == 2 (counting +
@@ -213,7 +215,7 @@ def from_base_and_adapters(
             )
             if config_kwargs.get("layer_types") is not None:
                 config_kwargs["layer_types"] = [
-                    *(["attention"] * SWITCH_CACHE_LAYERS),
+                    *([ATTENTION_LAYER_TYPE] * SWITCH_CACHE_LAYERS),
                     *list(config_kwargs["layer_types"]),
                 ]
 

src/granite_switch/composer/tokenizer_setup.py

@@ -288,6 +288,27 @@ def get_alora_first_invocation_token_id(adapter_path: str) -> int:
     return _load_alora_invocation_token_ids(adapter_path)[0]
 
 
+def _alora_invocation_drop_chars(first_token_id: int, tokenizer) -> int:
+    """Number of leading characters of the invocation text that its first
+    token spans, under *tokenizer*.
+
+    Pass 2 of the chat template writes the control token followed by the
+    invocation text with its first *token* removed (the runtime swaps the
+    control-token embedding for the first-invocation-token embedding, so
+    re-emitting that token would duplicate it at the swap site). The template
+    can only slice the invocation *string*, so it needs the character length of
+    the first token — not a hardcoded 1. For most Granite invocations the first
+    token is a lone ``<`` (1 char), but some (e.g. ``<context>`` ->
+    ``<context``) tokenize with a multi-character first token under newer
+    tokenizers, and dropping a single character there would leave a partial
+    token in the tail.
+
+    Computed by decoding the first invocation token id, so it stays consistent
+    with the sequence the runtime swaps and needs only ``tokenizer.decode``.
+    """
+    return len(tokenizer.decode([first_token_id], skip_special_tokens=False))
+
+
 #: A ``{{- ... }}`` emission in a Jinja template.
 _EMISSION_RE = re.compile(r"\{\{-?\s*(.*?)\s*-?\}\}", re.DOTALL)
 
@@ -939,6 +960,17 @@ def configure_chat_template(
                 sr_anchors.add(anchor_text)
             else:
                 entry["invocation_text"] = anchor_text
+                # Chars the first invocation token spans; Pass 2 drops exactly
+                # these from the tail so the runtime embedding swap does not
+                # duplicate the first token (see _alora_invocation_drop_chars).
+                # Re-encode the decoded invocation text with this tokenizer to
+                # get the first token id (mirrors the SR encode round-trip).
+                first_invocation_id = tokenizer.encode(
+                    anchor_text, add_special_tokens=False
+                )[0]
+                entry["invocation_drop_chars"] = str(
+                    _alora_invocation_drop_chars(first_invocation_id, tokenizer)
+                )
         adapter_mapping[adapter_name] = entry
 
     if len(sr_anchors) > 1:
@@ -955,7 +987,8 @@ def configure_chat_template(
             mapping_entries.append(
                 f"    '{adapter_name}': {{'token': '{info['token']}', "
                 f"'type': '{info['type']}', "
-                f"'invocation_text': '{info['invocation_text']}'}}"
+                f"'invocation_text': '{info['invocation_text']}', "
+                f"'invocation_drop_chars': {info['invocation_drop_chars']}}}"
             )
         else:
             mapping_entries.append(
@@ -969,11 +1002,13 @@ def configure_chat_template(
 {%- set adapter_token = '' %}
 {%- set adapter_type = '' %}
 {%- set adapter_invocation_text = '' %}
+{%- set adapter_invocation_drop_chars = 1 %}
 {%- if adapter_name is defined and adapter_name in adapter_map %}
 {%- set adapter_token = adapter_map[adapter_name]['token'] %}
 {%- set adapter_type = adapter_map[adapter_name]['type'] %}
 {%- if adapter_map[adapter_name]['type'] == 'alora' %}
 {%- set adapter_invocation_text = adapter_map[adapter_name]['invocation_text'] %}
+{%- set adapter_invocation_drop_chars = adapter_map[adapter_name]['invocation_drop_chars'] %}
 {%- endif %}
 {%- endif %}
 
@@ -1047,38 +1082,40 @@ def configure_chat_template(
     # in the message.
     #
     # Token drop (mirrors the role-marker skip-once flag used for LoRA /
-    # assistant-boundary ALoRA): we also omit the FIRST CHARACTER of the
-    # invocation text. The runtime embedding swap replaces the control-token
-    # embedding with the first-invocation-token's embedding; writing the full
-    # invocation text after the control token would then produce two copies
-    # of that first-invocation-token back to back — an OOD pattern at the
-    # swap site.
+    # assistant-boundary ALoRA): we omit the FIRST TOKEN of the invocation
+    # text. The runtime embedding swap replaces the control-token embedding with
+    # the first-invocation-token's embedding; writing the full invocation text
+    # after the control token would then produce two copies of that
+    # first-invocation-token back to back — an OOD pattern at the swap site.
     #
-    # For every granite_format ALoRA invocation text in the standard Granite adapter
-    # library (<requirements>, <certainty>, <guardian>, <context>, etc.) the
-    # first character is a single '<' that the tokenizer emits as its own token,
-    # and the tail of the string retokenizes identically to the tail of the full
-    # string. So dropping the first character on the string side is equivalent
-    # to dropping exactly the first token on the tokenized side. For ChatML the
+    # Jinja can only slice the invocation *string*, not its tokens, so the drop
+    # is expressed as a per-adapter character count baked into ``adapter_map``
+    # (``invocation_drop_chars`` = the length of the first invocation token's
+    # decoded text; see ``get_alora_invocation_drop_chars``). For most Granite
+    # invocations the first token is a lone '<' (drop 1), but some — e.g.
+    # ``<context>`` (first token ``<context``), ``<citation>``, ``<hallucination>``
+    # under tokenizers >= 0.23.1 — tokenize with a multi-character first token,
+    # where a hardcoded 1-char drop would leave a partial token in the tail.
+    # Computing the count at compose time keeps the string slice equivalent to a
+    # first-token drop regardless of the tokenizer's merges. For ChatML the
     # trained 4.2 adapters use the assistant-boundary invocation
     # (<|im_start|>assistant\\n) and therefore take the fallback path below, not
-    # Pass 2; a user-message ChatML ALoRA whose invocation text does not begin
-    # with a standalone-tokenizing character would need the first-token-drop
-    # invariant re-checked (see the property test in test_chat_template.py).
+    # Pass 2. The invariant is guarded by the property test in
+    # test_chat_template.py.
     #
     # ``content_var`` is ``content.val`` for the granite_format namespace-object content
     # or ``content`` for ChatML's plain-string content.
     alora_pass2 = (
-        """    {#- ALoRA Pass 2: inject activation token AND drop the first char of
-         the invocation text so the runtime-swapped embedding doesn't duplicate. -#}
+        """    {#- ALoRA Pass 2: inject activation token AND drop the first token's
+         chars of the invocation text so the runtime-swapped embedding doesn't duplicate. -#}
     {%- if loop.index0 == ns.alora_target_idx %}
         {%- set _parts = """
         + content_var
         + """.rsplit(ns.adapter_invocation_text, 1) %}
         {%- if _parts | length > 1 %}
             {%- set """
         + content_var
-        + """ = _parts[0] + ns.adapter_token + ns.adapter_invocation_text[1:] + _parts[1] %}
+        + """ = _parts[0] + ns.adapter_token + ns.adapter_invocation_text[ns.adapter_invocation_drop_chars:] + _parts[1] %}
         {%- endif %}
     {%- endif %}
 """
@@ -1126,6 +1163,7 @@ def configure_chat_template(
             "\n                       adapter_token=adapter_token,"
             "\n                       adapter_type=adapter_type,"
             "\n                       adapter_invocation_text=adapter_invocation_text,"
+            "\n                       adapter_invocation_drop_chars=adapter_invocation_drop_chars,"
             "\n                       alora_target_idx=-1,"
             "\n                       skip_next_role_marker=false"
             "\n                       )"